DeepSpeed

c6291b7e - ZeRO-3: copy offloaded master weights back to the device through pinned buffer (#8840)

ZeRO-3: copy offloaded master weights back to the device through pinned buffer (#8840) # ZeRO-3: pin the offloaded master weights and cast them on the device when copying them back **Part of:** #8672 (Q4 2026 roadmap), "Enable training of large scale models". ## Summary With `offload_optimizer.device: cpu`, ZeRO-3 copies each sub-group's updated float32 master partition back to the device as bf16 after every step. - `pin_memory: true` (the default since #8207) pins the float32 gradient partitions. It does not pin the master partitions, so this copy reads pageable host memory. - The copy is also a cast. A blocking `copy_` from float32 into bf16 casts on the CPU into a new pageable tensor and copies from that at a few GB/s. This PR follows ZeRO-1/2 and #5301: - with `pin_memory: true`, each CPU master partition is pinned when it is created, as ZeRO-1/2 does; - the copy-back moves the pinned float32 master to the device and casts it there, in pieces of 2**26 elements, so the device temporary is one piece (0.25 GiB) instead of a whole sub-group. Results: - Qwen2.5-32B, 8 H200, whole optimizer on the CPU: 12.08 s to 9.42 s per step (-22.0%), same peak device memory, same losses (the new unit test compares the parameters bitwise). - Host memory: the master partitions are already in host memory and become pinned. PyTorch's pinned allocator rounds each allocation up to a power of two: a 1e9-element sub-group (3.73 GiB in float32) reserves 4.00 GiB. ## Duplicate check Searched issues and PRs on 2026-10-13 ("offload pinned fp32 partition", "zero3 offload pin_memory copy", "fp32_partitioned_groups_flat", "offload ratio pinned", "pageable offload"). Related, not duplicates: #8207 (merged: pinned the ZeRO-3 gradient offload buffers, default `pin_memory: true`; master partitions unchanged), #8256 (merged: routes pin sites through `get_accelerator().pin_memory()`; this PR uses the same accelerator calls), #8618 (closed: a CPU-offload benchmark). Sibling path: ZeRO-1/2 pins its master partitions ([stage_1_and_2.py#L563-L575](https://github.com/deepspeedai/DeepSpeed/blob/0f8c9d7f13fc5d16479376392c67e7ca82ae9eed/deepspeed/runtime/zero/stage_1_and_2.py#L563-L575)). #5301 added the pinning and copied the pinned master to the device before the cast; ZeRO-1/2 now casts into a pinned bf16 buffer of each partition's size instead ([#L3093-L3096](https://github.com/deepspeedai/DeepSpeed/blob/0f8c9d7f13fc5d16479376392c67e7ca82ae9eed/deepspeed/runtime/zero/stage_1_and_2.py#L3093-L3096)). This PR uses #5301's form, which needs no second pinned buffer. ## Root cause What `offload_optimizer.pin_memory: true` pins in ZeRO-3 on master, checked at runtime on the earlier version of this PR, which pins the same buffers as master (one GPU, bf16, DeepSpeedCPUAdam, after `deepspeed.initialize` and after one step): | Buffer | Where | Pinned | |---|---|---| | `fp16_partitioned_groups_flat` (bf16 parameter partitions) | device | (device memory) | | `fp32_partitioned_groups_flat[i].grad` (float32 gradient partitions) | host | yes ([stage3.py#L1339-L1340](https://github.com/deepspeedai/DeepSpeed/blob/0f8c9d7f13fc5d16479376392c67e7ca82ae9eed/deepspeed/runtime/zero/stage3.py#L1339-L1340)) | | `fp32_partitioned_groups_flat` (float32 master partitions) | host | no | | Adam `exp_avg`, `exp_avg_sq` | host | no | With `offload_param`, `fp16_partitioned_groups_flat` is on the host and already pinned: it is a view of `param_groups_fp16_flat_cpu_memory` ([stage3.py#L906](https://github.com/deepspeedai/DeepSpeed/blob/0f8c9d7f13fc5d16479376392c67e7ca82ae9eed/deepspeed/runtime/zero/stage3.py#L906)). `_create_fp32_partitions` builds each CPU master partition with `.to('cpu', dtype=...)` and `.clone().detach()` ([stage3.py#L1157-L1160](https://github.com/deepspeedai/DeepSpeed/blob/0f8c9d7f13fc5d16479376392c67e7ca82ae9eed/deepspeed/runtime/zero/stage3.py#L1157-L1160)), which gives pageable memory. After the CPU optimizer step, `_reassign_or_swap_out_partitioned_parameters` runs `fp16_flat.data.copy_(fp32_flat.data)`: a blocking copy from CPU float32 to device bf16, which PyTorch does by casting on the CPU into a new pageable bf16 tensor of the whole partition, copying it from pageable memory and synchronizing. Found in a profile (MiniMax-M3, first 8 layers, 8 H200, 75% of the optimizer on the CPU, rank 0): 13 bare `aten::copy_` host-to-device copies of 576 MiB per step from pageable memory, at 3-7 GB/s, 1,648 ms per step. ## The change (`deepspeed/runtime/zero/stage3.py`, 43 added, 4 removed lines) | Where | What | |---|---| | `_create_fp32_partitions` | with `offload_optimizer.pin_memory`, pins each CPU master partition after `clone().detach()` (`get_accelerator().pin_memory`) | | `_reassign_or_swap_out_partitioned_parameters` | uses the new copy for a master pinned there whose destination is on the device; every other case keeps `dst.copy_(src)`: `pin_memory` false, sub-groups kept on the device (`ratio` < 1), NVMe-swapped sub-groups | | new `_copy_pinned_master_to_device(dst, src)` | per piece of 2**26 elements: `src[piece].to(device, non_blocking=True)`, then `dst[piece].copy_(...)`, which casts on the device; at the end, records an event and waits for it | | new `_master_pinned_at_construction(sub_group_id)` | the condition above (pinned, on the CPU, not NVMe-swapped); also used by `_unpin_offload_buffers` | | `__init__`, `_unpin_offload_buffers` | the piece size; release of the pinned masters (matters under `DS_PIN_MEMORY_BACKEND=native`) | Design notes: - **Why the cast moves to the device:** pinning alone does not make the copy faster. A blocking `copy_` from pinned float32 into bf16 still casts on the CPU into a new pageable tensor (136.2 ms against 136.5 ms from pageable memory, 300M elements, in an earlier measurement on the same setup). - **Why pieces:** `.to(device)` of a whole sub-group makes a float32 device temporary of the sub-group's size, 3.73 GiB at the default `sub_group_size` of 1e9. Pieces of 2**26 elements keep it at 0.25 GiB at the same speed. Measured on one H200, torch 2.8.0, 12 CPU threads, copying one partition back into its bf16 device partition: | Copy back | 300M elements | 1e9 elements | Extra device memory at 1e9 | |---|---|---|---| | `copy_` from pageable float32 (master) | 132.6 ms | 483.3 ms | 0 | | pinned float32, `.to(device)`, then `copy_` (#5301) | 22.1 ms | 73.6 ms | 3.73 GiB | | the same in pieces of 2**26 elements (this PR) | 22.1 ms | 73.7 ms | 0.25 GiB | | earlier version of this PR: two pinned 2**26-element bf16 staging buffers | 19.4 ms | 73.6 ms | 0 | - **Same values:** the device cast rounds to nearest even, as the CPU cast does. The new test compares the trained parameters bitwise against `pin_memory: false`. - **Stream order:** the pieces go on the current stream. The final wait keeps what callers had before: the partition is on the device when the function returns. The next forward's parameter all-gathers run on their own stream, and that stream does not wait for the main stream. - **Host synchronization:** one wait per sub-group, at the end of its copy. The old path also blocked for the whole copy. ## Measured Qwen2.5-32B (all 64 layers, real weights), 1 node of 8 H200 (143,771 MiB, driver 580.126.09), torch 2.8.0+cu126. Settings: - ZeRO-3 with `offload_optimizer: {device: cpu, pin_memory: true}`; - float32 master weights and Adam states, DeepSpeedCPUAdam with 12 threads per rank; - bf16, one 8,192-token PG-19 sequence per GPU, activation checkpointing, lr 1e-6; - mean of steps 3-8; one run per column; all runs in one job on the same node. | | master `0f8c9d7f1` | earlier version of this PR (staging buffers) | this PR | |---|---|---|---| | step time | 12.08 s | 10.20 s | **9.42 s** | | step time with `wall_clock_breakdown` | 11.96 s | - | 9.42 s | | rank 0: forward / backward | 1.19 / 3.71 s | - | 1.19 / 3.67 s | | rank 0: `step` / of which the sub-group loop (`optimizer_step`) | 7.06 / 6.05 s | - | 4.56 / 3.52 s | | peak device memory (rank 0) | 34.66 GB | 34.66 GB | 34.66 GB | | loss at steps 1 / 2 / 8 | 0.022649 / 0.976290 / 1.746374 | 0.022649 / 0.976256 / 1.746135 | 0.022649 / 0.976285 / 1.746306 | The losses differ from step 2 by run-to-run variation: two master runs differ by as much (step 2 is 0.976235 with `wall_clock_breakdown` and 0.976290 without it). The flash-attention backward adds with atomics. The bitwise check is the unit test below. This version is faster than the earlier one in training (9.42 against 10.20 s), although the two copies take the same time alone at 1e9 elements. A likely reason, not measured: the earlier version cast on the CPU, which the CPU optimizer of all 8 ranks on the node also uses; this version casts on the device. <details><summary>Earlier version of this PR (staging buffers) on MiniMax-M3 (DeepSpeed master + unmerged MiniMax-M3 model code: #8658, #8652)</summary> Settings for both: ZeRO-3 + AutoEP (expert parallel 8), real weights, each 65,536-token sequence split over 4 GPUs, 65% of the optimizer on the CPU (`offload_optimizer.ratio` 0.65), float32 master + Adam. - 4 nodes x 8 H200, all 60 layers, same nodes, back to back: 126.9 s without, 123.9 s with (mean of steps 3-10); peak device memory 132.4 GB in both. - 1 node, first 8 layers: 14.09 s without, 12.46 s with; losses equal at every printed step. These were not repeated with this version. </details> ## Tests Updated: `TestZero3OffloadPinnedCopyBack` in `tests/unit/v1/zero/test_zero.py` (2 ranks, `ratio` 1.0 and 0.5, `sub_group_size` 20,000 so partial offload has sub-groups on both sides). - It trains 3 steps with `pin_memory: false` (direct copy) and with `pin_memory: true`, and requires the gathered parameters to be bitwise equal. - It requires every CPU master partition to be pinned with `pin_memory: true`, and none with `pin_memory: false`. - It requires that the pinned run used the new copy. Otherwise the test would compare the direct copy with itself. - The piece size is set to 1,000 elements, so each sub-group is copied in many pieces. Commands, from `tests/`, on 1 node of 8 H200 (143,771 MiB, driver 580.126.09), torch 2.8.0+cu126, pytest 9.1.1: ``` # master 0f8c9d7f1 + only the new test: 2 failed, AttributeError: no attribute '_copy_pinned_master_to_device' # master + this PR: pytest unit/v1/zero/test_zero.py -k TestZero3OffloadPinnedCopyBack # 2 passed pytest unit/v1/zero/test_zero.py unit/v1/zero/test_zero_user_backward.py unit/v1/pin_memory/ \ -k "offload or Offload or pin or Pin" # 61 passed, 1 skipped (torch_npu) pytest unit/v1/zero/test_offload_states.py # 128 passed pytest unit/runtime/zero/test_zero_offloadpp.py unit/v1/pin_memory/ # 47 passed, 1 skipped (torch_npu) ``` The test is in `tests/unit/v1`, so PR CI selects it: the `modal-torch-latest` selector (`ci/tests_fetcher.py`) picked `tests/unit/v1/zero/test_zero.py` for the earlier version, which changes the same two files. ## Compatibility | Cell | Validated | |---|---| | ZeRO-3, CPU optimizer offload, `pin_memory` true, whole optimizer on the CPU | yes: new test (ratio 1.0); Qwen2.5-32B | | ZeRO-3, partial offload (`ratio` < 1) | yes: new test (ratio 0.5). MiniMax-M3 on 1 and 4 nodes with the earlier version only | | ZeRO-3, `pin_memory` false | unchanged path (not pinned, direct copy); the new test's reference | | bf16 | yes. fp16: not run (same path; the device cast also rounds to nearest even) | | ZeRO-3, NVMe optimizer offload | not run. Unchanged: NVMe-swapped sub-groups are not pinned here and keep the direct copy | | ZeRO-3 + ZenFlow, `DS_PIN_MEMORY_BACKEND=native` | not run | | ZeRO-1/2 | not affected | | Accelerators other than CUDA | not run; only `get_accelerator()` calls are used (`pin_memory`, `is_pinned`, `Event`, `unpin_memory`) | No config, checkpoint format or public API change. ## AI assistance This change, its tests and this description were prepared with AI assistance (Claude). --------- Signed-off-by: pengdurice <pengduhit@gmail.com> Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>