[ArjunPakhan]

Description

This PR introduces scoped support for CPU-offloaded optimizer states (master weights and Adam moments under ZeRO-2 and ZeRO-3) within DeepSpeed's live-state transfer executor and affine storage mapping, following up on proposal #8899.

Key Changes

1. Decoupled Residency & Readiness from Affine Geometry (zero_affine.py):

   • Lifted artificial NotImplementedError restrictions for CPU-offloaded optimizer states in storage affine maps and shard accessors while preserving strict residency checks.
   • Ensured Adam moments (exp_avg, exp_avg_sq) raise explicit errors if accessed prior to the first optimizer step, while master weights remain accessible immediately without blocking affine planning.
   • Configured ZeRO-3 shard retrieval to directly target the offloaded fp32 partition (_get_fp32_opt_state_partition) to support in-place mutations without allocating whole-tensor GPU replicas. NVMe offload remains strictly out-of-scope.

2. Bounded Host Pinned-Memory Staging (affine_transfer_executor.py):

   • Updated buffer validation and alias checks to safely handle CPU tensor addresses alongside accelerator device buffers.
   • Routed payloads strictly through accelerator arenas (disabling zero-copy direct transport when buffers reside on CPU).
   • Introduced bounded pinned host staging arenas (host_send_arena / host_recv_arena) constrained by bytes_per_peer budgets to prevent whole-model CPU-GPU allocations.
   • Added robust stream and event synchronization (get_accelerator().Event()) ensuring Host-to-Device (H2D) and Device-to-Host (D2H) copies synchronize safely on accelerator streams before host buffers are reused or unpacked.

3. Documentation & Testing:

   • Updated docs/source/affine-transfer.rst to document CPU offload buffer support and bounded pinned-memory staging semantics.
   • Added unit tests covering CPU-to-GPU, GPU-to-CPU, and CPU-to-CPU residency combinations with uneven partitions, storage guards, and in-place live state mutations (test_affine_transfer_executor.py and test_zero_affine.py).

Related Issues

Closes #8899

Testing

• Ran unit test suites for affine transfer executor and zero affine mappings (188 passed, 0 failed).
• Verified CUDA import neutrality.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-25T15:05:47.620973Z	d772166	Manual request

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6302d548f9

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: d7721664e6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[delock]

Hi @ArjunPakhan , can you resolve the conflicts? Thanks!

[0z5a]

Thanks for picking up #8899. The ZeRO bookkeeping looks right. I ran master's live-storage test, TestZeroOptimizerStorageMaps, with offload_optimizer: {device: cpu} added on this PR's zero_affine.py changes:

• On CPU/gloo, ZeRO-1, 2 and 3 all pass.
• On master the same test raises the offload NotImplementedError.

On a CPU-only run the offloaded storage and the "device" are the same memory. So this checks the offload code paths, the views into optimizer storage and the maps, not host/device residency.

I'd still like a few things changed before this goes in. Two are blocking.

Blocking

1. Base the executor on #8809 instead of carrying a copy of it.

   • affine_transfer_executor.py comes in as a new 589-line file. It is #8809 at an older revision (ba9e8b65).
   • That revision misses the later fix that takes rank and size from the group itself. test_affine_transfer_executor_external_group.py from #8809 fails on this head with target location 0 is not owned by this process.
   • zero_affine.py and test_zero_affine.py are also pre-merge copies of #8885. That is where the conflicts with master come from. At this head, test_zero_affine.py has 2 failures on CPU (test_nest_and_plan_a_scalar_parameter, test_nested_storage_plans_into_another_layout).
   • Could you rebase onto current master and stack on #8809 (or wait for it), so the diff shows only the offload delta?
   • I did that locally, keeping #8809's group-rank code in the one conflicting hunk. Everything except the new tests below passes on CPU.

2. The new TestB2ResidencyAndOffload tests cannot pass on any backend.

   • With world_size = 2, every rank passes the buffers of every location.
   • _validate_buffers requires each process to pass only the locations its endpoints bind to it. So all 8 cases stop in preflight with source location N is not owned by this process. I checked on gloo with enough CPU ranks.
   • Each rank should pass only {loc: buf for loc, buf in buffers.items() if endpoints.source[loc] == rank} (same for targets), and check only its own targets.
   • On a machine with fewer than two devices these classes skip. That may be why the run looked green.
   • Three of the four tests also bind every location to process 0, so they only cover local copies. The cross-process pinned path, which is the new part, is covered only by the first parametrized test.

Correctness

