FSDP and ZeRO-style sharding
Sharding reduces steady replicated state, but wrap boundaries decide the temporary memory and communication you pay to reconstruct it.
An eight-GPU job is slower than one GPU and loses progress on failure.
This chapter moves from set a sharding baseline to specify precision by role. We will make the idea concrete in a small code example before placing it back inside a production system with eight GPUs with restart under 15 minutes.
Declare parameter, reduction, buffer, and optimizer-state precision separately. Track overflow, convergence, and checkpoint format as part of the experiment.
Verify the behavior illustrated in fsdp_setup.py, then measure decision utility, slice reliability, and system latency.
One rank skips a batch and misses an all-reduce, leaving every other rank hung. An unsafe release must degrade to a named baseline or the last known-good version.
Build the idea one layer at a time.
Replication is wonderfully simple until model state fills every GPU. Sharding buys memory by making ranks assemble what they need just in time.
ZeRO-style stages progressively partition optimizer state, gradients, and parameters. Fully Sharded Data Parallelism applies the same central idea around module boundaries: all-gather a module’s parameters for computation, then reshard them; reduce-scatter distributes its gradients. Peak memory therefore depends on wrap granularity and the largest materialized unit, not only total model size.
Communication volume and lifetime replace replicated memory as the design pressure. Wrapping every tiny layer creates frequent collectives, while wrapping a huge stack creates a large peak all-gather. Mixed precision can reduce communication and compute cost, but the optimizer and reductions may still need higher precision for stability. Always construct the optimizer after wrapping so it references managed parameters.
- global batch
- FSDP
- tensor parallelism
- checkpoint resharding
Follow the reasoning before touching the code.
- 01
Set a sharding baseline
Measure plain DDP memory and time first. Apply full sharding with a clear module wrap boundary and verify loss/gradient parity on a small batch.
- 02
Balance gather size and frequency
Use profiler traces to locate all-gather and reduce-scatter stalls. Adjust wrap groups so communication can overlap substantial compute without creating a peak-memory spike.
- 03
Specify precision by role
Declare parameter, reduction, buffer, and optimizer-state precision separately. Track overflow, convergence, and checkpoint format as part of the experiment.
import os
import torch
import torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy
dist.init_process_group('nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = torch.nn.Sequential(
torch.nn.Linear(1024, 4096), torch.nn.GELU(),
torch.nn.Linear(4096, 1024)
)
precision = MixedPrecision(param_dtype=torch.bfloat16,
reduce_dtype=torch.float32,
buffer_dtype=torch.float32)
model = FSDP(model, device_id=local_rank, mixed_precision=precision,
sharding_strategy=ShardingStrategy.FULL_SHARD)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)Wrap a model with full parameter sharding and explicit mixed-precision roles.
What to notice
- The optimizer is created after FSDP wrapping, so it follows the parameters that FSDP manages.
- BF16 lowers parameter-compute traffic while FP32 reduction keeps more precision when gradients from ranks are combined.
YOUR TURN Pause and predict before revealing the answer.
Why can one very large FSDP wrap unit still run out of memory even though parameters are sharded at rest?
Need a hint?
Consider what must exist locally while that unit performs its forward or backward computation.
Check your reasoning
FSDP must materialize the wrap unit’s full parameters for computation. If the unit is too large, that temporary all-gather plus activations and buffers can exceed device memory.
Lessons 1, 2 in this course.
Connect the chapter to the course’s quantitative spine.
Sharding reduces steady replicated state, but wrap boundaries decide the temporary memory and communication you pay to reconstruct it. The scaling efficiency below is the shared course-level lens that helps quantify the surrounding system. Name each symbol, check its units or shape, and use only values that were available at the production cutoff.
Scaling efficiency
| Symbol | Meaning / shape / unit |
|---|---|
W | worker/GPU count |
throughput_W | distributed samples per second |
throughput_1 | single-worker baseline |
Open derivation and numerical substitution
Start from the production quantity being optimized, substitute the observed values with their declared units, then isolate the model-controlled term. Preserve shape annotations at each step so broadcasting or aggregation cannot silently change the result.
- Write the named inputs: W, throughput_W, throughput_1.
- Substitute one small, hand-checkable batch before vectorizing.
- Calculate an independent reference value and compare within a declared tolerance.
# equation → code contract
inputs = validate_shapes_and_units(batch)
value = compute_c13(inputs)
assert is_finite(value)Change one assumption and make the tradeoff visible.
Use the worked chapter as your prediction, then change one input in the lab and explain why the output moves. The lab runs predefined TypeScript and never executes learner code, so the focus stays on the concept rather than setup.
All-reduce scaling lab
Add workers and estimate delivered speedup after communication and stragglers.
Assumption: 120 ms local-compute step; ring communication and logarithmic straggler penalty.
Open nonvisual data table
| Item | Computed state | Interpretation |
|---|---|---|
| 1 workers | 120.3 ms | 1.0× |
| 2 workers | 125.7 ms | 1.9× |
| 4 workers | 131.4 ms | 3.7× |
| 8 workers | 137.8 ms | 7.0× |
| 16 workers | 145.6 ms | 13.2× |
Trace the complete operating path.
Carry the chapter’s central idea into the system: Sharding reduces steady replicated state, but wrap boundaries decide the temporary memory and communication you pay to reconstruct it. The path below shows where that reasoning needs an artifact, a test, an owner, or a fallback.
- 01
Turn “Set a sharding baseline” into a written invariant and a failing test.
- 02
Promote fsdp_setup.py from illustration to a versioned implementation with representative fixtures.
- 03
Use “Specify precision by role” as the release checkpoint: Declare parameter, reduction, buffer, and optimizer-state precision separately. Track overflow, convergence, and checkpoint format as part of the experiment.
- 04
Serve or materialize behind explicit eight GPUs with restart under 15 minutes.
- 05
Join telemetry to mature outcomes and retain a rollback path.
Observability
Join service health, input quality, prediction distributions, slice behavior, and mature outcomes by exact version.
Cost
Measure storage, preprocessing, compute, queueing, and human review under a representative arrival pattern.
Failure modes
One rank skips a batch and misses an all-reduce, leaving every other rank hung. Add a detector, owner, mitigation, and stop condition for this class of failure.
Alternatives
Compare a rule, a simpler statistical baseline, and a different system boundary before adding model complexity.
Explain the contract, not just the vocabulary.
Fault-tolerant distributed trainer
Specify DDP correctness, scaling experiments, worker failure, and complete checkpoint recovery.
- Gradient/data equivalence
- 1/2/4-GPU efficiency report
- Kill-and-resume proof
- Hang/straggler runbook
Read primary material with a purpose.
Return to the opening failure.
Log per-rank data/compute/collective time, preserve complete training state, and keep a one-GPU equivalence test.
For this lesson, the release evidence is the worked code example, a hand-checked formal result, deterministic simulation output, a ≥80% checkpoint, the production rubric, and a named fallback. The course resolves when the system can produce efficient ddp, parallelism choices, elastic sampling, and resumable checkpoints.