Lightning Talk: Lessons from Using Pytorch 2.0 Compile in IBM's Watsonx.AI Inference
summary
Open source LLaMA implementation for Watsonx.AI uses torch.compile() plus SDPA, tensor parallel and CUDA graphs and shows at least a 2X improvement on many models, with 70B gaining ~50% extra from CUDA graphs versus compile alone; 7B runs drop from ~40 ms per token to ~30–40 ms depending on batch and sequence length. Key workarounds include reimplementing rotor embeddings because inductor lacks complex support, fixing Flash attention causing spurious recompiles with dynamic shapes, and making tensor parallel work around missing dtensor compile support, all upstreamed to PyTorch and available on the project GitHub.
AI-written summary · about
chapters (6)
- 0:00PyTorch 2.0 compile in Watsonx.AI
- 1:26Performance improvements
- 3:13Overcoming graph breaks
- 5:20Tensor parallel with PyTorch compile
- 6:46Next steps and optimizations
- 9:16Conclusion
related talks
23mBlazing Fast GenAI Inference With Torch.compile
36mDL Compiler Panel Discussion
15mPyTorch 2.0: Unlocking the Power of Deep Learning with the Torch Compile API
24mAdvancements in High-Performance AI/ML through PyTorch's Python Compiler
14mLightning Talk: Accelerating PyTorch Models With Torch.compile's C++ Wrapper Mode
32mKeynote: Fireside Chat with Soumith Chintala, Co-Creator of PyTorch & Merve Noyan
34mSpeeding up Your ML Workload: Pytorch Compile and Distributed Training
29mKeynote