3. deepspeed/checkpoint/affine.py:691: the new map_logical_shape line breaks nest_affine_maps for any parameter with two or more dimensions that one process stores.

   • logical_shape is a tuple, so .numel() raises AttributeError.
   • Reproduction: nest_affine_maps(contiguous_split_map((4, 6), [4], partition_dim=0), {0: zero_storage_affine_map(24, [0])}). It works on master and raises here. Same result with holder [3].
   • This is a common case for weight sync: one process stores the whole parameter. Examples are a TP=1 rollout engine described the way _one_holder_each in test_zero_affine.py does, or any single-rank run.
   • Even with .numel() fixed, storage_shapes is keyed by storage rank, so storage_shapes[0] is the wrong lookup.
   • And a nested map should keep the parameter's logical shape. The rest of the planner relies on that.
   • I don't see what this line is needed for in the offload change. Could it be dropped?

4. affine_transfer_executor.py:332-333: device and stream_ordered are now keyed on accelerator.is_available() instead of on whether a process group exists.

   • Without a process group, every buffer used to be treated as host memory. Now, on a GPU machine, device buffers are accepted and the stream-ordered path runs.
   • Was that intended? If not, please keep the old condition. The offload paths all run inside a process group anyway.

5. zero_affine.py:98, :106, :181-192: these add if not dist.is_initialized() branches to library code.

   • They seem to exist only so the mocked tests can run without a process group.
   • The uninitialized ZeRO-1/2 branch also widens the parameter test to hasattr(param, '_hp_mapping'), which the initialized path doesn't accept.
   • I'd drop these branches and test against real engines instead (see 8).

Smaller

6. Staging bound, docstring at :40 and the docs page.

   • The budget is now bytes_per_peer per active peer per direction, plus up to two pinned host arenas and one extra device arena of bytes_per_peer each.
   • Please state that exactly rather than "bounded by the staging budget".
   • The pinned arenas are allocated on every call.

7. Host synchronizations, :517-524 and :533-552.

   • The receive side records and synchronizes an event for every entry, and each offloaded local copy does the same. That's one host round trip per chunk.
   • Issuing the round's D2H copies and waiting once before unpacking would keep the same ordering.
   • The CPU→device local copy could also go straight from the pinned stage into the target, without local_device_arena.
   • ZeRO-1/2 masters are already pinned when pin_memory is on, and ZeRO-3 masters are too since #8840. For those, the host bounce copy can be skipped. Adam moments are still pageable.

8. Tests for the ZeRO side.

   • test_zero2_cpu_offloaded_optimizer_storage and test_zero3_cpu_offloaded_optimizer_storage build SimpleNamespace stand-ins and reimplement _get_fp32_opt_state_partition inside the test, so they check the mock.
   • Parametrizing the existing TestZeroOptimizerStorageMaps.test_maps_describe_live_storage with offload_optimizer: {device: cpu} and zero_force_ds_cpu_optimizer: false covers stages 1 to 3 against real engines. It runs on CPU.

9. zero_affine.py:218: column 2 of the all-gather (use_offload) is still sent but no longer read. Drop it or use it.

10. zero_affine.py:261: _is_offloaded is not set anywhere in DeepSpeed. Only the _zero_optimizer.cpu_offload half of that condition can be true.

11. affine.py:691 is 188 columns. yapf's limit is 119, so the formatting hook will probably rewrap it. Running pre-commit run --files on the changed files would show it.

GPU coverage

None of the pinned staging or event ordering runs under DS_ACCELERATOR=cpu, because stream_ordered is false there. My local run says nothing about that code. Could you add the hardware and command for a two-GPU run of the fixed tests to the description?

[ArjunPakhan]

Hi @0z5a, thanks for the thorough review and feedback!

I have addressed the feedback on this PR:

• Rebased onto the latest master and cleaned up the commit diff.
• Removed the problematic affine.py map line (.numel()) to support single-rank and full-parameter storage properly.
• Restored the correct process-group check for streaming/device buffers in the transfer executor.
• Cleaned up the temporary mock branches in zero_affine.py.
• Verified that all unit tests pass cleanly locally (67 passed).

Ready for another look when you have a moment! Thanks again.

[0z5a]

Thanks for the update. At 6450ca1d, though, the PR diff against master is a single line in .gitignore (.venv/). None of the changes you list are in it, and master still raises NotImplementedError('Offloaded ZeRO-3 optimizer states have no affine storage map.') in zero_affine.py, so it looks like the rebase or squash dropped them. Could you push the branch that has them? When you do, building the executor part on top of #8809 instead of carrying a copy of it would keep the two from conflicting again.