Lightning Thunder: Supercharged PyTorch for Modern Hardware
summary
Lightning Thunder is an open source compiler that traces a PyTorch NNModule on the meta device and emits a prologue, a purely functional trace, and an epilogue so transforms can shard, fuse, and inject libraries before moving tensors to hardware. It targets FP8/FP6/FP4 via transformer engine and uses MD fuser or torch compile for JIT kernels, composes FSDP + DDP meshes, runs on MVL72 (144 GPUs), and shows bigger speedups on P200 than on H100.
AI-written summary · about
chapters (10)
- 0:00Introduction
- 1:05Hardware strategies
- 3:56What is Lightning Thunder
- 6:23Model optimization state
- 8:54Thunder compilation
- 11:00FP8 transforms and sharding
- 12:45FSDP and DDP composition
- 14:33Thunder Jet and plugins
- 16:21Automating transformations
- 17:04Current and future state
related talks
26mThunder: Distribute and Optimize Your PyTorch Models With Zero...
12mLightning Talk: Accelerating Inference on CPU with Torch.Compile
25mInto Generative AI with PyTorch Lightning 2.0
14mLightning Talk: Every Millisecond Counts: The Fine-tuning Journey of an Ultra-Eff...
24mAdvancements in High-Performance AI/ML through PyTorch's Python Compiler
08mLightning Talk: Enhancements Made to MPS Backend in PyTorch for Applications Running...
36mDL Compiler Panel Discussion
12mLightning Talk: State of PyTorch