Lightning Talk: Graph Based Pipeline Parallelism - Sanket Purandare, Meta & Simon Fan, Meta PyTorch

This talk discusses the challenges of pipeline parallelism in large models and how current PyTorch implementations struggle with advanced scheduling for state-of-the-art language models (LLMs). The speaker examines the issues related to model parallelism, including the communication overhead of mixture of experts (MoE) and the limitations of eager-mode implementations. They introduce a solution called Graph-Based Pipeline Parallelism, which compiles stages into reusable FX graphs executed through an explicit scheduling language, allowing users to write standard PyTorch code while managing schedules with greater ease and composability. The integration of Graph-PP into TorchTitan and AutoParallel demonstrates its effectiveness on real-world MoE workloads, with plans for upstream inclusion in torch.distributed.