[LiRunGuo]

Status: draft for discussion, not for merge as it is. Part of RFC #9026, which has not been accepted yet. I'm opening it so that the two-stage relay, the pipelined execution and the training-state move can be discussed on code that runs, next to #9045. The split I'd propose for upstream is at the end; each piece would be its own PR after the RFC is settled.

This PR does not propose a new rollout engine. The NativeRollout in the third commit is only the destination its tests write into; it is not registered with build_rollout and stays out of upstream. For the hybrid engine, the re-shard at the training/rollout boundary (ZeRO-3 partitions into the inference layout of HybridEngineRollout) is #9045's AffineWeightSync, which takes maps only and can be called from HybridEngine.eval(). What is here are the mechanisms beyond a fixed trainer on one node: the trainer's own layout changing with its optimizer state, ZeRO-3 over a subset of the ranks (#9076), and moves across nodes.

What it is. Moving model state between two sharded layouts in GPU memory, without a checkpoint and without gathering a whole parameter on any rank:

• weights from a ZeRO-3 trainer into a tensor-parallel inference copy on other ranks, or from one inference layout to another (other ranks, another TP degree);
• a ZeRO-3 trainer's fp32 master weights, optimizer moments and step counters from one live engine to another over other ranks (for example DP4 → DP2 → DP4 inside the same processes).

The four commits

1. [zero] ZeRO-3's rank-agreement checks and barriers run over the partition group. 31 lines added and 19 removed in five core files (stage3.py, partition_parameters.py, partitioned_param_coordinator.py, parameter_offload.py, zero/utils.py).
   • A ZeRO-3 engine whose data-parallel group is a subset of the world hung or failed its trace checks: the checks broadcast from global rank 0 over the world, and the barriers waited for every rank.
   • With the default group the same ranks take part as before: assert_ints_same_as_other_ranks(ints, group=None) compares over the world, and an engine's default dp_process_group holds every rank.
