# SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile

arXiv: https://arxiv.org/abs/2411.00284 · HTML: https://mlsanity.com/paper/2411.00284/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 36 references that are in the map.
- First posted on arXiv 1 Nov 2024.

## Builds on (most cited first, 20 of 20)

- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [BLOOM: A 176B-Parameter Open-Access Multilingual Language Model](https://mlsanity.com/paper/2211.05100/) (2022)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [NeMo: a toolkit for building AI applications using Neural Modules](https://mlsanity.com/paper/1909.09577/) (2019)
- [Multimodal C4: An Open, Billion-scale Corpus of Images Interleaved with Text](https://mlsanity.com/paper/2304.06939/) (2023)
- [Network Topology and Communication-Computation Tradeoffs in Decentralized Optimization](https://mlsanity.com/paper/1709.08765/) (2017)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://mlsanity.com/paper/2410.06511/) (2024)
- [TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models](https://mlsanity.com/paper/2102.07988/) (2021)
- [Optimus: Accelerating Large-Scale Multi-Modal LLM Training by Bubble Exploitation](https://mlsanity.com/paper/2408.03505/) (2024)
- [Slapo: A Schedule Language for Progressive Optimization of Large Deep Learning Model Training](https://mlsanity.com/paper/2302.08005/) (2023)

## Built on by (most cited first, 7 of 7)

- [MuonBP: Faster Muon via Block-Periodic Orthogonalization](https://mlsanity.com/paper/2510.16981/) (2025)
- [Entropy-Preserving Reinforcement Learning](https://mlsanity.com/paper/2603.11682/) (2026)
- [Characterizing the Efficiency of Distributed Training: A Power, Performance, and Thermal Perspective](https://mlsanity.com/paper/2509.10371/) (2025)
- [Hydraulis: Balancing Large Transformer Model Training via Co-designing Parallel Strategies and Data Assignment](https://mlsanity.com/paper/2412.07894/) (2024)
- [Bergson: An Open Source Library for Data Attribution](https://mlsanity.com/paper/2606.11660/) (2026)
- [AutoSP: Unlocking Long-Context LLM Training Via Compiler-Based Sequence Parallelism](https://mlsanity.com/paper/2604.27089/) (2026)
- [Lingua Custodi's participation at the WMT 2025 Terminology shared task](https://mlsanity.com/paper/2510.17504/) (2025)

## Abstract (arXiv)

Distributed training of large models consumes enormous computation resources and requires substantial engineering efforts to compose various training techniques. This paper presents SimpleFSDP, a PyTorch-native compiler-based Fully Sharded Data Parallel (FSDP) framework, which has a simple implementation for maintenance and composability, allows full computation-communication graph tracing, and brings performance enhancement via compiler backend optimizations. SimpleFSDP's novelty lies in its unique torch.compile-friendly implementation of collective communications using existing PyTorch primitives, namely parametrizations, selective activation checkpointing, and DTensor. It also features the first-of-its-kind intermediate representation (IR) nodes bucketing and reordering in the TorchInductor backend for effective computation-communication overlapping. As a result, users can employ the aforementioned optimizations to automatically or manually wrap model components for minimal communication exposure. Extensive evaluations of SimpleFSDP on Llama 3 models (including the ultra-large 405B) using TorchTitan demonstrate up to 28.54% memory reduction and 68.67% throughput improvement compared to the most widely adopted FSDP2 eager framework, when composed with other distributed training techniques.