LearnProject Deep DivesDeep Dive - DeepSpeed

⚡HardFine-Tuning & Training

Deep Dive - DeepSpeed

Trace DeepSpeed from ZeRO ownership and dtype-aware memory accounting through gathered working sets, sharded Adam updates, offload, and checkpoint recovery.

20 min read

Learning path

Step 187 of 196 in the full curriculum

Deep Dive - Megatron Deep Dive - Ray

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]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}

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.

We'll follow the memory accounting into the engine, a gathered module, and a checkpoint. The source tour uses commit c455031422641a825588c926455458c82b27b6c0 (August 6, 2026, before the 0.19.4 release); the current documentation consulted identifies itself as 0.19.6. These are distinct version references, not evidence that every feature combination works on both.^{[2]Reference 2DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}^{[3]Reference 3Training APIhttps://deepspeed.readthedocs.io/en/latest/training.html}

The marked examples run on a CPU with Python's standard library. They check ownership arithmetic and an Adam update, not DeepSpeed kernels, distributed execution, or GPU memory measurements.

Count state before choosing a stage

Let $P$ 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.^{[4]Reference 4DeepSpeed Configuration JSONhttps://www.deepspeed.ai/docs/config-json/}

Let $p$ , $g$ , and $o$ denote bytes per parameter for compute weights, gradients, and optimizer-plus-master state. For $N$ data-parallel ranks, idealized per-rank storage is:

\begin{aligned} M_{DP} &= P(p+g+o),\\ M_{Z1} &= P(p+g+o/N),\\ M_{Z2} &= P(p+(g+o)/N),\\ M_{Z3} &= P(p+g+o)/N. \end{aligned}

With $(p,g,o)=(2,2,12)$ and $N=4$ , those are $16P$ , $7P$ , $5.5P$ , and $4P$ 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]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}

Idealized model state per rank for four data ranks. Replicated data parallel: parameters two, gradients two, optimizer plus master twelve, total sixteen bytes per parameter. ZeRO-1 totals seven, ZeRO-2 five point five, and ZeRO-3 four. Every row uses the same scale. — ZeRO removes replication in stages. The large first reduction comes from sharding the twelve-byte optimizer-plus-master block; the remaining reductions change gradient and parameter ownership.

Use exact fractions to compare recipes without confusing GB ( $10^9$ bytes) and GiB ( $2^{30}$ bytes). This estimator deliberately excludes padding and live working sets.

zero_memory.py

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_000

Output

DP     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 GiB

Stage 1 partitions optimizer state and master weights; stage 2 also partitions reduced gradients; stage 3 also partitions compute parameters.^{[5]Reference 5ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html} 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.

Reduce first, then update the owned slice

ZeRO's central arithmetic 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 $(1^2+3^2)/2=5$ , while $((1+3)/2)^2=4$ .^{[6]Reference 6Adam: A Method for Stochastic Optimization.https://arxiv.org/abs/1412.6980}

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.

owned_adam.py

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")

Output

11 parameters, 4 owners, 5 updates: weights and moments match

The 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.

Let the engine own the microstep

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]Reference 2DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}

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]Reference 3Training APIhttps://deepspeed.readthedocs.io/en/latest/training.html}

engine_integration_sketch.py

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()

This four-data-rank configuration uses BF16 model parameters and explicitly requests FP32 gradient accumulation. It's a starting configuration to validate on compatible GPUs, not a benchmarked recommendation. The optimizer settings belong to your training recipe; ZeRO doesn't determine a suitable learning rate.

ds_config.json

{
  "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 $2\times8\times4=64$ 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.^{[4]Reference 4DeepSpeed Configuration JSONhttps://www.deepspeed.ai/docs/config-json/}

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.^{[5]Reference 5ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

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; it doesn't prove that a multi-rank training step or offload path works.^{[7]Reference 7DeepSpeed Installation Detailshttps://www.deepspeed.ai/tutorials/advanced-install/} Neither the integration sketch nor this GPU configuration was executed for this lesson.

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.^{[5]Reference 5ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

Illustrative ZeRO-3 schedule without persistent full parameters: gather, forward, release; then gather again, backward, and reduce gradients to owners. Local Adam updates occur at the accumulation boundary. — Forward and backward may need separate parameter gathers. The gradient reduction gives each owner its update slice; an optimizer step waits for the accumulation boundary. Persistence can skip some releases and gathers.

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]Reference 2DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0} 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.

Separate-allocation model for four eight-element tensors on four ranks. Each rank retains eight shard elements. One full gathered tensor adds eight elements for sixteen total; prefetching another adds eight more for twenty-four. The local values appear in both shard and full buffer. — Count allocations that coexist, not only values received from peers. Prefetch trades a larger live working set for a chance to hide the next gather; it doesn't guarantee that the transfer is hidden.

This storage model rounds each tensor's shard up to a whole element and allocates a separate padded gather buffer. It models one specified layout, not a universal DeepSpeed peak estimator.

gather_storage.py

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] == 24

