FSDP › Code Implementation
Language AI HandbookAll chapters FSDP › Code ImplementationOne section · Sections in this chapter9 / 13 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
This section walks through a complete FSDP training setup, starting from a small model definition and building up to a full distributed training loop.
Setup and Imports
We start with the necessary imports and a simple model definition.
Defining a Simple Transformer Model
We define a small transformer model with a configurable number of layers and dimensions. This is purely for demonstration; in practice you would use a model from Hugging Face or your own architecture.
import torch.nn as nn
class TransformerLayer(nn.Module):
"""A single transformer encoder layer."""
def __init__(self, d_model=256, nhead=4, dim_feedforward=1024):
super().__init__()
self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
self.ff = nn.Sequential(
nn.Linear(d_model, dim_feedforward),
nn.GELU(),
nn.Linear(dim_feedforward, d_model),
)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x):
attn_out, _ = self.self_attn(x, x, x)
x = self.norm1(x + attn_out)
x = self.norm2(x + self.ff(x))
return x
class SimpleTransformer(nn.Module):
"""A stack of transformer layers with an embedding and output head."""
def __init__(
self,
vocab_size=10000,
d_model=256,
nhead=4,
dim_feedforward=1024,
num_layers=6,
):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.layers = nn.ModuleList(
[
TransformerLayer(d_model, nhead, dim_feedforward)
for _ in range(num_layers)
]
)
self.output_head = nn.Linear(d_model, vocab_size)
def forward(self, input_ids):
x = self.embedding(input_ids)
for layer in self.layers:
x = layer(x)
return self.output_head(x)Counting Parameters
Let's inspect the parameter count to understand the memory requirements.
def count_parameters(model):
total = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
return total, trainable
model = SimpleTransformer(
vocab_size=10000,
d_model=256,
nhead=4,
dim_feedforward=1024,
num_layers=6,
)
total_params, trainable_params = count_parameters(model)Total parameters: 9,868,560 Trainable parameters: 9,868,560 Memory (FP32): 39.5 MB Memory (with Adam): 157.9 MB
The output shows parameter count and estimated memory footprint. The multiplier of 16 accounts for parameters (4 bytes), gradients (4 bytes), and Adam optimizer states (8 bytes: first moment plus second moment, both in float32).
FSDP Initialization
This function wraps a model with FSDP. In a real distributed job, it must be called after dist.init_process_group() so that the process group is available for communication.
import functools
from torch.distributed.fsdp import (
FullyShardedDataParallel as FSDP,
)
from torch.distributed.fsdp import (
MixedPrecision,
ShardingStrategy,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
def wrap_model_with_fsdp(model, sharding_strategy=ShardingStrategy.FULL_SHARD):
"""
Wrap a transformer model with FSDP.
Parameters:
- model: The PyTorch model to wrap
- sharding_strategy: The FSDP sharding strategy to use
Returns the FSDP-wrapped model.
"""
# Define the auto-wrap policy to wrap each TransformerLayer separately
auto_wrap_policy = functools.partial(
transformer_auto_wrap_policy,
transformer_layer_cls={TransformerLayer},
)
# Mixed precision configuration
mp_policy = MixedPrecision(
param_dtype=torch.float16,
reduce_dtype=torch.float32,
buffer_dtype=torch.float16,
)
# Wrap the model
fsdp_model = FSDP(
model,
auto_wrap_policy=auto_wrap_policy,
sharding_strategy=sharding_strategy,
mixed_precision=mp_policy,
device_id=torch.cuda.current_device(),
)
return fsdp_modelSimulating Memory Usage per Strategy
Since we cannot run distributed training in this notebook, we can compute the theoretical memory savings analytically for different sharding strategies.
def memory_per_gpu(total_params, num_gpus, strategy):
"""
Compute estimated GPU memory in bytes for each sharding strategy.
Parameters:
- total_params: total number of model parameters
- num_gpus: number of GPUs in the training job
- strategy: one of 'full_shard', 'shard_grad_op', 'no_shard'
Returns a dict with breakdown by component.
"""
bytes_per_param = 4 # float32
if strategy == "no_shard":
params = total_params * bytes_per_param
grads = total_params * bytes_per_param
optimizer = total_params * 2 * bytes_per_param # Adam: 2x
elif strategy == "shard_grad_op":
params = total_params * bytes_per_param # not sharded
grads = total_params * bytes_per_param / num_gpus
optimizer = total_params * 2 * bytes_per_param / num_gpus
elif strategy == "full_shard":
params = total_params * bytes_per_param / num_gpus
grads = total_params * bytes_per_param / num_gpus
optimizer = total_params * 2 * bytes_per_param / num_gpus
else:
raise ValueError(f"Unknown strategy: {strategy}")
total = params + grads + optimizer
return {
"params_mb": params / 1e6,
"grads_mb": grads / 1e6,
"optimizer_mb": optimizer / 1e6,
"total_mb": total / 1e6,
}
strategies = ["no_shard", "shard_grad_op", "full_shard"]
num_gpus_options = [1, 4, 8, 16]Strategy GPUs Params (MB) Grads (MB) Optim (MB) Total (MB) -------------------------------------------------------------------------------- no_shard 1 39.5 39.5 78.9 157.9 no_shard 4 39.5 39.5 78.9 157.9 no_shard 8 39.5 39.5 78.9 157.9 no_shard 16 39.5 39.5 78.9 157.9 shard_grad_op 1 39.5 39.5 78.9 157.9 shard_grad_op 4 39.5 9.9 19.7 69.1 shard_grad_op 8 39.5 4.9 9.9 54.3 shard_grad_op 16 39.5 2.5 4.9 46.9 full_shard 1 39.5 39.5 78.9 157.9 full_shard 4 9.9 9.9 19.7 39.5 full_shard 8 4.9 4.9 9.9 19.7 full_shard 16 2.5 2.5 4.9 9.9
The table reveals the progressive memory reduction across strategies. With full_shard on 16 GPUs, each GPU holds just 1/16 of each component, a 16x reduction in steady-state memory. Notice that shard_grad_op achieves a substantial reduction too because the optimizer states dominate: even though parameters are replicated, sharding the optimizer states ( bytes) alone saves significantly more memory than sharding parameters alone ( bytes).
Visualizing Memory Savings
Visualizing Communication Volume
With FULL_SHARD, each forward pass triggers an all-gather for every FSDP unit. The total communication volume per training step is larger than with standard DDP. Let's quantify this difference.
The communication overhead is the fundamental trade-off of FSDP: you gain memory capacity in exchange for more communication. Whether this trade-off is favorable depends entirely on your hardware's interconnect speed relative to compute speed. On a cluster where GPUs are connected by slow Ethernet, the extra communication turns a compute-bound problem into a communication-bound one. On an NVLink-connected server, the all-gather operations happen at memory speed and the overhead is negligible.
def communication_volume_gb(total_params, num_gpus, num_fsdp_units, strategy):
"""
Estimate communication volume per training step in GB.
For FULL_SHARD:
- Forward: num_fsdp_units all-gathers, each moving total_params/num_fsdp_units params
- Backward: num_fsdp_units all-gathers + num_fsdp_units reduce-scatters
For SHARD_GRAD_OP:
- No all-gathers during forward/backward
- One reduce-scatter at end of backward
For NO_SHARD (DDP):
- One all-reduce at end of backward
"""
bytes_per_param = 4
if strategy == "full_shard":
# Forward: M all-gathers, each reconstructing P/M params across N GPUs
# Communication volume per all-gather = P * (N-1)/N * 2 (send + receive)
# Simplified as P * 2 bytes (one full copy across ring)
fwd_volume = (
total_params * bytes_per_param * 2
) # all-gathers in forward
# Backward: M all-gathers + M reduce-scatters
bwd_volume = (
total_params * bytes_per_param * 2
+ total_params * bytes_per_param * 2
)
total_volume = fwd_volume + bwd_volume
elif strategy == "shard_grad_op":
# Only one reduce-scatter at end of backward
total_volume = total_params * bytes_per_param * 2
elif strategy == "no_shard":
# All-reduce = 2 * total_params * bytes (ring algorithm)
total_volume = total_params * bytes_per_param * 2
else:
raise ValueError(f"Unknown strategy: {strategy}")
return total_volume / 1e9 # GBCommunication volume per training step (model with 9,868,560 params): Strategy Volume (GB) ----------------------------------- no_shard 0.0789 GB shard_grad_op 0.0789 GB full_shard 0.2368 GB
The output illustrates the communication cost trade-off: FULL_SHARD incurs more communication per step because of the all-gather operations during both forward and backward passes. For this small demonstration model, the absolute volumes are tiny. For a 70B model on 8 GPUs, FULL_SHARD would move approximately TB per step (before counting the reduce-scatters), while SHARD_GRAD_OP and NO_SHARD would each move about GB. This is why interconnect bandwidth is the deciding factor: on high-bandwidth networks, even the larger communication volumes complete in a fraction of the compute time.
Visualizing Communication vs. Memory Trade-off
Sharding Strategy Selection Guide
To help you choose the right strategy for a given scenario, here is a decision framework based on model size and hardware.
def recommend_strategy(model_params, gpu_memory_gb, num_gpus):
"""
Recommend an FSDP sharding strategy based on model and hardware.
Parameters:
- model_params: number of model parameters
- gpu_memory_gb: memory per GPU in GB
- num_gpus: number of available GPUs
Returns a strategy recommendation with reasoning.
"""
# Memory requirements in GB (FP32 + Adam)
mem_no_shard_gb = model_params * 16 / 1e9
mem_shard_grad_op_gb = (
model_params * (4 + 4 / num_gpus + 8 / num_gpus) / 1e9
)
mem_full_shard_gb = model_params * 16 / (num_gpus * 1e9)
results = {
"NO_SHARD fits": mem_no_shard_gb <= gpu_memory_gb,
"SHARD_GRAD_OP fits": mem_shard_grad_op_gb <= gpu_memory_gb,
"FULL_SHARD fits": mem_full_shard_gb <= gpu_memory_gb,
"mem_no_shard_gb": mem_no_shard_gb,
"mem_shard_grad_op_gb": mem_shard_grad_op_gb,
"mem_full_shard_gb": mem_full_shard_gb,
}
if results["NO_SHARD fits"]:
recommendation = "NO_SHARD (use DDP for simplicity)"
elif results["SHARD_GRAD_OP fits"]:
recommendation = "SHARD_GRAD_OP (moderate savings, less communication)"
elif results["FULL_SHARD fits"]:
recommendation = "FULL_SHARD (maximum memory reduction)"
else:
recommendation = "FULL_SHARD + CPU offload (model may still not fit)"
results["recommendation"] = recommendation
return resultsScenario: 7B model, 8x A100 80GB NO_SHARD: 112 GB/GPU -> fits: False SHARD_GRAD_OP: 38.5 GB/GPU -> fits: True FULL_SHARD: 14.0 GB/GPU -> fits: True Recommendation: SHARD_GRAD_OP (moderate savings, less communication) Scenario: 13B model, 8x A100 40GB NO_SHARD: 208 GB/GPU -> fits: False SHARD_GRAD_OP: 71.5 GB/GPU -> fits: False FULL_SHARD: 26.0 GB/GPU -> fits: True Recommendation: FULL_SHARD (maximum memory reduction) Scenario: 70B model, 8x A100 80GB NO_SHARD: 1120 GB/GPU -> fits: False SHARD_GRAD_OP: 385.0 GB/GPU -> fits: False FULL_SHARD: 140.0 GB/GPU -> fits: False Recommendation: FULL_SHARD + CPU offload (model may still not fit) Scenario: 70B model, 32x A100 80GB NO_SHARD: 1120 GB/GPU -> fits: False SHARD_GRAD_OP: 306.2 GB/GPU -> fits: False FULL_SHARD: 35.0 GB/GPU -> fits: True Recommendation: FULL_SHARD (maximum memory reduction) Scenario: 175B model, 64x A100 80GB NO_SHARD: 2800 GB/GPU -> fits: False SHARD_GRAD_OP: 732.8 GB/GPU -> fits: False FULL_SHARD: 43.8 GB/GPU -> fits: True Recommendation: FULL_SHARD (maximum memory reduction)
The scenarios illustrate how FULL_SHARD becomes necessary as models grow. A 70B model still needs CPU offloading or more than 8 GPUs of 80 GB each because even full sharding leaves 140 GB of model state per GPU. At 32 GPUs, FULL_SHARD reduces that footprint to 35 GB per GPU. The 175B model on 64 GPUs ( GB per GPU for the model state) fits within 80 GB with FULL_SHARD, leaving room for activation buffers.
Visualizing Strategy Recommendations
Complete Training Loop
Here is a self-contained training loop that demonstrates how FSDP fits into the standard PyTorch training workflow. In a real distributed setting, this function would run on each GPU process.
def create_synthetic_batch(batch_size=4, seq_len=128, vocab_size=10000):
"""Generate a synthetic batch of token IDs and targets."""
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
targets = torch.randint(0, vocab_size, (batch_size, seq_len))
return input_ids, targets
def compute_loss(logits, targets):
"""Cross-entropy loss for language modeling."""
# logits: (batch, seq_len, vocab_size) -> reshape for cross-entropy
batch_size, seq_len, vocab_size = logits.shape
return nn.functional.cross_entropy(
logits.view(batch_size * seq_len, vocab_size),
targets.view(batch_size * seq_len),
)def simulate_training_step(model, optimizer, num_steps=20, seed=42):
"""
Overfit one fixed synthetic batch without distributed setup.
In practice this runs on each GPU process with an FSDP-wrapped model.
"""
torch.manual_seed(seed)
model.train()
losses = []
input_ids, targets = create_synthetic_batch()
for step in range(num_steps):
optimizer.zero_grad()
logits = model(input_ids)
loss = compute_loss(logits, targets)
loss.backward()
optimizer.step()
losses.append(loss.item())
return losses
# Run simulation on the (non-distributed) model for demonstration
optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
training_losses = simulate_training_step(model, optimizer, num_steps=30)Training completed: 30 steps Initial loss: 9.3501 Final loss: 1.9261 Loss reduction: 79.4%
The training loop is identical in structure whether you are using a plain model, DDP, or FSDP. The only differences in a real FSDP job are the process group initialization, the FSDP() wrapping call, and the checkpoint saving procedure. This is a deliberate design choice: FSDP should be a drop-in replacement for DDP in the training loop, with complexity hidden in the wrapper.
Training Loss Visualization
Checkpointing with FSDP
Checkpointing is more complex with FSDP than with standard models because parameters are sharded across GPUs. In a standard PyTorch model, model.state_dict() returns a dictionary of all parameter tensors. With FSDP, each GPU only holds shards, so a naive state_dict() call would return a dictionary of shards rather than complete tensors. PyTorch provides two checkpointing modes to handle this correctly.
Full state dict mode gathers all parameter shards to rank 0 before saving, producing a standard checkpoint file identical to what you would get from a non-sharded model. The checkpoint is universally loadable: you can load it into a non-FSDP model, into an FSDP model with a different number of GPUs, or into any framework that understands standard PyTorch checkpoints. The downside is that rank 0 must temporarily hold the entire model in memory to write the checkpoint, creating a brief memory spike.
Local state dict mode saves each GPU's local parameter shard separately. This is faster and avoids memory pressure on rank 0, but the resulting checkpoint files are only loadable with the exact same FSDP configuration (same number of GPUs, same wrapping policy). This mode is useful for frequent checkpointing during training where you care about training continuity rather than portability.
# Full state dict checkpointing (rank 0 saves, others wait)
# In a real distributed job:
#
# from torch.distributed.fsdp import FullStateDictConfig, StateDictType
#
# cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
# with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, cfg):
# state = fsdp_model.state_dict()
# if dist.get_rank() == 0:
# torch.save(state, "checkpoint.pt")
# For this notebook, demonstrate saving the non-distributed model
checkpoint = {
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"step": len(training_losses),
}
# torch.save(checkpoint, "model_checkpoint.pt") # commented out to avoid file creationCheckpoint contents: model_state: dict with 75 tensors optimizer_state: dict with 0 tensors step: 30
The offload_to_cpu=True flag in the full state dict config is important: it moves gathered tensors to CPU memory as they are assembled, preventing the GPU memory from doubling up. Without it, rank 0 would need to hold both its shard (from normal FSDP operation) and the full model (being assembled for saving) simultaneously in GPU memory.
Key Parameters
The key FSDP configuration parameters are:
- sharding_strategy: Controls what is sharded.
FULL_SHARDfor maximum memory reduction,SHARD_GRAD_OPfor moderate savings with less communication,NO_SHARDto disable sharding. - auto_wrap_policy: Determines how submodules are grouped into FSDP units. Use
transformer_auto_wrap_policyfor transformer architectures. - mixed_precision: A
MixedPrecisionobject specifying dtypes for parameters, gradients, and buffers. BF16 is the standard choice for large model training. - device_id: The GPU to place the model on. Should be
torch.cuda.current_device()after setting the appropriate device. - cpu_offload: A
CPUOffloadobject that, when enabled, keeps parameters on CPU and only moves them to GPU when needed. Reduces GPU memory further at the cost of PCIe bandwidth.