Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

[BUG] ZeRO-3 resume can fail when a resident adapter becomes an external dependency after tracing

Open
  • Is something wrong?
#8,874 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

@DimensionSTP is already working on this.

Since Oct 21, 2026.

  • #8875 by @DimensionSTP — open

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
25/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Stale
Tech stack
python, pytorch
Domain
distributed-systems, machine-learning

Research direction

Save the supplied reproducer as reproduce.py and run it with the stated two GPU torchrun command. Read the automatic external parameter discovery and parameter trace code implicated by the report, then compare behavior when resident external parameters are discovered on first access. Done means the full resume completes the second optimizer step and prints FULL_RESUME_PASS; linked pull request #8875 is already open.

Written by the indexing model from the issue text.

Description

Describe the bug

A full ZeRO-3 checkpoint resume can fail when a resident adapter weight is accessed outside its owning module, for example to read its dtype before calling that module. The first access is skipped by automatic external-parameter discovery because the parameter is already AVAILABLE. Once it is partitioned, a later access adds a dependency after the parameter trace has been cached. Reentrant checkpoint recomputation then finds the wrong parameter at the front of the fetch queue.

This was observed in a 27B LoRA training run with GA 4 and reentrant checkpointing after upgrading the training stack. It reproduces with DeepSpeed 0.19.8 and current master 68faa67c; I have not identified a first-bad version. The minimal example below requires only PyTorch and DeepSpeed, not PEFT, TRL, model weights, or a dataset.

To Reproduce

Save the following as reproduce.py, then launch on two GPUs with a new checkpoint directory:

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 reproduce.py /tmp/zero3-resident-adapter-repro
Minimal full-resume reproducer 
import argparse
import os

import torch
from torch.utils.checkpoint import checkpoint

import deepspeed
from deepspeed.accelerator import get_accelerator


class Adapter(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.base = torch.nn.Linear(8, 8, bias=False)
        self.base.weight.requires_grad_(False)
        self.adapter_A = torch.nn.Linear(8, 2, bias=False)
        self.adapter_B = torch.nn.Linear(2, 8, bias=False)

    def forward(self, value):
        base = self.base(value)
        value = value.to(self.adapter_A.weight.dtype)
        return base + self.adapter_B(self.adapter_A(value))


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.adapter = Adapter()

    def forward(self, value):
        return checkpoint(self.adapter, value, use_reentrant=True).square().mean()


def initialize():
    model = Model()
    optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-3)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)
    config = {
        'train_micro_batch_size_per_gpu': 2,
        'gradient_accumulation_steps': 4,
        'zero_optimization': {
            'stage': 3,
            'stage3_param_persistence_threshold': 16,
            'stage3_prefetch_bucket_size': 0,
            'stage3_max_reuse_distance': 0,
            'offload_param': {'device': 'cpu', 'pin_memory': False},
        },
    }
    engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, lr_scheduler=scheduler, config=config)
    return engine


def train_step(engine):
    for _ in range(4):
        value = torch.ones(2, 8, device=engine.device, requires_grad=True)
        loss = engine(value)
        engine.backward(loss)
        engine.step()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('checkpoint_directory')
    args = parser.parse_args()
    get_accelerator().set_device(int(os.environ['LOCAL_RANK']))
    deepspeed.init_distributed(dist_backend=get_accelerator().communication_backend_name())
    torch.manual_seed(42)
    engine = initialize()
    train_step(engine)
    engine.save_checkpoint(args.checkpoint_directory, tag='seed')
    engine.destroy()
    engine = initialize()
    loaded, _ = engine.load_checkpoint(args.checkpoint_directory, tag='seed')
    assert loaded is not None and engine.global_steps == 1
    train_step(engine)
    assert engine.global_steps == 2
    engine.destroy()
    print('FULL_RESUME_PASS', flush=True)


if __name__ == '__main__':
    main()

The adapter A has 16 elements and is persistent at the configured threshold; the frozen base has 64 elements. The original source fails on both ranks after restoring the checkpoint, with the adapter weight expected but the base weight found in the fetch queue. The same script completes after discovering resident external parameters on their first access. The local runs used a torchrun progress wrapper, without forking or mocking the engine.

Expected behavior

The same module execution should not acquire new external dependencies solely because a parameter changes between AVAILABLE and NOT_AVAILABLE. Full resume should complete the second optimizer step without a fetch-trace mismatch.

ds_report output

The original failure was checked on DeepSpeed 0.19.8 and master 68faa67c. The master source build reports 0.19.9+68faa67c; this is not a claim about a released 0.19.9 wheel. The relevant report from the source validation environment is:

torch version .................... 2.13.0+cu129
deepspeed info ................... 0.19.9+68faa67c, 68faa67c
torch cuda version ............... 12.9
torch hip version ................ None
nvcc version ..................... 12.8
deepspeed wheel compiled w. ...... torch 2.13, cuda 12.9
shared memory (/dev/shm) size .... 1.00 GB
ninja ............................ available
fused_adam ....................... not prebuilt, compatible

Screenshots

Not applicable; the reproducible failure is a Python RuntimeError.

System info (please complete the following information):

  • Ubuntu 24.04.1 LTS, Python 3.12.13.
  • Single H100 host; two GPUs for the minimal reproduction and four for the original 27B comparison. The cloud exposes the device name as CUDA GPU, compute capability 9.0, 81,559 MiB each.
  • Driver 535.161.08, PyTorch 2.13.0+cu129.
  • Original integration: PEFT 0.21.2, Transformers 5.19.0, TRL 1.15.0, Accelerate 1.15.0, Liger disabled. These extra packages are not needed by the small reproduction.

Launcher context

torchrun for the two-rank reproduction and official engine regression; Accelerate for the four-GPU SFT integration.

Docker context

Backend.AI-managed cloud container. An exact public image tag was not recorded.

Additional context

In the real model, adapter A contains 655,360 elements. Keeping all libraries and the checkpoint fixed, reducing the persistence threshold from 1,000,000 to 100,000 made that weight non-persistent and allowed resume. This is a model-specific workaround, not a requirement to reject higher thresholds. Discovering the external access before caching the trace lets the same checkpoint resume at 1,000,000 with unmodified PEFT.

Existing #7205 concerns a generation-unwrap trace failure without a minimal reproduction; I am not assuming it has this same cause. #6718 discusses related weight getter behavior under compilation, whereas this reproduction uses eager execution.

Dominant language
Python
Stars
43.3k
Forks
5.1k
Avg merge
4d 20h
Merged PRs (30d)
69

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepspeedai/DeepSpeed

All issues in deepspeedai/DeepSpeed

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.