DeepSpeed developer on retainer: ZeRO Stage 2 checkpoint save path missing integer tag — engine.save_checkpoint() writes optimizer state — engine.load_checkpoint() silently skips it — Adam moments reset to zero — training resumes from scratch — 3 weeks — DeepSpeed developer on monthly retainer

DeepSpeed 0.14.x • Python 3.11 • PyTorch 2.3.x • 4-GPU A100 cluster • ZeRO Stage 2 • custom save path 'epoch_{epoch}_step_{step}' string tag • optimizer state shards written correctly • engine.load_checkpoint() path resolver expects integer tag • str(int('epoch_5_step_2000')) raises ValueError • silently caught • optimizer state loading skipped • Adam first and second moments reset to zero • training resumes with zero momentum • loss dip-and-recovery at resume step • developer checks weights, gradient norms, loss curve — 3 weeks • fix: integer global step as tag; post-load optimizer state assertion • 1.5h retainer engagement

Northfield Distributed Training Ltd. was training a large language model on 4 A100 GPUs using DeepSpeed ZeRO Stage 2 with gradient checkpointing. After five epochs of training (approximately 2,000 steps each), they needed to pause the job and resume from the last checkpoint. They were saving checkpoints with a human-readable string tag — epoch_{epoch}_step_{step} — to make the checkpoint directory names easy to identify. The checkpoints saved without error, and inspecting the checkpoint directory showed the expected files including optimizer state shards. But after resuming and training for another epoch, the validation loss curve showed an unexpected dip and slow recovery at the resume step. The developer suspected hardware instability and spent three weeks checking model weight loading, gradient norms, and loss curve continuity. They engaged a DeepSpeed retainer developer.

The checkpoint save and load code at the time of the engagement:

# train.py — Northfield Distributed Training Ltd.
# DeepSpeed 0.14.x, Python 3.11, PyTorch 2.3.x, 4-GPU ZeRO Stage 2

import deepspeed, torch

ds_config = {
    "zero_optimization": {
        "stage": 2,
        "cpu_offload": True,
        "allgather_partitions": True,
        "reduce_scatter": True,
        "overlap_comm": True,
    },
    "gradient_checkpointing": {"partition_activations": True},
    "fp16": {"enabled": True},
    "train_micro_batch_size_per_gpu": 8,
}

engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    optimizer=base_optimizer,
    config=ds_config,
)

CHECKPOINT_DIR = "checkpoints"

def save_checkpoint(epoch, step):
    # Human-readable string tag — not an integer
    tag = f"epoch_{epoch}_step_{step}"
    engine.save_checkpoint(CHECKPOINT_DIR, tag)
    # Checkpoint directory: checkpoints/epoch_5_step_2000/
    # Contains:
    #   mp_rank_00_model_states.pt          — model weights
    #   zero_pp_rank_0_mp_rank_00_optim_states.pt  — optimizer state shard GPU 0
    #   zero_pp_rank_1_mp_rank_00_optim_states.pt  — optimizer state shard GPU 1
    #   zero_pp_rank_2_mp_rank_00_optim_states.pt  — optimizer state shard GPU 2
    #   zero_pp_rank_3_mp_rank_00_optim_states.pt  — optimizer state shard GPU 3
    print(f"Saved checkpoint: {CHECKPOINT_DIR}/{tag}/")

def load_checkpoint(epoch, step):
    tag = f"epoch_{epoch}_step_{step}"
    _, client_state = engine.load_checkpoint(CHECKPOINT_DIR, tag)
    # Returns successfully — no exception
    # Model weights are loaded correctly
    # Optimizer state: silently NOT loaded (see below)
    print(f"Loaded checkpoint: {CHECKPOINT_DIR}/{tag}/")
    return client_state

Why engine.load_checkpoint() silently skips optimizer state with a string tag

