[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.

[jinyouzhi]

Thanks @0z5a for the thorough check, and for the adversarial cases. All three are addressed in 7f92794:

1. _all_shards now gathers shard.detach().cpu().clone(). You were right: on CPU, .cpu() returns the view itself, so the pickled bias kept its offset into the flat buffer. With DS_ACCELERATOR=cpu LOCAL_SIZE=4, all 27 tests in test_zero_affine.py pass; without the change, the 4 TestAutoTPStorageMaps cases fail as you described. They also pass on 8× B70.
2. Added test_flat_storage_is_written_as_well_as_read[flat_target_cut_mid_row]: the target is ZeRO-style flat storage cut at element 13, inside row 1. Its expected shards are slices of the flattened parameter, so the oracle does not depend on the maps.
3. Mixing flat and N-D shards in one map is intended: for example, ranks that keep a ZeRO flat partition next to ranks that hold an unpartitioned shard. [flat_and_matrix_shards] documents it.

On replica selection: I agree the planner should not hard-code a hardware cost model (links, contention, topology), because costs depend on the machine and on the goal. But the planner is the right place to apply one: given a pluggable cost or policy, it sees every request and can balance and split globally, which a per-request selector cannot. Concretely:

• The default. Lowest-rank is the worst policy we measured. Local-first, then fewest assigned bytes, needs no hardware information, stays deterministic, and roughly halved the ZeRO-2 sync time (0.49 s vs 1.05–1.35 s).
• The interface. A per-request selector cannot balance load (it sees one request at a time, so results depend on call order) and cannot split a region across holders. A batch form, where the selector receives every request and returns an assignment that may split regions while the planner checks the split tiles exactly, would let external policies do both.

It is not strictly necessary to achieve the absolute optimum right away; my view is that we should provide as much room as possible for the cost model to perform optimizations, while avoiding premature optimization. Happy to open a separate issue for this if you think it's worth discussing.

[jinyouzhi]

And I want share some follow-up ideas for the planner from #8623, which I hit while building storage maps. @0z5a Could you have a look on it? Thanks!

Check	Current approach	Geometric problem	Better algorithm
_shards_tile: a rank's pieces don't overlap in shard addresses	Pairwise, O(n²)	Box disjointness detection	Sweep line with an active set. In 1-D, sort and compare neighbours: O(n log n)
_validate_holders: one rank's logical boxes don't overlap, and the declared holders are correct	O(n²), plus n × ranks coverage sums	Disjointness detection plus box deduplication	Group by box in a hash map, so each rank's coverage query becomes a key lookup
_sources: distinct source regions don't overlap	O(R²)	Same as above	Sweep line
_covers_union: target pieces cover the whole parameter	Subtract box by box; fragments can grow exponentially	Klee's measure problem (volume of a union of boxes)	Deduplicate, verify disjointness, then compare the summed volume: O(n log n)
Source × target intersection	All S×T pairs	Red–blue box intersection reporting	Sweep line with an interval tree: O((S+T) log + K), where K is the number of segments produced

Structural fast paths

• Grid layouts (TP/EP, each axis cut independently): merge breakpoints per axis. Each intersection is the product of per-axis interval overlaps, O(Σ per-axis breakpoints + K), with no general-purpose data structure needed.
• ZeRO flat fragments: in N-D they become full rows plus partial rows at each end. A sweep sorted on axis 0 is close to linear.

[jinyouzhi]

Thanks @0z5a for the thorough check, and for the adversarial cases. All three are addressed in 7f92794:

1. _all_shards now gathers shard.detach().cpu().clone(). You were right: on CPU, .cpu() returns the view itself, so the pickled bias kept its offset into the flat buffer. With DS_ACCELERATOR=cpu LOCAL_SIZE=4, all 27 tests in test_zero_affine.py pass; without the change, the 4 TestAutoTPStorageMaps cases fail as you described. They also pass on 8× B70.
2. Added test_flat_storage_is_written_as_well_as_read[flat_target_cut_mid_row]: the target is ZeRO-style flat storage cut at element 13, inside row 1. Its expected shards are slices of the flattened parameter, so the oracle does not depend on the maps.
3. Mixing flat and N-D shards in one map is intended: for example, ranks that keep a ZeRO flat partition next to ranks that hold an unpartitioned shard. [flat_and_matrix_shards] documents it.

On replica selection: I agree the planner should not hard-code a hardware cost model (links, contention, topology), because costs depend on the machine and on the goal. But the planner is the right place to apply one: given a pluggable cost or policy, it sees every request and can balance and split globally, which a per-request selector cannot. Concretely:

• The default. Lowest-rank is the worst policy we measured. Local-first, then fewest assigned bytes, needs no hardware information, stays deterministic, and roughly halved the ZeRO-2 sync time (0.49 s vs 1.05–1.35 s).
• The interface. A per-request selector cannot balance load (it sees one request at a time, so results depend on call order) and cannot split a region across holders. A batch form, where the selector receives every request and returns an assignment that may split regions while the planner checks the split tiles exactly, would let external policies do both.

It is not strictly necessary to achieve the absolute optimum right away; my view is that we should provide as much room as possible for the cost model to perform optimizations, while avoiding premature optimization. Happy to open a separate issue for this if you think it's worth discussing.

@0z5a I'd like to clarify one point, since my earlier changes and demo may have given the wrong impression. I don't think replica_selector should pick replicas using only local information. It would be better to give the planner the full selection space and the information it needs, so it can decide with a global view. Please treat lowest / balanced / first only as demos showing how much the choice can affect performance, not as the design I'm proposing.

That said, this is just my personal view, so feel free to take it or leave it. Of course, the above are just my personal thoughts—feel free to take them as a reference.

[Achyuthan-S]

@jinyouzhi @0z5a On replica selection: +1 to a batch interface where the policy sees every request and may split regions, with the planner checking the split tiles exactly. That's the policy layer I took on in #8252 (replica policy above the plan, meeting at the hook), and your local-first / fewest-bytes measurement is a great baseline for it.

I'm on the correctness proof and exactness evaluation until mid-November. Could we write the interface up in an issue now so we agree on it, and I'll pick up the policy work after that? Happy to pair on it if either of you wants to start sooner.

[0z5a]

Thanks @jinyouzhi, and sorry for the slow reply.

Planner complexity. Your table matches the code, and one point makes it more than a speed question: each of these checks is guarded by PlanBudget.max_pair_checks (2^18 by default), so the default budget is also a scale limit. A parameter with more than 512 non-empty pieces on one side fails _validate_holders, which for a ZeRO-3 parameter means more than 512 ranks, and fewer when partitions are cut mid-row. plan_transfer also refuses more than 2^18 source × target pairs. A caller can raise the budget, but then pays the quadratic cost in Python.

The order I'd take them in:

1. _validate_holders: group by box in a hash map. It is the check that binds first, and the change is local.
2. Source × target: the grid fast path for TP/EP, and a sort on axis 0 for ZeRO flat fragments (full rows plus partial ends). A general interval tree can wait until a layout needs it.
3. _covers_union: dedupe, check disjointness, compare volumes. Target boxes that overlap without being equal need a decision: keep the current subtraction as their fallback, or refuse them. Either way a test should pin it.
4. _shards_tile and _sources: sort and compare neighbours. These are small in practice, since a rank holds only a few pieces of one parameter.

On the shared affine_geometry.py you proposed on Slack: +1 from the planner side, and I'm happy for you to take it. Two requests. Keep the pure move separate from the algorithmic changes, so the move can be reviewed as a no-op. And for each new check, use the current implementation as the oracle on randomized maps (the adversarial cases from my earlier comment included), plus a planning-time curve over the rank count. I'll review the planner side.

Replica selection. Thanks for the clarification; I had read lowest / balanced / first as the proposal. I agree with the batch form:

• the planner hands the policy every request with its validated candidates;
• the policy returns an assignment that may split a region;
• the planner checks that each request's parts tile it exactly and use only validated holders.

That keeps the policy pluggable, which is what I meant earlier, and gives it the global view. The default can be your local-first, then fewest-assigned-bytes rule, which is deterministic and needs no hardware information.

@Achyuthan-S I can open the issue with a strawman of that interface this week, so it is agreed before you pick up the policy in mid-November.