---
title: 'SimpleFSDP: Efficient FSDP in PyTorch'
url: https://www.emergentmind.com/topics/simplefsdp
type: topic
---

# SimpleFSDP: Efficient FSDP in PyTorch

SimpleFSDP is a PyTorch-native, compiler-based framework for Fully Sharded Data Parallel (FSDP) training, distinguished by its succinct implementation, deep integration with the $torch.compile$ ecosystem, and ability to expose the entire computation-communication dependency graph for aggressive compiler backend optimizations. Its design leverages native PyTorch primitives—DTensor for device placement and sharding, parametrizations for collective communication, and selective activation checkpointing—to enable full tracing and composable distributed training. SimpleFSDP unlocks new capabilities in computation-communication overlap and memory efficiency, especially for training ultra-large language models, by eliminating reliance on Python hooks and by exploiting advanced bucketing and reordering optimizations within the TorchInductor backend [2411.00284].

## 1. High-Level Architecture and Workflow

SimpleFSDP refactors the conventional PyTorch FSDP2 eager approach, which relies on backward hooks, into a model where all communication and computation are embedded in the FX graph and thus traceable by $torch.compile$. The core workflow is as follows:
1. **Parameter Sharding**: Each $nn.Parameter$ is wrapped as a DTensor, sharded across the data-parallel mesh: 
   $$
   \text{layout} = \text{Shard}(dim=0, \text{mesh}="data\_parallel")
   $$
2. **Forward All-Gather**: During forward computation, a parametrization module calls $DTensor.redistribute$, gathering the full parameter into activations.
3. **Activation Cleanup**: Activations are immediately released via checkpointing, reducing peak memory consumption.
4. **Backward Gradient Reduction**: In the backward pass, $DTensor.redistribute$ all-gathers shards as needed and reduces/scatters gradients.

All operators are inside the graph, enabling inspection, rewriting, and optimization at the IR level. DTensor metadata self-manages device, rank, and mesh placement, obviating explicit $dist.get\_rank()$ or NCCL communicator orchestration.

## 2. Collective Communications via PyTorch Primitives

Collective operations are transparently mapped to autograd-aware PyTorch primitives:
- **Parametrization API**: Each parameter uses a custom parametrization. Its forward logic is,
  ```python
  def replicate_compute(self, x):
      y = x.redistribute(
          placements=(Replicate(),),
          forward_dtype=self.param_dtype,
          backward_dtype=self.reduce_dtype
      ).to_local(grad_placements=(Partial("avg"),))
      return y
  ```
  $Replicate()$ denotes replication across the mesh, and $Partial(reduce\_op="avg")$ codes reduce-scatter for gradients.
- **Selective Activation Checkpointing**: Only the parametrization module is checkpointed, ensuring that (a) forward all-gather and local compute run once with immediate activation release, and (b) backward gathers and reduces gradients on demand.
- **DTensor Role**: DTensor encapsulates placement and coordinates; on initialization, each parameter is sharded. This abstraction removes manual rank and communication management.

*This approach guarantees full end-to-end tracing with all communication and computation both visible and optimizable for $torch.compile$.*

## 3. TorchInductor Backend Optimizations

After FX→IR lowering, SimpleFSDP presents a long IR node sequence:
$$
AG_1; C_1; Wa_1; AG_2; C_2; Wa_2; \ldots RS_1; Wr_1; RS_2; Wr_2; \ldots
$$
where $AG$ is all-gather, $C$ is compute, $Wa$ is wait for AG, $RS$ is reduce-scatter, $Wr$ is wait for RS.

### Bucketing IR Nodes
- **Bucketing Objective**: Communications over multiple small shards are grouped (“bucketed”) so that one NCCL collective is issued per bucket.
- **Implementation**: Input shards $in_1, in_2$ are concatenated, collectively gathered, and outputs split. Pseudocode:
  ```python
  flat_in = torch.cat([in1, in2], dim=0)
  flat_out = Mesh.redistribute(flat_in, Replicate())
  out1, out2 = split(flat_out, [n1, n2])
  ```

### Reordering for Overlap
- **Async Streams**: All-gather/reduce-scatter are asynchronous, allowing movement of IR nodes for maximal overlap under dependency constraints.
- **Example Reorder**: $AG_2 \rightarrow C_1 \rightarrow Wa_1 \rightarrow C_2 \rightarrow Wa_2$ can be reordered as $AG_2 \rightarrow AG_1_{wait} \rightarrow C_1 \rightarrow Wa_2 \rightarrow C_2 \rightarrow Wait\,AG_1_{wait}$.

*This exposes computation-communication concurrency and minimizes communication bubbles, subject to dependency analysis validity.*

## 4. Mathematical Formalisms

