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

[Proposal] Extend affine live-state transfer to CPU-offloaded optimizer storage

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

Maintainers usually reply within 3 days

@ArjunPakhan is already working on this.

Since Oct 25, 2026.

  • #8905 by @ArjunPakhan — open

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
25/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
distributed-systems, machine-learning

Research direction

This is a design proposal with no implementation files or tests named. First read #8885, #8809, and #8230 and ask the parent maintainers to agree on scope and a residency/readiness contract. Done means an accepted design defining safe live storage access, bounded CPU↔accelerator staging, and collective preflight before implementation begins.

Written by the indexing model from the issue text.

Description

Is your feature request related to a problem? Please describe.

#8885 describes live ZeRO parameter/master/optimizer storage for use with #8809's bounded transfer executor. It explicitly rejects CPU-offloaded optimizer states. The executor also requires buffers to reside on its accelerator transport device, so removing the map guard alone would not provide correct CPU-offload support.

I would like to contribute this as a scoped follow-up to #8230, #8885 and #8809, after agreement on the design and ownership. This is a proposal, not a feature implementation or a claim that the parent PRs are incomplete bug fixes.

Describe the solution you'd like

Initial scope: CPU-offloaded optimizer masters and Adam moments under ZeRO-2/3, with training paused at a documented safe boundary. Preserve existing device-resident behavior. Exclude NVMe, quantized storage, MiCS/hpZ, concurrent training/transfer and topology-changing optimizer resumes initially.

Proposed phases:

  1. Describe the same logical ranges while representing the shard's residency and readiness separately from the affine geometry. Decide with the parent authors whether this belongs in an endpoint descriptor or in a checked storage-view/lease API.
  2. Obtain a stable, live view after pending offload copies complete; keep the owner and staging buffers alive until transfer completion. Refuse unavailable/swapped storage rather than returning a placeholder.
  3. Add bounded CPU↔accelerator staging around existing transport rounds. NCCL payloads remain on the accelerator. Pinned-memory availability, copy events and buffer lifetimes must be explicit; budgets must bound both host and device staging. No whole-model gather or implicit CPU→GPU allocation proportional to total model size.
  4. Preserve collective preflight: all participants reject unsupported residency or allocation failures before target writes. Do not claim rollback or continued training after a transport failure has partially overwritten target storage.

Describe alternatives you've considered

  • Gather complete FP32 tensors, copy them to the accelerator and use the current executor: straightforward but duplicates memory and full gathers.
  • Read serialized checkpoints instead of live state: useful for restart, but less suitable for repeated training/rollout synchronization.
  • Use CPU/Gloo-only transport: a possible backend option, but it does not directly solve mixed CPU optimizer storage and GPU rollout targets. It also has different performance and stream-completion contracts.

Additional context and baseline experiments

On 2×RTX 3090 24 GiB, driver 580.105.08, torch 2.12.1+cu130 / CUDA 13.0, I exercised frozen #8885 head e34367a7f790eb3f39c62f8b36978cbf3186f0aa with the existing executor file from #8809 head ba9e8b65219efdde2833abda3d66ceb03f907d6d. This combined experimental tree was not published or presented as a new implementation.

A small BF16 model's live ZeRO-3 parameter shards transferred into replicated target buffers with bitwise agreement. The same transfer also passed with bitwise agreement under ZeRO-2 and ZeRO-3 CPU optimizer offload. In both CPU-offload cases, optimizer map access raised the documented NotImplementedError. These runs used 13×13 weights plus a 13-element bias, exercising uneven partitions/padding, and a 128-byte per-peer staging budget. These are correctness smoke tests, not end-to-end throughput claims.

Planned acceptance experiments, once maintainers accept a design:

  • Compare every restored master and Adam moment with an independently gathered reference, including straddling parameters, uneven shards and padding.
  • Check the next optimizer update against an uninterrupted control at the same topology.
  • Exercise CPU→GPU, GPU→CPU and CPU→CPU residency combinations, rejected unsupported states, event ordering and bounded staging under a small budget.
  • Begin with two-GPU small models; extend to four-GPU DP/TP topology integration and a representative model for repeated-sync latency, CPU RSS, peak GPU allocation, transfer bytes and run-to-run variability. Report both memory and throughput and avoid treating a microbenchmark as a general speedup.

Before implementation, could the parent maintainers confirm whether this scope is useful, whether related work already exists, and the preferred residency/readiness contract? AI assistance was used to investigate existing code and prepare this proposal.

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.