[Achyuthan-S]

Refs #9026 (open question 8), #9045.

Why

A rollout replica that changes placement can often keep a tensor-parallel shard that is still current, while the trainer's ZeRO-3 partitions hold the whole parameter. plan_transfer can read a block from any of its holders only if the source map's blocks are equal or disjoint and each block lists every holder. A map that simply lists both layouts' pieces breaks both rules, since a partition can start or end inside a row of the shard. So today the planner refuses it, and #9045 falls back to its own planner for these weights.

What

• union_affine_maps(maps) in deepspeed/checkpoint/affine.py returns one map of a parameter that several layouts hold at once.
  • Each piece is cut at the edges of every block held over it, in any layout (replicas within one layout may store a region whole or in parts), and each cut piece lists the holders from every layout.
  • A cut keeps the piece's strides and moves its offsets, so every location still stores exactly what it did.
  • Layouts use disjoint locations, one per buffer. TransferEndpoints already lets several locations map to one process.
  • More than two layouts are merged pairwise.
• subtract moves unchanged from affine_transfer._subtract into affine_geometry, the module [AffineIR] Move shared affine index and box geometry into affine_geometry #9010 added for shared box arithmetic. The planner imports it under its old name.
• Spec §10.8 and a section in affine-transfer.rst, after the replica-policy section from [checkpoint] Batch replica selection for affine transfer plans #8936.

Scope

Cuts along rows are planned today. A ZeRO-3 block of several rows cut at a column edge keeps its row-major strides and its partition's offset, so it is not dense in its flat shard, and plan_transfer refuses it. @0z5a plans to make the planner accept that case on top of #9011 (see #9026). The union already describes it exactly, and the tests check that at the map level.

Tests

test_affine_transfer_plan.py:

• A surviving row shard plus ZeRO-3 partitions, in 4 configurations (uneven TP sizes, one or two surviving ranks, 1 to 5 partitions), planned into a row-split and a column-split target. Checked:
  • each element's holders are exactly the union of both layouts' holders;
  • every location extracts what its own layout says;
  • every plan reproduces extract exactly;
  • the surviving shard is read.
• Three layouts merged in all six orders: the same holders, and exact plans.
• A column cut: holders and exactness at the map level.
• A layout whose replicas store a block whole and in two halves, merged with ZeRO-3 partitions: holders, and an exact plan into a column-split target.
• Refusals: no maps, a location in both layouts, different shapes, a transposed piece that is not a block, and a row-major run that wraps into the next row.

test_affine_transfer_executor.py: on 2 ranks, one process holds a ZeRO-3 partition and a surviving row shard under two locations. The plan runs through execute_transfer, and every target shard equals extract.

Not merging the holders, dropping the part of a block that a later block doesn't overlap, or leaving the destination offset unmoved each makes 7 of these 9 tests fail.

Validation

Rebased on master 988b2e3, which includes #9010 and #8936. The code and test changes are the same as before the rebase; the only conflict was this docs section next to #8936's. macOS arm64 CPU, Python 3.9.6, torch 2.8.0, gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4):

• All eight affine test files, now including [checkpoint] Batch replica selection for affine transfer plans #8936's test_affine_replica_policy.py: 412 passed.
• The four union configurations above, planned through [checkpoint] Batch replica selection for affine transfer plans #8936's plan_transfers with local_first_least_loaded and with first_holder into both targets: all 16 plans reproduce extract exactly.
• Against the previous commit (110ff4b, before the two fixes from review), on random layouts: 300 unions with a row-split surviving shard, 300 with a column-split one, and 300 of three or four layouts. Every element has the same holders, and every plan into a row- or column-split target has the same outcome as before: exact, except 82 of the column-split cases, which both commits refuse (the column cut under Scope).
• pre-commit passes on the changed files.

Before the rebase, on f9fe4df with #9010 applied:

• All seven affine test files: 346 passed.
• With [AffineIR] Check affine holders and source regions without pairwise comparisons #9011 on top: 648 passed. The only conflicts were adjacent additions: import lines and appended definitions.
• 300 random two-layout cases (uneven TP sizes, any subset of surviving ranks, 1 to 6 partitions):
  • holders and map-level exactness hold in every case;
  • cuts along rows: 137 of 137 planned and exact;
  • cuts along columns: 79 planned and exact, and 84 refused by the current planner.
• 300 random unions of three or four layouts, merged in random order: all planned and exact.

On GPU, also before the rebase: 4× A100-SXM4-40GB (NVLink), driver 595.71.05, torch 2.12.1+cu130, NCCL 2.29.7, Python 3.12.11, with pytest --forked over six of the affine test files (all but test_autoep_affine_universal.py):

• master (8a6b53f): 317 passed, 1 skipped;
• this branch: 325 passed, 1 skipped, so the 8 new tests also pass over NCCL, including the executor test.
• The one skip is the external-group test that runs only on the CPU accelerator.

cc @0z5a @LiRunGuo @jinyouzhi

[chatgpt-codex-connector]

Codex Review Summary: ✅ Completed

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-12T20:08:36.647260Z	110ff4b	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.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 110ff4b8b3

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

[chatgpt-codex-connector]

Reject blocks that cross logical-axis bounds

When an otherwise executable piece wraps across a logical-axis boundary (for example shape=(1, 4), source_offset=0, and source_strides=(3, 1) for a (2, 3) parameter), matching row-major strides and a valid starting coordinate make this branch classify it as axis-aligned even though origin + shape exceeds the logical shape. _cut then models offset 3 as nonexistent coordinate (0, 3) rather than (1, 0), and the returned union advertises different locations for the two real copies; both validate_coverage() and plan_transfer() reject it. Check that LogicalBox(origin, piece.shape) is contained in the logical parameter before accepting it.

Useful? React with 👍 / 👎.

[Achyuthan-S]

Thanks, confirmed: for a (2, 3) parameter the (1, 4) run from offset 0 was cut as the box (0, 0)+(1, 4), and the union came back with a piece that validate_coverage and plan_transfer then rejected. _block_of now refuses a block that doesn't fit inside the parameter, the same way it refuses other pieces that aren't blocks, and test_a_union_refuses_layouts_it_cannot_merge covers it. Fixed in 45652c6.

[chatgpt-codex-connector]

Refine replica boundaries within each input map

When one valid input has replicas with different piece boundaries—for example location 0 stores one (2, 3) box while location 1 stores two (1, 3) boxes in swapped shard rows, all declaring {0, 1}—_blocks_held contains a full box plus its halves. This loop appends overlaps against the original block rather than a common disjoint refinement, so unioning a third full-copy layout gives that location the full box and both halves; affine_map.validate() then raises because 12 elements are assigned to a 6-element shard. Include boundaries within each input map when building the disjoint refinement before emitting pieces.

Useful? React with 👍 / 👎.

[Achyuthan-S]

Thanks, confirmed: that input made the union raise "Rank 10 holds a shard of 6 elements but its pieces account for 12". Every piece is now cut into the disjoint refinement of all the blocks both layouts hold, including one layout's own replica edges, so each part appears once and lists every holder. test_a_union_cuts_replicas_that_split_a_block_differently covers it. On 900 random layouts the holders match the previous commit for every element, and every plan is exact or refused exactly as before. Fixed in 45652c6.