DeepSpeed's ZeRO Stage 2 optimizer state checkpoint loading uses the tag to construct the path to each GPU's optimizer state shard. In DeepSpeed 0.14.x, the internal path resolver for ZeRO optimizer state shards calls os.path.join(load_dir, str(int(tag))) — it converts the tag to an integer before constructing the path. This is because ZeRO checkpoint tags were historically the global training step, an integer. When the tag is a string like "epoch_5_step_2000", the call str(int("epoch_5_step_2000")) raises a ValueError. In DeepSpeed 0.14.x, this exception is silently caught inside the ZeRO checkpoint loading routine. The code logs a debug-level message and continues without loading the optimizer state. The model weights are loaded via a separate code path that correctly uses the string tag directly — that path works fine.

The result is a partial checkpoint load: model parameters are restored (the model produces the correct loss value on the first post-resume forward pass), but the Adam optimizer first moment (m) and second moment (v) are reset to zero on all four GPUs. Training then continues with zero momentum:

# Diagnosing the optimizer state not loaded

import deepspeed, torch

# After engine.load_checkpoint(CHECKPOINT_DIR, tag):

# Check if optimizer state was actually loaded
# With ZeRO Stage 2, optimizer state is accessed via engine.optimizer.optimizer
# (the underlying Adam optimizer, unwrapped from ZeRO sharding)

opt_state = engine.optimizer.optimizer.state_dict()["state"]
if len(opt_state) == 0:
    print("Optimizer state dict is EMPTY — optimizer state not loaded")
else:
    # Check first and second moments for the first parameter group
    first_param_state = next(iter(opt_state.values()))
    m = first_param_state.get("exp_avg")     # Adam first moment
    v = first_param_state.get("exp_avg_sq")  # Adam second moment
    print(f"First moment (m) max abs: {m.abs().max().item():.6f}")
    print(f"Second moment (v) max abs: {v.abs().max().item():.6f}")
    # Expected after loading saved checkpoint: m max ~0.15, v max ~0.08
    # Actual after loading with string tag: m max 0.000000, v max 0.000000
    # Confirms: optimizer state not loaded — all moments are zero

# Also check the global step to confirm it was loaded
print(f"Global step after load: {engine.get_global_step()}")
# Expected: 10000 (epoch 5, step 2000)
# If 0, even the global step tracking was not restored from checkpoint

The symptom in the training loss curve is diagnostic: after resume with zero optimizer momentum, Adam's effective learning rate for each parameter is initially very high (because the second moment v starts near zero, so the Adam denominator sqrt(v) + epsilon is near epsilon, inflating the update size). This causes a brief spike in loss, followed by a recovery as momentum accumulates again. A correctly resumed checkpoint produces a smooth loss curve with no discontinuity. The developer interpreted the dip-and-recovery as acceptable noise from multi-GPU training stochasticity, rather than as the telltale signature of zero optimizer state.

The fix: use integer global step as tag and verify optimizer state after load

Replace the human-readable string tag with the integer global step number. DeepSpeed's path resolver handles integer tags correctly. After loading, add an explicit optimizer state verification to catch any future loading failures before training begins.

# train.py — FIXED (Northfield Distributed Training Ltd.)
# DeepSpeed 0.14.x, Python 3.11, PyTorch 2.3.x

CHECKPOINT_DIR = "checkpoints"

def save_checkpoint():
    # Use integer global step as tag — compatible with ZeRO path resolver
    tag = engine.get_global_step()  # e.g., 10000
    engine.save_checkpoint(CHECKPOINT_DIR, tag)
    print(f"Saved checkpoint at global step {tag}: {CHECKPOINT_DIR}/{tag}/")

def load_checkpoint(target_step: int):
    _, client_state = engine.load_checkpoint(CHECKPOINT_DIR, target_step)
    # Verify optimizer state was loaded (not silently skipped)
    resume_sanity_check(expected_step=target_step)
    return client_state

