Conclusion
Accelerated PyTorch 2 Transformers were designed to make the training and production deployment of state-of-the-art transformer models affordable and integrated with PyTorch 2.0 model JIT compilation.  The newly introduced PyTorch SDPA operator provides improved performance for training Transformer models and is particularly valuable for the expensive Large Language Model training. In this post we demonstrate a number of optimizations on the exemplary nanoGPT model  including:

Over 26% training speedup, when compared against the baseline with constant batch size
An additional speedup achieved with padded vocabulary, bringing the total optimization to approximately 64% compared to the baseline
Additional numerical stability

Appendix A: Analyzing Attention Numeric Stability