Efficient MoE Pre-training at Scale on AMD GPUs With TorchTitan
summary
Claims an end-to-end PyTorch pipeline (TorchTitan + Primus Turbo) that uses FP8 blockwise quantization and grouped GEMM via torch._scaled_mm plus a custom Triton kernel removes CPU–GPU sync to speed the forward pass ~3x, get ~2x from FP8 GEMM and up to 277% with FP8 grouped GEMM, and run DeepSeek-R1 and LLaMA4 pretraining on AMD MI300X hardware. Demonstrates near-linear strong scaling to 1,000+ MI300X GPUs on Tensorwave, with code and kernels open-sourced.
AI-written summary · about
chapters (11)
- 0:00TorchTitan and AMD GPUs
- 1:25TorchTitan MoE overview
- 3:55MoE pre-training challenges
- 6:02AMD GPU training benefits
- 8:10Parallel scaling strategy
- 8:54Primus Turbo library
- 11:23Primus Turbo optimizations
- 14:12DeepSeek V3 scaling results
- 16:23Training accuracy validation
- 17:06Conclusion and future plans
- 18:56Q&A
related talks
13mFrom Silicon to Scale: Training at Scale
32mPowering Your Generative AI Workloads with AMD and Open-Source ROCm
11mLightning Talk: Trinity Large - Torchtitan on 2000+ B300s
24mEfficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
01h 18mOptimizing LLM Training and Inference Performance on GPUs
20mMetis: Fast Automatic Distributed Training on Heterogeneous GPUs
08mGPU Programming with Triton Kernels
21mCost Effectively Deploy Thousands of Fine Tuned Gen AI Models Like...