2. [rollout] relayout.py, affine_plan.py, zero3_state.py (new files; no existing file changes).
   • RelayoutPlan: both sides are declared per rank as affine pieces. Source and destination ranks may be disjoint, equal or overlapping, and a weight may have several holders. Planning runs identically on every rank from the gathered descriptions; execution is one all_to_all_single per bucket of weights and stage.
   • topology="relay": a block that crosses nodes is fetched once per node and copied to the node's other ranks in a second stage. spread=True has the replicas of a node fetch bands of each large block and exchange them, so the traffic uses every NIC of the node. Over buckets the stages are pipelined: stage one of a bucket travels with stage two of the bucket before.
   • planner="affine" (affine_plan.py): the same moves planned by plan_transfer. A relay is two plans through a relay map whose shards are bands of the destinations themselves, so there is no relay buffer. Two cases still fall back, per weight, to the module's own rectangle planner: holders in two layouts (waiting for Describe a parameter held in several layouts as one affine map #9052) and partial targets (RFC [RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026, open question 1).
   • Zero3StateMove(src_engine, dst_engine): moves the training state through zero_optimizer_storage_affine_maps ([AffineIR] Affine storage maps for live ZeRO parameters and optimizer states #8885). release_training_state / restore_training_state give back and re-allocate the inactive engine's memory.
3. [rollout] native/ and the tests. A small tensor-parallel decoder whose weights are views of one buffer per rank, written in place, so CUDA graphs captured over one version of the weights stay valid for the next. NativeRollout is the destination the tests move weights into. RolloutBatch gains an optional response_logprobs.
4. [rollout] NativeRollout leaves build_rollout, which is as on master again. The tests construct NativeRollout directly.

How it relates to open work

• [rollout] Affine weight-sync building block: ZeRO storage into a tensor-parallel inference copy, shard to shard #9045 (AffineWeightSync) is the same weight move on the IR alone (plan_transfers + execute_transfers), validated on one node. This draft holds what [rollout] Affine weight-sync building block: ZeRO storage into a tensor-parallel inference copy, shard to shard #9045 leaves out: the relay, the pipelined execution, holders in two layouts, and the training state. native/arena.py duplicates [rollout] Affine weight-sync building block: ZeRO storage into a tensor-parallel inference copy, shard to shard #9045's WeightArena; only one of them should land.
• feat(rollout): add AutoTP rollout engine with in-memory weight sync #9023 (AutoTPRollout) is a tensor-parallel rollout engine already under review. NativeRollout overlaps with it and is not proposed to land next to it. It is here because the relayout and trainer-switch tests run through it.
• [checkpoint] Expose the affine transfer schedule as schedule_transfers #9027 (schedule_transfers) is what the pipelined execution should be built on: with the schedule exposed, stage two of one bucket can share a collective with stage one of the next. The executor in relayout.py does this with its own bucketing today.
• Describe a parameter held in several layouts as one affine map #9052 (a parameter held in several layouts as one map) removes one of the two fallbacks to the rectangle planner.
• [AffineIR] Calibrate affine transfer links per path (Plan analysis & calibration Part-2) #8966 / [AffineIR] Analyze affine transfer batches before execution (Plan analysis & calibration Part-1) #8967 (analysis and calibration): the relay is chosen by a flag here; the RFC's point is to choose it by the calibrated cost model.
• I found no open PR that covers ZeRO-3 over a subgroup of the world (commit 1) or a live move of the optimizer state.

Validation

One node with 4× RTX 3090 (24 GB), driver 595.58.03, torch 2.13.0+cu130, NCCL 2.29.7. Commands from tests/, each as python -m pytest -q -rfEs <files>. The two rollout rows were rerun on the head commit; the affine and ZeRO rows and the end-to-end run below are from the third commit (the fourth changes only rollout/__init__.py and a docstring):

Tests	Result
New: unit/runtime/rollout/test_relayout.py, test_native_decoder.py, test_native_rollout.py	167 passed
All of unit/runtime/rollout	262 passed, 1 skipped, 1 failed (below)
Affine IR: unit/checkpoint/test_affine_shard_map.py, test_affine_transfer_plan.py, test_affine_transfer_executor.py, test_affine_transfer_executor_external_group.py, test_affine_replica_policy.py, test_zero_affine.py	350 passed, 1 skipped (CPU-only test)
ZeRO, for commit 1: unit/v1/zero/test_zero.py, unit/v1/zero/test_zero_init_pg_contradiction.py, unit/runtime/zero/test_zero_context.py	136 passed, 1 skipped (marked "not working" on master)

• The one failure is test_hybrid_engine_rollout.py::TestZero3SynchronizedEarlyStop. It needs the transformer_inference extension, which does not build on this machine (the CUDA toolkit there has no cuda_profiler_api.h); the same cause skips one more test in that file. Neither touches this change.
• A ZeRO-3 engine over a subgroup of the world (commit 1) is exercised by TestNativeRolloutApart and TestTrainerSwitch in test_native_rollout.py. TestTrainerSwitch also compares training across two trainer layouts with a plain-torch AdamW reference.

End to end, a ZeRO-3 trainer going DP4 → DP2 → DP4 inside the same four processes (Qwen3-0.6B, bf16, Adam; 310 parameters; 6.66 GiB of fp32 master weights and two moments). Three rounds per planner, a training step after every move:

Planner	Plan, once (s)	DP4 → DP2 (s)	DP2 → DP4 (s)	Peak on a DP2 rank (GiB)
rectangle	0.04	0.65	0.80	9.41
affine (plan_transfer)	0.37	0.62	0.79	9.41

• 12 of 12 moves are bitwise equal by a layout-independent digest of every parameter's master weight and moments, taken before and after each move. The step counters are carried. One flipped bit in one moment is caught and named, and the next check is clean.
• Times are medians of 3 on RTX 3090s, which exchange data over PCIe, so they say little about faster links. A DP2 rank holds 3.17 GiB before the move and 5.39 GiB after; the peak covers both layouts plus the receive buffers.
• The second engine is built in 2.3 s inside the running processes.
• The benchmark script is not part of this PR.

What I could not run, and known limits

• ZeRO-3 only, bf16 with fp32 master weights, Adam / AdamW, no parameter or optimizer offload, no NVMe. ZeRO-1/2, fp16 loss scaling, pipeline parallelism and MoE / expert parallelism are not covered.
• CUDA with NCCL only. Not run on CPU (gloo) or on another accelerator. The new code goes through get_accelerator() and deepspeed.comm, but the native decoder's graph capture is CUDA-specific.
• After Keep the parameters a model names for fp32 in fp32 under ZeRO-3 #8843 a ZeRO-3 model may hold partitions of two dtypes (keep_in_fp32_modules). NativeRollout assumes one dtype and has no test for a mixed model.
• LoRA / PEFT students are refused by NativeRollout.bind (NotImplementedError).
• Both layouts are resident on a rank during a state move, so a rank needs room for its old and new partitions at once.
• Open issue on two nodes, below this code. With a 32B model on 2 nodes × 4 H100, the bitwise check after a cross-node move failed in 2 of 162 checked moves. An element-wise dump shows blocks on a 128 KiB grid holding the bytes of another block of the same all_to_all_single message. Replaying the same plans on CPU is exact for all 707 weights. A stress test of the bare collective (torch.distributed.all_to_all_single with self-describing data, no DeepSpeed, no model) reproduces the same signature at the message sizes of these moves: 6 misplaced 128 KiB blocks in 370 TB between two nodes, with NCCL's default settings (1 in 121 TB) as well as with NCCL_IB_QPS_PER_CONNECTION=4 or NCCL_NCHANNELS_PER_NET_PEER=8; none in 159 TB in an earlier run whose messages all stayed at or below 256 MiB per pair. So the cause is below this executor (torch 2.13.0, NCCL 2.29.7, IB on that cluster); which component it is has not been established. The digest after each move is what detects it. An 8B model's digests were clean in all 39 two-node configurations I ran.
• No documentation page and no tutorial yet. They belong with the pieces that go upstream (the weight path would extend affine-transfer.rst from [rollout] Affine weight-sync building block: ZeRO storage into a tensor-parallel inference copy, shard to shard #9045).

Size, and the split I'd propose. 2,311 new lines under deepspeed/ and 1,052 in tests is too much for one review. For upstream I would not send it in this form:

1. Commit 1 as its own small PR (it is useful to any ZeRO-3 engine over a subgroup): now [zero] Run ZeRO-3's rank-agreement checks and barriers over the partition group #9076.
2. The training-state move rewritten on the IR alone (plan_transfers + execute_transfers, as [rollout] Affine weight-sync building block: ZeRO storage into a tensor-parallel inference copy, shard to shard #9045 does for weights), with DP4 → DP2 → DP4 as its test.
3. The relay as an addition to RFC [RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026: two plan_transfers calls through a relay map. The planner needs nothing new; the open questions are where the relay-map constructor lives and how the choice meets the cost model of [AffineIR] Calibrate affine transfer links per path (Plan analysis & calibration Part-2) #8966 / [AffineIR] Analyze affine transfer batches before execution (Plan analysis & calibration Part-1) #8967.
4. The pipelined execution after [checkpoint] Expose the affine transfer schedule as schedule_transfers #9027.
5. The rectangle planner stays out of upstream. It goes away once Describe a parameter held in several layouts as one affine map #9052 and partial targets cover its two remaining cases.
6. NativeRollout stays out; the destination upstream is feat(rollout): add AutoTP rollout engine with in-memory weight sync #9023's engine through [rollout] Affine weight-sync building block: ZeRO storage into a tensor-parallel inference copy, shard to shard #9045's sync.

Two questions for the RFC: does this split look right, and is the relay-map formulation in affine_plan.py what #9026 should describe?

AI assistance: AI assistance was used. The code, tests and this description were drafted with an AI coding agent (Claude Code). The submitter reviewed every changed line and ran the tests above.

Maintenance contact: @LiRunGuo. cc @Achyuthan-S @0z5a @jinyouzhi

[Achyuthan-S]

Thanks @LiRunGuo, this is great to see running end to end. A DP4 → DP2 → DP4 state move in 0.6–0.8 s with every move digest-checked is the comparison the RFC needed against the checkpoint path: the same Qwen3-0.6B ZeRO-3 changes through a universal checkpoint took me 46 s (DP4 → DP2) and 51 s (DP2 → DP4) on 4× A100, most of it the relaunch.

To your two questions, as the RFC's author:

The split: +1.

• Commit 1 on its own is a good first PR; it is useful beyond this work.
• Moving the state move onto plan_transfers + execute_transfers also routes its replica choice through [checkpoint] Batch replica selection for affine transfer plans #8936's batch policy, which is where the cost-aware policy plugs in. Today affine_plan.py calls plan_transfer per weight with its own selector.
• Agreed that the rectangle planner and NativeRollout stay out. Describe a parameter held in several layouts as one affine map #9052 is for the two-layout fallback and is ready for review; it covers row splits, and column splits still need @0z5a's planner change ([RFC] Two-stage relay for cross-node weight sync, priced by a calibrated cost model #9026, Q8).

The relay-map formulation: yes. Bands of the destinations as relay shards is the alternative already recorded under open question 1 in #9026, from your earlier comment: no partial targets and no relay buffer. On the execute_transfers path it still needs #8809 to accept the stage where a relay rank's source lies inside its own target (_validate_aliases). The RFC lists relay_map and the pipelined execution under @0z5a. @0z5a, would you take this formulation as the constructor's definition, or keep relay_map(source_map, relay_locations, min_split_bytes) and add bands as an option?

The choice. I'll take what is a flag here: whether a weight goes direct, through a relay, or through a relay with spread, and the thresholds (min_relay and min_split, both 65,536 elements now), chosen by the calibrated model in §3 as the cost-aware policy on #8936. Your relay becomes one of the candidates the policy prices, and your InfiniBand runs are its first check.

Before and after. I can run your state move on the same 4× A100 node as my checkpoint-path numbers, with the same model and changes, so there is one comparison on one machine. I'll update #9026's body with the relay-map formulation and the split once @0z5a has weighed in.

[jinyouzhi]

@LiRunGuo, thanks for the split at the end of the description. A note on step 2 (the training-state move rewritten on plan_transfers + execute_transfers).

The description says no open PR covers a live move of the optimizer state. deepspeedai/DeepSpeedExamples#1021 (open since 10-22) already does this on the IR, beyond ZeRO-3:

• this comment moves bf16 params, the fp32 master and the Adam moments through the shipped zero_affine / autotp_storage_affine_maps(..., optimizer_states=True) maps, mid-training, across:
  • training world size: Z3x4 → Z3x2 → Z3x4, and 4 → 8 → 4;
  • ZeRO stage: Z3 ↔ Z2, Z3 ↔ Z1, and stage and DP changed together (for example Z2x8 → Z3x4 → Z1x8);
  • AutoTP size: TP2 ↔ TP4 under ZeRO-1/2/3, and up to TP8 combined with a stage change;
  • AutoEP size: EP2 ↔ EP4 ↔ EP8 under ZeRO-1/2/3.
• this comment adds ZeRO-3 with aliased masters (ZeRO-3: reuse the weight partition as the master copy when it already has the master dtype #8842).
• Every switch is checked against an oracle that shares no code with the IR: safe_get_full_* / safe_set_* plus a plain all-gather over the TP or EP group. The 60-step loss curves and the final logits are compared bitwise between the two paths.
• Hardware: 8× NVIDIA H20.

Suggestion. Since step 2 is a rewrite anyway, could it target this wider matrix rather than ZeRO-3 only? The IR already describes every one of these layouts. What differs per case is:

• which maps the engine exposes;
• what has to be refreshed after the move (low-precision partitions, optimizer scalars, counters).

Zero3StateMove would become one case of a move between two engines. The scripts in #1021 (z3x.py, tpep.py) can be turned into unit tests, with the gather path as the oracle. I'm glad to help with that part, or to take the ZeRO-1/2 and AutoTP cases on top of your ZeRO-3 piece, whichever you prefer.

One gap the TP cases expose. A second engine with another TP size overwrites the process-global groups in deepspeed/utils/groups.py (_TENSOR_MODEL_PARALLEL_GROUP, _DATA_PARALLEL_GROUP, mesh_device, mpu). #1021 works around this by saving and restoring them per engine, which is not something core should do. #9076 fixes the same class of problem for ZeRO-3's checks and barriers; for TP and EP it would need each engine to hold its own groups. I think that is worth agreeing on in #9026 before step 2 goes beyond DP changes.

This comment was drafted with AI assistance and checked by me against #9074's description and the #1021 comments.

[LiRunGuo]

Thanks @Achyuthan-S.

One machine, both paths. Thanks for the offer. The benchmark script is not in the PR and I'm not publishing it yet, so I can't hand you the run for now. I have the same run queued on a 4× A100 node and will post it here when it finishes; that is another machine than yours, so it puts the same GPU model next to your checkpoint numbers, not one machine. @jinyouzhi's run on #9026 already gives one pair on one machine, measured independently (4× L20, ZeRO-1 DP4 → ZeRO-3 DP4: 0.62 s in the running processes against 27.4 s through a universal checkpoint). The difference to this PR is the check: sums there, a bitwise digest here.

Per-weight planning. Right: affine_plan.py calls plan_transfer once per weight with its own selector. The selector keeps its load across the weights of a move, so the balance is over the batch, but it is not #8936's policy interface. Item 2 of the split moves the state move onto plan_transfers. The order the selector implements today is the target itself, then a holder on the target's node, then the caller's order, then the least loaded. That is local_first_least_loaded with two keys before the load, so a policy that expresses it needs each location's node, or the link class from #8966.

_validate_aliases. Agreed that this is what the relay map needs from #8809. In stage two a relay rank's source is a band inside its own destination shard. The check compares whole buffer ranges, so it refuses the batch, although no element is both read and written: the band onto itself is the same memory, and the other bands the rank receives lie outside it. The narrowest change I see is to drop a local segment whose source and target are the same memory with the same strides, and to check overlap on the regions the remaining segments read and write. relayout.py's executor relies on the same property today. @0z5a, your call.

The choice. Good, it is yours. Data for it, on 2 nodes × 4 H100 with Qwen3-8B (a ZeRO-3 trainer on node 0, replicas on node 1): I ran 30 configurations with both planners: three syncs from the trainer (into TP4, TP2 × 2 and TP1 × 4) and seven transitions (four layout changes, read from the old replica or from both sources), each direct, relayed, and relayed with spread.

• Every move passed a per-replica digest (270 checks).
• The two planners send the same inter-node bytes in all 30 configurations, and their medians agree within 1% in 22.
• Sync from the trainer, relayed with spread: 0.176 / 0.186 / 0.213 s for TP4 / TP2 × 2 / TP1 × 4 with plan_transfer, 0.177 / 0.187 / 0.213 s with the rectangle planner.

Both planners run on this PR's executor, not on #8809. I can post the per-configuration table (bytes inside and across nodes, times) on #9026 as points for §3.

The two-layout fallback. In those runs 217 of 399 weights planned from both sources fell back to the rectangle planner. I'll rerun that case on #9052, and on #9070 for the column case, and report how many each removes.

#9076 now carries its regression test and is marked ready for review.