[jinyouzhi]

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 Preserve non-floating buffers in ZeRO checkpoint conversion #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 (Add AutoTP training to rollout weight sync example DeepSpeedExamples#1021): Qwen3-8B, 8× NVIDIA L20, AutoTP-2 training engine → AutoTP-4 rollout engine via autotp_storage_affine_maps + plan_transfer + executor ([checkpoint] Execute affine transfer plans with bounded peer staging #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).

[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-22T14:29:00.833416Z	2d8d873	Draft marked ready

ℹ️ 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: 2d8d873e4f

ℹ️ 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".

[0z5a]

Thanks @jinyouzhi, I went through the plan_transfer relaxation. From the planner side it is safe and I'm happy with it as written.

Why the invariant from #8624 still holds. For a flat shard, a piece's local footprint becomes the run [dest_offset, dest_offset + numel). _shards_tile still checks that the runs of one rank are pairwise disjoint and add up to the shard size. So two pieces still cannot write the same shard address, and a hole is still caught. The dense check (dest strides must equal the box's own row-major strides) refuses everything that is not one run. The source side keeps its existing check against the logical row-major strides. Nothing beyond a dense run is newly accepted.

I added adversarial cases locally against 42bafc5 and all of them behave as expected:

• refused: two flat pieces whose runs overlap, a hole in the shard, a run past the end of the shard, transposed / padded / gapped dest strides, a 1-D piece describing a matrix, a 2-D shard of a 3-D parameter;
• planned and checked element by element: flat storage on the target side (a ZeRO cut inside a row on the write side), flat storage on both sides (40 random TP×ZeRO-3 → TP×ZeRO-3 relayouts with partitions that start mid-row), ZeRO-2 replicas with a random replica_selector, and one map that mixes (numel,) and N-D shards across ranks;
• Qwen3-8B dense shapes (254 matrices), TP2×ZeRO-3 DP4 → TP4×ZeRO-3 DP2 through nest_affine_maps: 2,032 pieces per side, maps + plans in 0.28 s on a laptop CPU.

Executor (#8809). No change is needed for flat segments. Segments on the flat side carry the piece's own row-major strides. _validate_buffers already bounds the reach of those strides, and _region already adds the buffer's own storage offset. The shards that zero_param_storage_shard / zero_optimizer_storage_shard return are packed views, which is what _flat requires. I cherry-picked #8809 onto 42bafc5. The only conflict is in docs. On CPU (gloo, 4 ranks) the executor, planner and ZeRO-map tests pass together, apart from the four TestAutoTPStorageMaps cases in 1 below.

Three small requests:

1. TestAutoTPStorageMaps fails on CPU because of the test, not the code. With DS_ACCELERATOR=cpu (4 ranks, gloo), all four cases fail at the planned-copy check. I reproduced this on two machines. _all_shards gathers shard.detach().cpu(). On CPU that is the view itself, so the unpickled bias keeps its storage offset into the flat buffer (72 elements in the first case). torch.as_strided(..., segment.source_offset) then counts from the start of the storage and reads the weight's bytes. On XPU, .cpu() makes a fresh copy, which hides this. Gathering shard.detach().cpu().clone(), or adding storage_offset() the way the executor's _region does, makes all 27 tests in the file pass on CPU.
2. A target-side flat test. test_a_flat_shard_holds_a_dense_box_of_a_matrix covers flat storage only as the source. A restore or reverse sync writes into ZeRO storage, so a case with a flat target cut inside a row would pin the write direction. I can push mine if that helps.
3. Mixed shard ranks in one map. _normalized now lets one map have (numel,) shards on some ranks and N-D shards on others. That plans correctly in my test. If it is intended, a one-line test would document it. If not, it is cheap to refuse.

On the numbers: balanced beating lowest (0.49 s vs 1.05–1.35 s for repeated ZeRO-2 syncs) is a nice confirmation that holder choice belongs in a replica_selector policy and not in the planner.