### Memory Usage
For layer $\ell$ with parameter tensor size $S_\ell$ (after fp16/bf16 reduction):
- **Eager (no sharding)**: $M_\ell^{before} = 2S_\ell$ (weights + gradients)
- **FSDP2**: $M_\ell^{after} \approx \frac{2S_\ell}{P} + O(\text{overlap\_temp})$
- **SimpleFSDP**: Immediate freeing and delayed regather yields
  $$
  M_\ell^{after} = \max(S_\ell, S_\ell/P) + \text{overhead\_checkpoint}
  $$
  Usually $S_\ell \gg S_\ell/P$, so SimpleFSDP peaks at $S_\ell$, not $2S_\ell$ or $S_\ell/P + S_\ell$.

### Communication Volume & Overlap Efficiency
- **Per-Shard All-Gather**: $V_\ell = \frac{(P-1)}{P}S_\ell$ bytes.
- **Bucketing**: For $B$ layers, $V_{bucket} = \frac{(P-1)}{P}\sum_\ell S_\ell$
- **Overlap Efficiency**: $\eta = 1 - \frac{\text{exposed\_comm\_time}}{\text{total\_comm\_time}}$

*These formulations quantify the memory and communication efficiency gains achievable by SimpleFSDP compared to legacy approaches.*

## 5. Benchmark Results on Llama 3

Extensive evaluation on real Llama 3 models (8B, 70B, 405B) using TorchTitan orchestration, PyTorch 2.x with TorchInductor, and advanced hardware (16-node, 8×NVIDIA H100/node, NVLink, InfiniBand):

| Model   | FSDP2-eager   | FSDP2-compile   | SimpleFSDP      |
|---------|---------------|----------------|-----------------|
| 8B      | 47K tok/s, 86.8GB | 52K tok/s, 79.3GB | 52K tok/s, 62.8GB |
| 70B     | 8.5K tok/s, 224GB | 9.3K tok/s, 208GB | 10.3K tok/s, 187GB |
| 405B    | 1.1K tok/s, 650GB | 1.2K tok/s, 610GB | 1.8K tok/s, 544GB |

- **Memory reduction**: Up to 28.54% vs FSDP2-eager; 8.37% vs FSDP2-compile.
- **Throughput improvement**: Up to 68.67% vs FSDP2-eager; 6.06% vs FSDP2-compile.

A plausible implication is that, when composed with distributed training and parallelism techniques, SimpleFSDP enables scalable training of ultra-large models with reduced memory footprint and improved token throughput [2411.00284].

## 6. Usage Patterns and Composability

Several usage recipes are described:

### Automatic & Manual Wrapping
Configuration of automatic bucketing and reorder:
```python
import torch
torch._inductor.config.simplefsdp.bucket_mode = "auto"
torch._inductor.config.simplefsdp.enable_reorder = True
model = simple_fsdp(model)
model = torch.compile(model, fullgraph=True)
```
Manual wrapping of model components:
```python
from simple_fsdp import ModuleWrapList
wrap_list = ["TransformerBlock1", "TransformerBlock2", ...]
model = simple_fsdp(model, wrap_modules=wrap_list)
model = torch.compile(model, fullgraph=True)
```

### Integration with Distributed Techniques
SimpleFSDP is composable with tensor parallel and pipeline parallel setups such as TorchTitan:
```python
from torch_titan import tensor_parallel, pipeline_parallel
model = meta_init(model)
model = model.half()
model = activation_checkpoint(model)
model = tensor_parallel(model, tp_degree=8)
model = pipeline_parallel(model, pp_degree=4)
model = simple_fsdp(model)
model = torch.compile(model, fullgraph=True)
```

*Composability with parallelism orchestrators is achieved with minimal API surface.*

## 7. Limitations and Prospective Directions

- Auto-wrapping employs an $\alpha + \beta n$ NCCL time model and greedy bucketing may overselect for memory-heavy buckets. Future work could employ global ILP or topology-aware models.
- Graph breaks, typically caused by Python data-dependent control flow, limit global comm/compute overlap. FX graph knitting is identified as a future research direction.
- Extension of DTensor mesh abstraction for cross-data-center or heterogeneous device meshes is unexplored.
- Potential integration with fully-automated planners (Alpa, nnScaler, Slapo) could exploit SimpleFSDP’s exposed computation-communication graph for joint data/tensor/pipeline planning and bubble minimization.

In summary, SimpleFSDP is a PyTorch-native re-implementation of ZeRO-3 fully sharded data parallelism, purpose-built for end-to-end compiler tracing, unlocking compiler-level optimizations such as IR bucketing and reordering, and delivering significant improvements in memory and throughput for large-scale distributed model training [2411.00284].

Source: https://www.emergentmind.com/topics/simplefsdp