PyTorch Symmetric Memory: A New Programming Paradigm for Distributed AI
summary
PyTorch adds Symmetric Memory: a global GPU address space that lets device kernels read/write peer GPU memory directly and author custom communication logic in Triton. Implementations include fused compute+comm kernels, RDMA via NVIDIA GDS/IBGDA to bypass CPU, and a tiled multi-root reduce. Results show 4x inference speedup for dynamic MoE shuffles and 2x faster for small all-reduce compared with NCCL, and the Triton plugin + NVISHMAN primitives prove the approach works end-to-end.
AI-written summary · about
chapters (13)
- 0:00PyTorch Symmetric Memory intro
- 1:26Symmetric memory in PyTorch
- 2:52Symmetric memory challenges
- 5:30Symmetric memory example
- 7:16Capabilities unlocked
- 8:01Custom MoE patterns
- 11:19Tiling for blockwise opt
- 13:02Triton custom kernels
- 16:41Fusing compute and comm
- 18:30Scaling across nodes
- 20:19Cross-node with Enishman
- 21:49Future plans
- 22:53Q&A
related talks
19mSymmetric Memory: A New Paradigm for Programming Distributed AI
25mPyTorch Symmetric Memory + NCCL Device APIs: A New Path Towards Multi-GP...
34mSymmetric Memory in PyTorch: 10x Faster GPU Communication for AI
15mColocating ML Inference and Training with Fast GPU Memory Handover
38mMPI Meets Machine Learning: Unlocking PyTorch distributed for scaling AI workloads
08mPyTorch Distributed: Towards Large Scale Training
58mKeynote: Enabling the AI Revolution: Next-Generation Memory Solutions for AI Computing(…)
25mRace to Superintelligence: Building Low-Friction AI Infrastructure for a Heterogen...