Personalize this lesson
Adapt explanations and teaching visuals to your background and preferred voice.
A model can fit its weights on one GPU and still run out of memory at the first optimizer update. A 1.5-billion-parameter model needs 3 GB for two-byte weights alone. In the ZeRO paper's mixed-precision Adam accounting, training state takes 24 GB before activations or temporary buffers. Where did the other 21 GB come from?[1]
The Megatron lesson split computation across parallel axes. DeepSpeed also asks who must retain each tensor between computations. Its Zero Redundancy Optimizer (ZeRO) partitions optimizer state, gradients, and eventually parameters across data-parallel ranks. Each rank still processes its own examples; sharding state doesn't mean sharding the examples' computation as tensor parallelism would.
A saved directory can still look complete and resume a different update. The source tour uses commit c455031422641a825588c926455458c82b27b6c0 (August 6, 2026, before the 0.19.4 release). On October 10, 2026, the latest training documentation identifies itself as 0.19.8 and the stable ZeRO documentation as 0.19.7. Treat these as distinct version references; validate each feature combination against the source and documentation you deploy.[2][3][4]
The marked examples run on a CPU with Python's standard library and check ownership arithmetic plus an Adam update. Validate the corresponding behavior separately with DeepSpeed kernels, distributed execution, and GPU memory measurements.
Count state before choosing a stage
Let be the parameter count represented by a data-parallel group. If tensor or pipeline parallelism already splits the model, use the relevant local model partition, not the entire model's parameter count.
The paper's 16-byte recipe is:
| Tensor | Bytes per parameter | Role |
|---|---|---|
| Low-precision parameter | 2 | Forward and backward computation |
| Low-precision gradient | 2 | Gradient awaiting the update |
| FP32 master parameter | 4 | Higher-precision update state |
| FP32 first moment | 4 | Adam's running gradient average |
| FP32 second moment | 4 | Running average of squared gradients |
| Total | 16 | Model state only |
The second moment isn't an estimate of centered variance. Also, BF16 compute doesn't by itself determine the gradient-accumulation dtype or whether a master copy exists. For example, replacing the two-byte gradient with a four-byte accumulation array changes this simplified total to 18 bytes. Inspect the actual optimizer and buffer dtypes. DeepSpeed exposes data_types.grad_accum_dtype separately.[5]
Let , , and denote bytes per parameter for compute weights, gradients, and optimizer-plus-master state. For data-parallel ranks, idealized per-rank storage is:
With and , those are , , , and bytes. The formulas describe balanced persistent state, not peak allocator usage. Padding, persistent small parameters, gathered tensors, activations, communication buffers, workspaces, and fragmentation can add memory.[1]
Use exact fractions to compare recipes without confusing GB ( bytes) and GiB ( bytes). This estimator deliberately excludes padding and live working sets.
from fractions import Fraction
def state_bytes(parameters, ranks, param=2, grad=2, optimizer=12):
values = (parameters, ranks, param, grad, optimizer)
if any(type(x) is not int or x <= 0 for x in values):
raise ValueError("counts and byte widths must be positive integers")
shard = Fraction(parameters, ranks)
return {
"DP": parameters * (param + grad + optimizer),
"ZeRO-1": parameters * (param + grad) + shard * optimizer,
"ZeRO-2": parameters * param + shard * (grad + optimizer),
"ZeRO-3": shard * (param + grad + optimizer),
}
for name, size in state_bytes(1_500_000_000, 4).items():
print(f"{name:6s} {float(size / 10**9):5.2f} GB | {float(size / 2**30):5.2f} GiB")
assert set(state_bytes(17, 1).values()) == {17 * 16}
assert state_bytes(1_000_000_000, 4, grad=4)["ZeRO-3"] == 4_500_000_000DP 24.00 GB | 22.35 GiB
ZeRO-1 10.50 GB | 9.78 GiB
ZeRO-2 8.25 GB | 7.68 GiB
ZeRO-3 6.00 GB | 5.59 GiBStage 1 partitions optimizer state and master weights. Stage 2 also partitions reduced gradients. Stage 3 also partitions compute parameters.[4] Higher stages aren't quality levels. If stage 2 fits, stage 3's additional gathers may or may not pay off. If saved activations dominate, investigate activation checkpointing, microbatch size, or sequence length before expecting parameter sharding to solve the problem.
Partitioning the bytes does not settle Adam's second moment. That depends on whether each rank updates from its own gradient or from the reduced one.
Reduce first, then update the owned slice
It doesn't if each rank updates from its own gradient. The square of a mean isn't the mean of squares, which is easiest to see with two ranks. Suppose their gradients for one parameter are 1 and 3. The averaged gradient is 2. The owner applies Adam once to that gradient and its local moment state. It mustn't apply Adam separately to each rank's gradient and then average the updated moments: the second moment involves a square, and , while .[6]
In an unsharded run, every rank can hold the reduced gradient and optimizer state. In a sharded run, the owning rank holds the relevant slice. The mathematical update can match, although different reduction orders and low-precision arithmetic can change floating-point results.
This scalar implementation runs five updates on eleven parameters, partitioned over four logical owners. The comparison covers every weight and moment, including uneven partitions. The reference uses ordinary Adam without weight decay, not the later AdamW configuration, and doesn't simulate NCCL's execution or DeepSpeed's padding layout.
from math import isclose, sqrt
def adam_step(weight, first, second, gradient, step, lr=0.01):
first = 0.9 * first + 0.1 * gradient
second = 0.999 * second + 0.001 * gradient * gradient
corrected_first = first / (1 - 0.9**step)
corrected_second = second / (1 - 0.999**step)
weight -= lr * corrected_first / (sqrt(corrected_second) + 1e-8)
return weight, first, second
def compare_owned_updates(parameters, ranks, steps):
if any(type(x) is not int or x <= 0 for x in (parameters, ranks, steps)):
raise ValueError("parameters, ranks, and steps must be positive integers")
reference = [(0.1 * i, 0.0, 0.0) for i in range(parameters)]
# Contiguous ownership; some owners may be empty when ranks > parameters.
owners = [range(parameters * r // ranks, parameters * (r + 1) // ranks)
for r in range(ranks)]
sharded = [{i: reference[i] for i in ids} for ids in owners]
assert sorted(i for owner in sharded for i in owner) == list(range(parameters))
for step in range(1, steps + 1):
local = [[((i + 2 * r + step) % 7 - 3) / 10
for i in range(parameters)] for r in range(ranks)]
mean = [sum(local[r][i] for r in range(ranks)) / ranks
for i in range(parameters)]
reference = [adam_step(*state, mean[i], step)
for i, state in enumerate(reference)]
for rank, owner in enumerate(sharded):
# A reference reduce-scatter: reduce across ranks, retain owned indices.
reduced = {i: sum(row[i] for row in local) / ranks for i in owner}
sharded[rank] = {i: adam_step(*state, reduced[i], step)
for i, state in owner.items()}
reconstructed = [state for owner in sharded for state in owner.values()]
assert all(isclose(a, b, rel_tol=1e-12, abs_tol=1e-12)
for expected, actual in zip(reference, reconstructed)
for a, b in zip(expected, actual))
return reference
result = compare_owned_updates(11, 4, 5)
# Independently check the first-step scalar update for gradient 2.
w, m, v = adam_step(1.0, 0.0, 0.0, 2.0, 1)
assert isclose(w, 1.0 - 0.01 * 2.0 / (2.0 + 1e-8))
assert isclose(m, 0.2) and isclose(v, 0.004)
assert (1**2 + 3**2) / 2 != ((1 + 3) / 2)**2
print(f"{len(result)} parameters, 4 owners, 5 updates: weights and moments match")11 parameters, 4 owners, 5 updates: weights and moments matchThe example averages equally weighted rank gradients. Variable microbatch sizes or different numbers of unmasked tokens need the appropriate loss weighting; “average each rank's average” isn't automatically the global token-average objective. If one rank averages one token's gradient of 1 and another averages three token gradients of 3, their rank-average is 2. The four-token average is . Sharding can preserve the update you specify, but it can't repair a wrongly weighted objective. The arithmetic is settled for one parameter. Which object in the training loop is allowed to apply it?
Let the engine own the microstep
That object is the engine. deepspeed.initialize(...) selects and configures a DeepSpeedEngine, wrapped optimizer, optional data loader, and scheduler. The engine coordinates backward, accumulation, reduction, precision, stepping, and checkpoint state.[2]
Here is an integration sketch, not a standalone program. It assumes a model whose forward returns a scalar mean loss, a prepared batch on the correct device, and a distributed launcher. Under the default managed-accumulation behavior, call step() after every microbatch; the engine applies an optimizer update at its accumulation boundary. Don't divide the loss by the accumulation count again when using the default engine scaling.[3]
engine, optimizer, _, scheduler = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
config="ds_config.json",
)
for batch in prepared_batches:
loss = engine(batch)
engine.backward(loss)
engine.step()With accumulation set to eight, this loop calls backward() and step() eight times for one optimizer update. Calling step() only on every eighth batch under the default managed behavior would advance the engine's microstep counter only once. Dividing each loss by eight as well would instead scale an already-averaged gradient a second time.
The current API also offers managed_gradient_accumulation: false. In that mode, each step() completes one update and the caller chooses when to call it. If the actual accumulation count differs from the configured count, use backward(..., scale_wrt_gas=False) and supply the intended loss weighting yourself. The pinned August source rejects unmanaged accumulation with optimizer or parameter offload, while the current training documentation supports those paths. Check the installed engine before combining them.[2][3]
This four-data-rank configuration uses BF16 model parameters and explicitly requests FP32 gradient accumulation. Validate it on compatible GPUs before adopting it. The optimizer settings belong to your training recipe; ZeRO doesn't determine a suitable learning rate.
{
"train_batch_size": 64,
"train_micro_batch_size_per_gpu": 2,
"gradient_accumulation_steps": 8,
"bf16": { "enabled": true },
"data_types": { "grad_accum_dtype": "fp32" },
"zero_optimization": {
"stage": 2,
"contiguous_gradients": true,
"overlap_comm": true,
"reduce_bucket_size": 200000000
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": 0.00002,
"betas": [0.9, 0.999],
"eps": 1e-8,
"weight_decay": 0.01
}
}
}Check two different units. The batch identity is samples per optimizer update. The last factor is the data-parallel degree, not necessarily the total number of GPUs when model parallelism is present. Holding all three batch fields fixed while launching eight data ranks is inconsistent, not a harmless scale-out change.[5]
reduce_bucket_size is a count of elements, not bytes. Two hundred million elements occupy 400 MB at two bytes each or 800 MB at four bytes each, before other buffers. Communication dtype and accumulation dtype needn't match. Larger buckets can amortize collective overhead but increase temporary storage and delay the first reduction; inspect the timeline rather than copying a bucket size.[4]
For a real launch, record the installed DeepSpeed and PyTorch versions, accelerator and driver, the resolved configuration, and ds_report output. The report checks operator/environment compatibility. Run a multi-rank training step and exercise the offload path on the target setup.[7] Neither the integration sketch nor this GPU configuration was executed for this lesson. A config that fits the persistent estimate still has to gather a full matrix from the shards each rank stores.
Stage 3 needs a live working set
A sharded parameter can't directly supply a conventional full-layer matrix multiplication. ZeRO-3 gathers the parameter view needed for a module, executes it, and can release that full view afterward. Backward can require another gather. Prefetch, persistence, reuse, and module grouping affect the actual schedule.[4]
Don't subtract the local shard from a gathered allocation unless the implementation actually aliases that storage. In the pinned source's sequential gather path, param.ds_tensor is the input shard and a separate torch.empty(...) tensor receives the full padded gather.[2] Logical values and allocated bytes are different counts.
For four tensors of eight parameters on four ranks, one rank retains eight shard elements in total. A separate full eight-element gather brings the parameter-storage count to sixteen, not fourteen. Keeping a second gathered tensor for prefetch raises this example to twenty-four. These counts exclude gradients, optimizer state, activations, and communication-library workspace.
This storage model rounds each tensor's shard up to a whole element and allocates a separate padded gather buffer. To estimate peak memory for a run, include its DeepSpeed stage, gather schedule, prefetch policy, gradients, optimizer state, activations, and communication workspace.
def parameter_storage(layer_sizes, ranks, live_layers):
if type(ranks) is not int or ranks <= 0:
raise ValueError("ranks must be a positive integer")
if not layer_sizes or any(type(n) is not int or n <= 0 for n in layer_sizes):
raise ValueError("tensor sizes must be positive integers")
if (any(type(i) is not int or not 0 <= i < len(layer_sizes) for i in live_layers)
or len(set(live_layers)) != len(live_layers)):
raise ValueError("live tensor indices must be distinct and in range")
partitions = [(n + ranks - 1) // ranks for n in layer_sizes]
resident = sum(partitions)
gathered = sum(partitions[i] * ranks for i in live_layers)
return resident, gathered, resident + gathered
sizes = [8, 8, 8, 8]
for live in ([], [0], [0, 1]):
resident, gathered, total = parameter_storage(sizes, 4, live)
print(f"{len(live)} gathered: shards {resident} + full buffers {gathered} = {total} elements")
assert parameter_storage([5, 9, 2], 4, [1]) == (6, 12, 18)
assert parameter_storage(sizes, 4, [0])[2] == 16
assert parameter_storage(sizes, 4, [0, 1])[2] == 240 gathered: shards 8 + full buffers 0 = 8 elements
1 gathered: shards 8 + full buffers 8 = 16 elements
2 gathered: shards 8 + full buffers 16 = 24 elementsThe largest module's forward/backward working set must still fit. deepspeed.zero.Init() can partition parameters during construction so initialization doesn't first materialize the entire model on every GPU, but it doesn't make an oversized live operator fit. Tiling or model parallelism may still be needed.[4]
External parameters and dynamic routing are different problems
A tied embedding reused outside its owning module is an external parameter. DeepSpeed can discover common cases automatically; otherwise use deepspeed.zero.register_external_parameter() to identify the consuming module, or deepspeed.zero.GatheredParameters for an explicit access scope. This coordinates access to an existing parameter.[4]
Dynamic routing is different. Two data ranks may select different existing experts and enter incompatible gather sequences. Marking the enclosing module as a ZeRO leaf gathers its descendants as a unit before executing that module, at the cost of a larger working set. Set the boundary with deepspeed.utils.set_z3_leaf_modules() before engine initialization, or the zero_optimization.leaf_module configuration. This doesn't register newly created trainable parameters with an already-built optimizer.[3]
Communication volume isn't elapsed time
A simplified ring model makes the paper's communication comparison concrete. For communicated elements per full vector on ranks, one all-gather or reduce-scatter sends elements per rank. Count sends or receives consistently, not their sum on one side of the comparison.
With equal element widths and no extra gathers, the basic schedules are:
| Schedule | Collective work | Sent elements per rank |
|---|---|---|
| Replicated data parallel | Gradient reduce-scatter + all-gather | |
| ZeRO-2 | Gradient reduce-scatter + updated-parameter all-gather | |
| ZeRO-3 | Forward parameter gather + backward parameter gather + gradient reduce-scatter |
This explains the paper's approximately versus memory model for large .[1] Step time also depends on gradient/parameter dtypes, repeated gathers during accumulation or recomputation, parameter persistence, bucket sizes, overlap, and collective algorithms. In stage 1, gradient storage stays replicated; inspect the implementation before inferring a particular reduce-scatter path from its ownership formula.
An accumulation window also changes what “per step” means. By default, ZeRO-2 and ZeRO-3 reduce and partition gradients on each backward call, even though Adam waits for the final microbatch. Keeping only the owned gradient between calls saves memory but can repeat communication. Don't wrap those stages in a DDP-style engine.no_sync(): the engine rejects that combination. The current engine.coalesce_grad_reduction() API can defer reduction across several backward calls, with a full local gradient copy retained until the scope ends and restrictions on the engine, precision wrapper, and accumulation window. Treat it as a separate schedule to validate, not a free communication reduction.[3]
On NVIDIA GPUs, NCCL commonly executes these collectives. DeepSpeed chooses the groups and schedules their use. A slow gather could be an exposed dependency, a small-message problem, or a network route problem. Read the per-rank timeline and communication counters together. Volume and time still assume the state lives on the GPU. What changes when it doesn't fit there at all?
Offload exchanges memory pressure for movement
Then the state has to move. ZeRO-Offload moved optimizer-related state and work to CPU memory. ZeRO-Infinity extends the hierarchy to CPU and NVMe storage.[8][9] In current configuration, optimizer offload is available for stages 1 through 3, parameter offload for stage 3. Optimizer computation remains on the CPU even when its state is stored on NVMe.[4]
This fragment shows a stage-3 choice, not a complete launch configuration. It requires suitable local storage, available space, compatible I/O operators, and host memory for staging.
{
"zero_optimization": {
"stage": 3,
"offload_optimizer": { "device": "cpu", "pin_memory": true },
"offload_param": {
"device": "nvme",
"nvme_path": "/local_nvme/deepspeed"
}
}
}A bandwidth bound is a useful first check. If a step must move 12 GB through a bottleneck sustaining 6 GB/s, that path needs at least two seconds of service time. Perfect overlap with three seconds of independent compute could hide it; a dependency requiring the data before compute can't. Multiple reads/writes and shared-link contention can make the real step slower. Use these supplied values to check the arithmetic, then measure the target NVMe path.
Measure GPU memory, host memory, pinned-memory use, CPU optimizer time, transfer volume, and storage queueing. Compare tokens per second at the same effective batch and sequence lengths. A capacity fit still needs throughput and time-to-completion measurements. A JSON key can name the offload device and still hide which function gathers a parameter or refuses a pipeline combination.
Trace how the engine uses each config value
Follow the call, not the key name. The pinned source organizes the mechanisms as follows:[2]
| Entry point | What to inspect |
|---|---|
deepspeed/__init__.py | Initialization and engine selection |
deepspeed/runtime/engine.py | Backward, accumulation boundaries, step, save/load |
deepspeed/runtime/zero/stage_1_and_2.py | Gradient and optimizer ownership |
deepspeed/runtime/zero/stage3.py | Sharded optimizer and parameter coordination |
deepspeed/runtime/zero/partition_parameters.py | Partition storage, gathers, and parameter status |
deepspeed/runtime/swap_tensor/ | Offload staging and state movement |
deepspeed/runtime/pipe/engine.py | Pipeline-specific scheduling and restrictions |
Those branches are cooperating responsibilities, not three sequential phases. When a parameter is unexpectedly resident, follow its gather/release code; when a step occurs too early, follow accumulation bookkeeping.
Compatibility also belongs to a specific path. The pinned PipelineEngine rejects ZeRO-2 and ZeRO-3. Current AutoTP training documentation separately supports stages 0 through 3, with model-specific partition rules and topology-aware checkpoint handling. Validate each wrapper combination against its specific engine and documentation.[2][3]
FSDP is another route to sharded training, using PyTorch-native interfaces and a different engine/config boundary. Megatron supplies model-parallel building blocks and schedules. Compare the integration your model requires, then test loss parity and checkpoint portability. Whichever stack you choose, verify that a saved directory resumes the intended trajectory.
Resume training and export weights separately
A distributed checkpoint contains more than inference weights. Treat two outcomes separately:
| Outcome | State needed | Verification |
|---|---|---|
| Resume training | Parameters, optimizer, scheduler, counters, required RNG/data progress | Next update agrees with an uninterrupted reference |
| Export a model | Consolidated weights plus model/tokenizer configuration | Reference inputs produce matching outputs within tolerance |
Every participating process must call engine.save_checkpoint with a consistent tag. Calling it only on rank zero can hang. For a ZeRO-3 save/reload test, construct a fresh model and engine before loading; the pinned implementation warns against loading immediately into the same already-partitioned engine.[2]
Use client_state for application-owned progress that DeepSpeed doesn't automatically manage. Restore the sampler or data-loader position and any required per-rank RNG state deliberately. A saved step counter alone doesn't ensure the next microbatch is the same.
For supported dense ZeRO layouts, get_fp32_state_dict_from_zero_checkpoint reconstructs FP32 weights. Plan CPU memory for consolidation; lazy conversion changes when tensors are materialized, not the size of the final weights. A weights-only export doesn't contain enough state to reproduce the next Adam update.[10]
Use the scalar Adam state from our ownership lab to isolate the difference. Both restores begin with the identical weight and use the same next gradient and step number. Predict whether omitting only the moments changes the next weight.
import json
checkpoint = json.loads(json.dumps({"next_step": 6, "parameter_state": result[0]}))
state = checkpoint["parameter_state"]
next_gradient = 0.4
reference = adam_step(*result[0], next_gradient, 6)
resumed = adam_step(*state, next_gradient, checkpoint["next_step"])
weights_only = adam_step(state[0], 0.0, 0.0, next_gradient, checkpoint["next_step"])
assert state[0] == result[0][0]
assert resumed == reference
assert not isclose(weights_only[0], reference[0], rel_tol=1e-12, abs_tol=1e-12)
print(f"full-state next weight={resumed[0]:.8f}")
print(f"weights-only next weight={weights_only[0]:.8f}")
print(f"next-update gap={abs(weights_only[0] - reference[0]):.8f}")full-state next weight=0.01977263
weights-only next weight=0.01900607
next-update gap=0.00076656Identical saved weights can therefore pass an inference check while resuming a different training trajectory. This CPU fixture round-trips one parameter's weight, moments, and counter; it doesn't exercise DeepSpeed checkpoint files, topology conversion, RNG restoration, or data progress.
Universal Checkpointing adds a conversion/load path for supported topology changes, not arbitrary model conversion. The current tutorial requires compatible parameter names and shapes and documents separate AutoTP/AutoEP layouts. In particular, partition-native AutoEP checkpoints aren't supported by zero_to_fp32.py; the documented path uses ds_to_universal.py instead.[11]
Before a long run, save at the intended topology, stop every worker, reload into a fresh run, and compare the next loss, update, scheduler state, and data progress. Test export separately. These distributed recovery operations weren't executed here.
Project context and contribution path
The original ZeRO authors worked at Microsoft. The project's pinned governance defines technical oversight through a Technical Steering Committee and its relationship with Linux Foundation AI & Data.[1][12] The pinned committer roster includes contributors affiliated with Snowflake, Microsoft, Anyscale, AMD, Google, Intel, UIUC, and Argonne National Laboratory. Those affiliations describe the dated roster, not everyone's current employer.[13]
The code license is Apache-2.0. The charter specifies Developer Certificate of Origin sign-off for contributions and CC BY 4.0 for project documentation.[14][12] Check the current contribution instructions before submitting a change.
For a source-level investigation, reduce the failing run first. Include the version/commit, resolved config, smallest reproducer, topology, and earliest rank error. Report whether the failure is initialization, correctness, performance, or recovery; a peer's collective timeout can be a consequence of an earlier failure elsewhere.
Turn a failed run into a specific hypothesis
Common failure modes
| Observation | Next investigation |
|---|---|
| Longer sequences OOM despite a small state estimate | Separate saved activations and workspaces from model state |
| OOM during a stage-3 gather | Count simultaneous gathered buffers, padding, and prefetch |
| Loss changes after scaling out | Check effective batch, token weighting, precision, and update count |
| Hang with routed experts | Compare gather order and the configured leaf boundary |
| Lower GPU peak but slower steps after offload | Measure exposed transfers and CPU/storage service time |
| Saved directory exists but resume diverges | Restore optimizer, scheduler, RNG, and data progress in a fresh engine |
Evaluation rubric
- Derive each stage's storage from tensor ownership and dtype, then explain what the estimate excludes.
- Trace one parameter and its gradient through gather, compute, reduction, and the accumulation-boundary update.
- Separate CPU reference checks from GPU measurements, and distinguish a resumable checkpoint from a weights-only export.
Follow-up questions
The state estimate fits, but enabling prefetch causes an OOM. Why isn't that a contradiction?
Answer
The estimate counts persistent state. Prefetch can keep multiple full gathered tensors live alongside their shards. Add the simultaneous allocations, activations, and workspaces before comparing with available device memory.
Two ranks route to different experts. Should you create parameters lazily and register them as external parameters?
Answer
No. The routing problem concerns compatible gathers for existing submodules. A suitable leaf boundary gathers descendants together, with a memory cost. External-parameter coordination concerns access outside an existing parameter's owning module; it doesn't make post-initialization parameter creation safe for the optimizer.