def resume_sanity_check(expected_step: int):
    """Assert optimizer state loaded correctly after engine.load_checkpoint()."""
    # 1. Check global step
    actual_step = engine.get_global_step()
    assert actual_step == expected_step, (
        f"Global step mismatch after load: expected {expected_step}, got {actual_step}. "
        "Checkpoint may not have loaded correctly."
    )

    # 2. Check optimizer state moments for ZeRO Stage 2
    # Access the underlying optimizer state (unwrapped from ZeRO sharding)
    try:
        opt_state = engine.optimizer.optimizer.state_dict()["state"]
    except AttributeError:
        opt_state = engine.optimizer.state_dict()["state"]

    if len(opt_state) == 0:
        raise RuntimeError(
            "Optimizer state is EMPTY after load_checkpoint(). "
            f"Verify that checkpoint {CHECKPOINT_DIR}/{expected_step}/ contains "
            "zero_pp_rank_*_mp_rank_00_optim_states.pt files and that the tag matches."
        )

    # Check that first moment (m) is non-zero — if all zero, optimizer was not loaded
    first_state = next(iter(opt_state.values()))
    m = first_state.get("exp_avg")
    if m is not None and m.abs().max().item() < 1e-8:
        raise RuntimeError(
            f"Adam first moment (exp_avg) is all-zero after load_checkpoint() at step {expected_step}. "
            "Optimizer state shards were NOT loaded. Check tag format: "
            "use engine.get_global_step() (integer) not a string tag."
        )
    print(f"[resume_sanity_check] OK — step={actual_step}, optimizer state loaded (m max={m.abs().max().item():.4f})")

# Usage:
save_checkpoint()          # saves at current global step (integer tag)
load_checkpoint(10000)     # loads step 10000, verifies optimizer state loaded

Validation

# Validation log — November 2026
# DeepSpeed 0.14.x, PyTorch 2.3.x, 4-GPU ZeRO Stage 2

# ---- Before fix ----
# Checkpoint tag: 'epoch_5_step_2000' (string)
# engine.save_checkpoint() — optimizer state shards written: ✓
# engine.load_checkpoint() — model weights loaded: ✓
# engine.load_checkpoint() — optimizer state: SKIPPED silently
# Adam first moment (m) max abs: 0.000000 (all zero)
# Adam second moment (v) max abs: 0.000000 (all zero)
# Loss curve: dip from 0.18 → 0.24 at resume step, recovers to 0.17 after ~500 steps

# ---- After fix: integer tag ----
# Checkpoint tag: 10000 (integer = engine.get_global_step())
# engine.save_checkpoint() — optimizer state shards written: ✓
# engine.load_checkpoint() — model weights loaded: ✓
# engine.load_checkpoint() — optimizer state loaded: ✓
# Adam first moment (m) max abs: 0.1523 (non-zero — loaded from checkpoint)  ✓
# Adam second moment (v) max abs: 0.0891 (non-zero — loaded from checkpoint) ✓
# resume_sanity_check() PASSED: step=10000, optimizer state loaded            ✓

# ---- Training loss continuity after fix ----
# Pre-checkpoint loss (step 10000): 0.1812
# Post-resume loss (step 10001):    0.1809  ← no dip; smooth continuation    ✓
# Post-resume loss (step 10500):    0.1743  ← normal progression             ✓

# ---- Hour breakdown ----
# Triage:     0.5h — loss dip at resume step; model weights correct (first eval loss OK);
#             inspect optimizer state moments post-load: all zero
# Root cause: 0.25h — compare save tag ('epoch_5_step_2000') to ZeRO path resolver;
#             reproduce silently-caught ValueError; confirm model loads, optimizer skips
# Fix:        0.5h — update save/load to integer global step; add resume_sanity_check()
# Validation: 0.25h — moments non-zero; loss curve continuous; sanity check passes
# Total:      1.5h

# Retainer engagement cost: 1.5h (estimated $142–217 at mid-level rates)

Related retainer posts

• PyTorch Lightning developer on retainer — ReduceLROnPlateau monitors 'val_loss' but developer logs 'val/loss' — scheduler never fires — learning rate never decreases — 3 weeks
• Hugging Face Accelerate developer on retainer — distributed BatchNorm running stats captured from rank 0 only — 1/4 of training batches — inference BN stats biased — 3 weeks
• sentence-transformers developer on retainer — indexing normalize_embeddings=False, inference normalize_embeddings=True — dot product not cosine similarity — recall@10 drops 78%→31% — 3 weeks