Output

gathered: shards 8 + full buffers 0 = 8 elements
gathered: shards 8 + full buffers 8 = 16 elements
gathered: shards 8 + full buffers 16 = 24 elements

The 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.^{[5]Reference 5ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

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 its explicit gathering/registration mechanisms. This coordinates access to an existing parameter.^{[5]Reference 5ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

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. This doesn't register newly created trainable parameters with an already-built optimizer.^{[3]Reference 3Training APIhttps://deepspeed.readthedocs.io/en/latest/training.html}

Communication volume isn't elapsed time

A simplified ring model makes the paper's communication comparison concrete. For $P$ communicated elements per full vector on $N$ ranks, one all-gather or reduce-scatter sends $(N-1)P/N$ 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	$2(N-1)P/N$
ZeRO-2	Gradient reduce-scatter + updated-parameter all-gather	$2(N-1)P/N$
ZeRO-3	Forward parameter gather + backward parameter gather + gradient reduce-scatter	$3(N-1)P/N$

This explains the paper's approximately $2P$ versus $3P$ model for large $N$ .^{[1]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054} It doesn't promise that stage 3 takes 1.5 times as long. Different gradient/parameter dtypes, repeated gathers during accumulation or recomputation, parameter persistence, bucket sizes, overlap, and collective algorithms change the comparison. In stage 1, gradient storage stays replicated; don't infer a particular reduce-scatter implementation from its ownership formula.

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.

Offload exchanges memory pressure for movement

ZeRO-Offload moved optimizer-related state and work to CPU memory; ZeRO-Infinity extends the hierarchy to CPU and NVMe storage.^{[8]Reference 8ZeRO-Offload: Democratizing Billion-Scale Model Traininghttps://www.usenix.org/conference/atc21/presentation/ren-jie}^{[9]Reference 9ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning.https://arxiv.org/abs/2104.07857} 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.^{[5]Reference 5ZeRO Configurationhttps://deepspeed.readthedocs.io/en/stable/zero3.html}

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.

offload_fragment.json

{
  "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. These are illustrative values, not measured NVMe performance.

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. “It fits” answers a capacity question, not whether the training run finishes sooner.

Read the engine path, not just the config

The pinned source organizes the mechanisms as follows:^{[2]Reference 2DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}

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

Diagram showing initialize config + process groups, Selected engine, backward / step reduction + owned update, and module hooks gather + release. — initialize config + process groups, Selected engine, backward / step reduction + owned update, and module hooks gather + release.

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. AutoTP isn't PipelineEngine, and neither claim establishes arbitrary combinations with other wrappers.^{[2]Reference 2DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}^{[3]Reference 3Training APIhttps://deepspeed.readthedocs.io/en/latest/training.html}

FSDP is another route to sharded training, using PyTorch-native interfaces rather than the same engine/config boundary. Megatron supplies model-parallel building blocks and schedules. Compare the integration your model requires, then test it; a feature list doesn't establish loss parity or checkpoint portability.

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]Reference 2DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}

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]Reference 10DeepSpeed Model Checkpointinghttps://deepspeed.readthedocs.io/en/latest/model-checkpointing.html}

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]Reference 11Universal Checkpointing with DeepSpeed: A Practical Guide.https://www.deepspeed.ai/tutorials/universal-checkpointing/}

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]Reference 1ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.https://arxiv.org/abs/1910.02054}^{[12]Reference 12DeepSpeed Project Charter and Governancehttps://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/GOVERNANCE.md} 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]Reference 13DeepSpeed TSC Committershttps://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/COMMITTERS.md}

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]Reference 14DeepSpeed Apache License 2.0https://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/LICENSE}^{[12]Reference 12DeepSpeed Project Charter and Governancehttps://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/GOVERNANCE.md} 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

Next Step

Continue to Deep Dive - Ray

DeepSpeed coordinates tensors and optimizer state within a distributed training job. Ray adds scheduling for tasks, actors, and workflows across a cluster.

PreviousDeep Dive - Megatron

Share this article

X Facebook LinkedIn Bluesky Reddit Hacker News Email

References

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.

Rajbhandari, S., et al. · 2020 · SC 2020

https://arxiv.org/abs/1910.02054

DeepSpeed Source Repository

DeepSpeed Contributors · 2026

https://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0

Training API

DeepSpeed Team · 2026

https://deepspeed.readthedocs.io/en/latest/training.html

DeepSpeed Configuration JSON

DeepSpeed Contributors · 2026

https://www.deepspeed.ai/docs/config-json/

