A minimal and educational implementation of Fully Sharded Data Parallel (FSDP).
Built by digging through PyTorch's FSDP codebase and isolating the core mechanics of sharding, gathering, and state transitions for educational clarity.
This implementation mirrors some of the core concepts of torch.distributed.fsdp:
- Parameter Sharding: Implements a
FULL_SHARDstrategy, where model parameters, gradients, and optimizer states are sharded across data-parallel ranks. - Auto Wrapping: Provides a policy-based function (
wrap_fn) for automatically wrapping and sharding. - Activation Checkpointing: Supports activation checkpointing for memory efficiency during forward and backward passes.
This implementation mainly aims to show FSDP states and communication. As such, it lacks many optimizations and features of the official PyTorch FSDP, for example:
- Sync Collectives: All collective operations (
all_gather,reduce_scatter) are blocking, with no overlap between communication and computation (i.e., no prefetching). - No Mixed Precision: Does not support mixed precision; computations/communication are performed using the provided dtype (e.g., master weights are not forced to be
float32). - No Meta Device: Doesn't use the meta device to optimize sharding.
- No CPU Offloading: Sharded parameters remain on the GPU, with no mechanism for CPU offloading to save memory.
Ensure you have uv installed, then:
uv syncAlternatively, you can use standard pip and venv:
python3.11 -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install .You can run a minimal training loop comparing edufsdp with the standard torch FSDP, by training a fresh Qwen2-0.5 on TinyStories:
# Two ranks, Qwen2-0.5B and EduFSDP
PROCESSES=2 MODEL=qwen2 BACKEND=nccl MBS=8 ACC=2 STEPS=3 EVAL=0 FSDP=edufsdp ./scripts/train.shRanks = 2, FSDP = edufsdp, Acc = 2
1/3: Loss = 12.067031, Peak = 13006.0 MB, Time = 6.7548s
2/3: Loss = 11.321454, Peak = 15382.0 MB, Time = 1.0896s
3/3: Loss = 10.492050, Peak = 15382.0 MB, Time = 1.0921s# ...same with FSDP=torch
Ranks = 2, FSDP = torch, Acc = 2
1/3: Loss = 12.067031, Peak = 14238.0 MB, Time = 7.7054s
2/3: Loss = 11.321454, Peak = 16614.0 MB, Time = 1.1203s
3/3: Loss = 10.492050, Peak = 16614.0 MB, Time = 1.1080sLet it run long enough if you only care about stories (EVAL=20 STEPS=1000):
280/1000: Loss = 3.615633, Peak = 15382.0 MB, Time = 1.1021s
281/1000: Loss = 3.493644, Peak = 15382.0 MB, Time = 1.1021s
=== eval ===
Once upon a time there was a young little girl named Lily. She loved to sing all day. One day
282/1000: Loss = 3.390410, Peak = 15382.0 MB, Time = 2.4894s
283/1000: Loss = 3.459949, Peak = 15382.0 MB, Time = 1.1043s
If you don't have access to an NVIDIA GPU, you can try with the GLOO backend on CPU (also with a very tiny model to just make the code run):
PROCESSES=2 MODEL=tiny MBS=16 ACC=2 STEPS=3 EVAL=0 FSDP=edufsdp ./scripts/train.shRanks = 2, FSDP = edufsdp, Acc = 2
1/3: Loss = 1.021882, Peak = N/A, Time = 0.0362s
2/3: Loss = 1.010545, Peak = N/A, Time = 0.0035s
3/3: Loss = 1.040771, Peak = N/A, Time = 0.0027s
Memory visualization for a single training step (GBS=8, SEQ=256) of Qwen2-0.5B across two GPUs, generated using the PyTorch Memory Visualizer:
| EduFSDP | PyTorch FSDP |
|---|---|
 |
 |