Add device-agnostic ZeRO-3 CPU-offload end-to-end benchmark by delock · Pull Request #8618 · deepspeedai/DeepSpeed

delock · 2026-09-21T07:30:05Z

What does this PR do?

Adds benchmarks/pin_memory/zero3_offload_bench.py, a companion to the existing h2d_d2h_bench.py:

h2d_d2h_bench.py measures isolated H2D/D2H copy bandwidth per pinning arm;
this new script measures the end-to-end effect on ZeRO-3 CPU-offload training step time, comparing DS_PIN_MEMORY_REGISTER_DEVICE=1 vs 0 on the native pin backend (DS_PIN_MEMORY_BACKEND=native) with offload_optimizer + offload_param (both cpu, pin_memory).

Design

Each arm (registered / unregistered) runs in its own subprocess with a fresh rendezvous port, so device state, JIT caches, and any lingering registrations from one arm can never contaminate the other.
Model selection: a real HF architecture via --model (instantiated from the config with random weights, which suffice for timing) or a self-contained synthetic stack via --hidden / --layers (no network access needed).
Device-agnostic: the device API is resolved via getattr(torch, device.type), so the same script runs on cuda and xpu; synchronization and peak-memory reporting adapt automatically.
Prints a per-arm table (avg/min step time, GPU peak memory), the speedup, per-step saving, throughput in tok/s, and a machine-readable DRIVERRESULT= JSON line.

Usage

# real model architecture
python benchmarks/pin_memory/zero3_offload_bench.py --model Qwen/Qwen2.5-7B --batch 4 --seq 512

# synthetic stack, no network needed
python benchmarks/pin_memory/zero3_offload_bench.py --hidden 2048 --layers 12 --batch 4 --seq 128

Note: the native backend mlocks host memory; multi-GB models need a raised RLIMIT_MEMLOCK (ulimit -l) or root.

Notes forReviewers

Follow-up to the pin-memory registration hooks (#8283) and the XPU implementation (#8315): this benchmark is the end-to-end counterpart of the copy-level numbers, runnable on any accelerator that implements the native pin backend.

🦉

Adds benchmarks/pin_memory/zero3_offload_bench.py alongside the existing h2d_d2h_bench.py copy-bandwidth micro-benchmark. While the latter measures isolated H2D/D2H throughput, this script quantifies the end-to-end effect of native host-memory registration on ZeRO-3 CPU-offload training step time: it drives a full ZeRO-3 loop (offload_optimizer + offload_param, both with pin_memory) on the native pin backend and compares DS_PIN_MEMORY_REGISTER_DEVICE=1 vs 0. Each arm runs in its own subprocess with a fresh rendezvous port so device state never leaks between arms. Supports real HF architectures (--model, random weights suffice for timing) or a self-contained synthetic stack (--hidden/--layers); the device API is resolved via getattr(torch, device.type) so the same script runs on cuda and xpu. Signed-off-by: Ma, Guokai <guokai.ma@intel.com>

sfc-gh-truwase · 2026-09-21T20:03:39Z

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0


Can we move to https://github.com/deepspeedai/DeepSpeedExamples/tree/master/benchmarks

delock · 2026-09-22T09:19:11Z

Sounds good — moved it to DeepSpeedExamples in deepspeedai/DeepSpeedExamples#1013 (script unchanged, plus a README). Closing this in favor of that PR.

…ed buffer (deepspeedai#8840) # ZeRO-3: pin the offloaded master weights and cast them on the device when copying them back **Part of:** deepspeedai#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 deepspeedai#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 deepspeedai#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: deepspeedai#8207 (merged: pinned the ZeRO-3 gradient offload buffers, default `pin_memory: true`; master partitions unchanged), deepspeedai#8256 (merged: routes pin sites through `get_accelerator().pin_memory()`; this PR uses the same accelerator calls), deepspeedai#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)). deepspeedai#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 deepspeedai#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_` (deepspeedai#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: deepspeedai#8658, deepspeedai#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 0f8c9d7 + 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>

delock requested a review from tjruwase as a code owner September 21, 2026 07:30

sfc-gh-truwase reviewed Sep 21, 2026

View reviewed changes

delock mentioned this pull request Sep 22, 2026

Add device-agnostic ZeRO-3 CPU-offload end-to-end benchmark deepspeedai/DeepSpeedExamples#1013

Merged

delock closed this Sep 22, 2026

pengdurice mentioned this pull request Oct 16, 2026

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

Merged

Navigation Menu

Repository navigation

Add device-agnostic ZeRO-3 CPU-offload end-to-end benchmark - #8618

Add device-agnostic ZeRO-3 CPU-offload end-to-end benchmark#8618
delock wants to merge 1 commit into
deepspeedai:masterfrom
delock:gma/zero3-offload-bench

delock commented Sep 21, 2026

Uh oh!

sfc-gh-truwase Sep 21, 2026

Uh oh!

delock commented Sep 22, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

delock commented Sep 21, 2026

What does this PR do?

Design

Usage

Notes forReviewers

Uh oh!

sfc-gh-truwase Sep 21, 2026

Choose a reason for hiding this comment

Uh oh!

delock commented Sep 22, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants