The vocab size determines the dimensions of matmuls in the output layer of GPT, and these are so large that they were taking a majority of the time for the entire training loop!  We discovered that they were achieving performance significantly below the peak throughput achievable on the A100 GPU, and guessed from NVIDIA's matmul documentation that 64-element alignment would yield better results.  Indeed, padding these matmuls achieves nearly a 3x speedup!  The underlying cause is that unaligned memory accesses significantly reduce efficiency.  A deeper analysis can be found in this Twitter thread.
With this optimization we were able to further reduce training time from ~113 ms (using flash attention) to ~87 ms per batch.
Results
The figure below demonstrates the performance gained using Pytorch custom kernels. Here are the exact figures: