Repository navigation
Conversation
autotp_rollout_e2e.py trains Qwen3-style dense models with AutoTP + ZeRO and moves the weights into rollout engines with a different TP size through a Universal checkpoint, through gathered parameters, and through ZeRO storage affine maps (plan_transfer + execute_transfers), checking every path bit for bit and recording time and peak memory per sync. A --replica_policy chooses which ZeRO-0/1/2 replica each target block is read from; the default reads locally when possible and otherwise from the least-loaded holder. Needs DeepSpeed with deepspeedai/DeepSpeed#8809 and deepspeedai/DeepSpeed#8885. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
A checkpoint saved below the sharding threshold has model.safetensors and no index; map every key to that file. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
ZeRO-3 keeps persistent parameters gathered between steps, and the forward reads that copy, so zero_param_storage_shard now rejects a gathered parameter. Release the gathered copies of both engines before each sync. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>
banxingmjj
pushed a commit
to openanolis/DeepSpeed
that referenced
this pull request
Oct 28, 2026
… states (deepspeedai#8885) ## Summary Describes, as affine maps in global ranks, where a live ZeRO engine stores each element of a parameter and its optimizer states, so `plan_transfer` can read from or write into training storage directly (e.g. training → rollout weight sync) without gathering full parameters first. - `zero_affine.py` - `zero_param_storage_affine_map`: low-precision replicas under ZeRO-0/1/2 (DP or expert-DP group), `ds_tensor` partitions under ZeRO-3, padding excluded. - `zero_optimizer_storage_affine_maps`: fp32 master + optimizer states (ZeRO-1/2, bf16 optimizer, ZeRO-3). - `zero_param_storage_shard` / `zero_optimizer_storage_shard`: flat views of ZeRO's own buffers matching the maps. - `autotp_storage_affine_maps`: AutoTP layout × ZeRO storage, one all-gather over the TP group. - `affine.py`: `nest_affine_maps` composes a layout map (TP/EP shard placement) with per-shard storage maps. A process's stored run of an N-D shard is described as at most `2·ndim−1` N-D boxes in the parameter's own shape, so plans stay small (row-parallel shards are not cut per row). - `affine_transfer.py`: `plan_transfer` accepts an N-D piece in a 1-D (flat) shard only when the piece is dense there; every other rank mismatch is still refused. - `affine_ir_spec.md`: phase-2 section on layout vs storage maps. - Fix `test_convert_checkpoint` `stat().size` → `st_size` (from deepspeedai#8515). No change to ZeRO, AutoTP or the executor. Offloaded optimizer states raise `NotImplementedError`; MiCS / hpZ are not covered. ## Tests - `tests/unit/checkpoint/test_zero_affine.py`, `test_affine_transfer_plan.py`: ZeRO-0/1/2/3 param + optimizer maps checked against ZeRO's own buffers (bitwise rebuild), AutoTP nest, N-D box count, flat-shard acceptance/refusal. - ≤2 ranks: Intel Arc B70 (XPU) ×2, pass. >2 ranks: CPU (gloo), pass. - End-to-end (deepspeedai/DeepSpeedExamples#1021): Qwen3-8B, 8× NVIDIA L20, AutoTP-2 training engine → AutoTP-4 rollout engine via `autotp_storage_affine_maps` + `plan_transfer` + executor (deepspeedai#8809), ZeRO-2 and ZeRO-3, weights bitwise equal after every sync. Plans are built once and reused. | Qwen3-8B, 8×L20 | ZeRO-2 gather | ZeRO-2 affine `lowest` | ZeRO-2 affine `balanced` | ZeRO-3 gather | ZeRO-3 affine | |---|---|---|---|---|---| | maps + plan (once) | - | 0.63 + 0.57 s | 0.70 + 0.70 s | - | 0.82 + 0.60 s | | first sync | 1.01 s | 1.16 s | 0.48 s | 1.71 s | 1.18 s | | repeated syncs | 0.78-0.82 s | 1.05-1.35 s | 0.49-0.55 s | 1.22-1.28 s | 0.45-0.74 s | | device memory added | 3.48 GB | 0.30 GB | 0.16 GB | 3.48 GB | 0.66 GB | | sent / local copy | - | 35.2 / 5.0 GB | 15.1 / 25.1 GB | - | 34.0 / 6.3 GB | | executor rounds | - | 96 | 41 | - | 32 | Each map has 3,192 pieces on both sides; before N-D boxes (one piece per stored row) the same ZeRO-3 maps had 592k source / 2.36M target pieces and planning took 207 s. Under ZeRO-2 every DP rank holds a replica: `lowest` is the planner's default holder choice, `balanced` is a `replica_selector` in the example (local copy first, else the least-loaded holder). --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
training/deepspeed_finetune_demo/autotp_rollout_e2e.pyandrun_autotp_rollout_e2e.sh: AutoTP + ZeRO training → rollout engines at a different TP size, three ways:ucp: DeepSpeed checkpoint → Universal checkpoint → load.gather: gather whole parameters, slice the rollout shard.affine:autotp_storage_affine_maps(ZeRO's own storage) →plan_transfer→execute_transfers; no full parameter is materialized.Rollout engines start from NaN weights; every path is checked bit for bit against whole parameters rebuilt with plain collectives. Plans are built once and reused for repeated syncs.
--replica_policy(lowest/local/balanced) chooses which ZeRO-0/1/2 replica each target block reads from.Companion to #1020 (AutoEP); this one covers dense AutoTP weights.
Depends on deepspeedai/DeepSpeed#8809 (executor) and deepspeedai/DeepSpeed#8885 (ZeRO / AutoTP storage maps).
Results
Qwen3-8B, 8× NVIDIA L20, train TP 2 → rollout TP 4,
--paths gather,affine, 3 training steps + 3 repeated syncs, all bitwise equal:lowestbalancedEach map has 3,192 pieces on both sides.