ZeRO Configuration

DeepSpeed Team · 2026

https://deepspeed.readthedocs.io/en/stable/zero3.html

Adam: A Method for Stochastic Optimization.

Kingma, D. P., Ba, J. · 2015 · ICLR 2015

https://arxiv.org/abs/1412.6980

DeepSpeed Installation Details

DeepSpeed Team · 2026

https://www.deepspeed.ai/tutorials/advanced-install/

ZeRO-Offload: Democratizing Billion-Scale Model Training

Ren, J., Rajbhandari, S., Aminabadi, R. Y., et al. · 2021 · USENIX ATC 2021

https://www.usenix.org/conference/atc21/presentation/ren-jie

ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Rajbhandari, S., et al. · 2021 · SC 2021

https://arxiv.org/abs/2104.07857

DeepSpeed Model Checkpointing

DeepSpeed Contributors · 2026

https://deepspeed.readthedocs.io/en/latest/model-checkpointing.html

Universal Checkpointing with DeepSpeed: A Practical Guide.

DeepSpeed Team · 2026

https://www.deepspeed.ai/tutorials/universal-checkpointing/

DeepSpeed Project Charter and Governance

DeepSpeed Project · 2026

https://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/GOVERNANCE.md

DeepSpeed TSC Committers

DeepSpeed Project · 2026

https://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/COMMITTERS.md

DeepSpeed Apache License 2.0

DeepSpeed Project · 2026

https://github.com/deepspeedai/DeepSpeed/blob/c455031422641a825588c926455458c82b27b6c0/LICENSE

Discussion

Questions and insights from fellow learners.

Discussion loads when you reach this section.

Back to Topics

LearnProject Deep DivesDeep Dive - DeepSpeed

⚡HardFine-Tuning & Training

Deep Dive - DeepSpeed

Trace DeepSpeed from ZeRO ownership and dtype-aware memory accounting through gathered working sets, sharded Adam updates, offload, and checkpoint recovery.

20 min read

Learning path

Step 187 of 196 in the full curriculum

Deep Dive - Megatron Deep Dive - Ray

Personalize this lesson

Adapt explanations and teaching visuals to your background and preferred voice.

The marked examples run on a CPU with Python's standard library. They check ownership arithmetic and an Adam update, not DeepSpeed kernels, distributed execution, or GPU memory measurements.

Count state before choosing a stage

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

Let $p$ , $g$ , and $o$ denote bytes per parameter for compute weights, gradients, and optimizer-plus-master state. For $N$ data-parallel ranks, idealized per-rank storage is:

\begin{aligned} M_{DP} &= P(p+g+o),\\ M_{Z1} &= P(p+g+o/N),\\ M_{Z2} &= P(p+(g+o)/N),\\ M_{Z3} &= P(p+g+o)/N. \end{aligned}

Use exact fractions to compare recipes without confusing GB ( $10^9$ bytes) and GiB ( $2^{30}$ bytes). This estimator deliberately excludes padding and live working sets.

zero_memory.py

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_000

Output

DP     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 GiB

Reduce first, then update the owned slice

owned_adam.py

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")

Output

11 parameters, 4 owners, 5 updates: weights and moments match

Let the engine own the microstep

engine_integration_sketch.py

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()

ds_config.json

{
  "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
    }
  }
}

Stage 3 needs a live working set

This storage model rounds each tensor's shard up to a whole element and allocates a separate padded gather buffer. It models one specified layout, not a universal DeepSpeed peak estimator.

gather_storage.py

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] == 24

Output

gathered: shards 8 + full buffers 0 = 8 elements
gathered: shards 8 + full buffers 8 = 16 elements
gathered: shards 8 + full buffers 16 = 24 elements

External parameters and dynamic routing are different problems

Communication volume isn't elapsed time

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	$2(N-1)P/N$
ZeRO-2	Gradient reduce-scatter + updated-parameter all-gather	$2(N-1)P/N$
ZeRO-3	Forward parameter gather + backward parameter gather + gradient reduce-scatter	$3(N-1)P/N$

Offload exchanges memory pressure for movement

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.

offload_fragment.json

{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": {
      "device": "nvme",
      "nvme_path": "/local_nvme/deepspeed"
    }
  }
}

Read the engine path, not just the config

The pinned source organizes the mechanisms as follows:^{[2]Reference 2DeepSpeed Source Repositoryhttps://github.com/deepspeedai/DeepSpeed/tree/c455031422641a825588c926455458c82b27b6c0}

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

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

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

Next Step

Continue to Deep Dive - Ray

DeepSpeed coordinates tensors and optimizer state within a distributed training job. Ray adds scheduling for tasks, actors, and workflows across a cluster.

PreviousDeep Dive - Megatron

Share this article