Lightning Talk: Large-Scale Distributed Training with Dynamo and...
summary
Experimental PyTorch/XLA SPMD uses XLA GSPMD to transform single-device PyTorch programs into partitioned ones with mark_sharding annotations and integrates with DTensor APIs for XLA meshes. Benchmarks on LLaMA2 (7B–70B) show ~28% speedup over FSDP and maintain ~55% MFU across sizes; scaling is near-linear with sequence length, with performance drop within 10% (as low as 4% for 70B). Implemented 2D sharding of weights and activations to avoid gathering full tensors and achieve high memory and compute efficiency.
AI-written summary · about
chapters (4)
- 0:00Introduction to PyTorch SPMD
- 0:26Using the SPMD API
- 2:31Sharding Mechanics
- 4:18Performance and Implementation
related talks
08mPyTorch Distributed: Towards Large Scale Training
27mTips and tricks for distributed large model training
24mTraining large scale models using PyTorch
38mMPI Meets Machine Learning: Unlocking PyTorch distributed for scaling AI workloads
01h 03mScaling up machine learning without tears (and what do programming languages ha...
10mLightning Talk: Scaling Recommendation Systems To 2K GPUs and Beyond
19mNo GPU Left Behind: Scaling Online LLM Training With...
26mThunder: Distribute and Optimize Your PyTorch Models With Zero...