[Achyuthan-S]

Conversion trusts a map further than it checks it.

• rebuild (what ds_to_universal calls on convert) runs only validate(). That checks each rank's element count and where pieces write in the shard, from metadata, but it never looks at the source side. So a map can pass it while two ranks claim the same block of the parameter and no rank holds another.
  • rebuild then returns whatever torch.empty left in the elements nobody holds.
  • It also lets one rank's block overwrite the other's with no comparison.
• extract (what the AutoTP restore path calls) runs no check at all.
• A destination that wraps past the end of a shard row is not a box, so validate() cannot place it. Check where column-split pieces write in a shard #8922 made validate_coverage() complete for such destinations, but conversion does not call validate_coverage(), because that check walks elements in Python.

Reproduction on master:

1. Take two ranks with shards of 4 elements each, for an (8,) parameter.
2. Point rank 1's piece at source offset 0 instead of 4.
3. validate() passes. rebuild returns rank 1's values in elements 0–3 and uninitialised memory in 4–7, with no error.

How often this happens:

• I fuzzed the in-tree constructors (contiguous split, replicated, segmented, block gather). Each case is a random 2-D shape, TP 2–4 with uneven widths, and one offset or stride of one piece corrupted.
• On master, 671 of 2,334 corrupted maps converted without an error and produced a wrong tensor. With this change, none do: each one is refused or still round-trips exactly.
• Every map that passes validate_coverage() round-trips exactly on both (801 of 801).
• With the stricter oracle from @0z5a's review (shards built from the uncorrupted map, then rebuilt through the corrupted one; 2,000 cases): master converts 421 silently wrong, this PR none.
• One limit: a corruption that leaves a valid but different layout, for example two equal-shaped pieces swapping destinations, passes every map check, because the map is self-consistent and disagrees only with the memory it describes. Catching that takes a runtime digest of the data, not a check on the map.

Changes

• rebuild
  • Records who wrote each element of the full tensor (nobody, a single holder, or a replica), one byte per element, and refuses:
    • an element no piece writes;
    • an element written by two pieces unless both are replicas, whichever rank is visited first;
    • replicas that disagree on any element. That includes replicas cut at different boundaries on different ranks. Before, only a piece starting exactly where an earlier one started was compared, so the rest of the replica overwrote it unchecked.
  • Also tracks the shard addresses it reads on each rank, so a destination collision is refused for any stride.
• extract
  • Runs rebuild's metadata check for its own rank: validate() restricted to one rank, through a new _validate_rank.
  • Uses the same shard-address mask, so a restored shard never holds uninitialised memory.
• Both, through _validate_rank (and so validate())
  • Refuse a piece whose source or destination view reaches one address twice within itself.
  • Such a view loses a shard value, and which value survives depends on the device.
  • The check is metadata only: taking the axes that move in order of stride, each stride must clear what the axes before it reach. Every in-tree layout satisfies it.
  • A zero stride now raises a ValueError instead of torch's own RuntimeError.
• No change for valid maps. Error types and messages follow the existing ones (ValueError, "twice", "uncovered", "different data"), plus "itself" for the new check.

Cost:

• One byte per element of the full tensor and one per shard element, with vectorised tensor ops and no per-element Python.
• rebuild already materialises the full tensor, so the fp32 states that ds_to_universal converts need at most a quarter more memory, for the duration of one parameter's conversion.

Tests

tests/unit/checkpoint/test_affine_shard_map.py:

Test	Contract	Fails on master
test_conversion_refuses_a_destination_collision_of_any_stride	extract and rebuild refuse the collision from the full-check test of #8922, which validate() accepts	yes (did not raise)
test_replicas_cut_at_different_boundaries_are_compared_in_full	a replica that one rank holds whole and another holds in halves round-trips, and is refused when the second half drifts	yes (did not raise)
test_conversion_refuses_a_map_that_does_not_cover_the_parameter	rebuild refuses a map that leaves an element unheld, and one where two ranks own one block	yes (did not raise)
test_a_replica_overlapping_a_single_holder_is_refused_in_either_order	a block declared both replicated and single-held is refused whichever rank comes first	yes (accepted in both orders)
test_a_view_that_overlaps_itself_is_refused	extract and rebuild refuse a source view that reads one address twice, a zero stride, and a destination view that writes one address twice	yes (silently wrong)
test_a_corrupted_map_is_refused_or_still_exact	fixed seed: 300 random layouts, with shards built before one offset or stride is corrupted, are either refused or rebuilt bit-exact	yes (silently wrong tensor)

The colliding map of #8922's full-check test is now a helper shared by both tests.

Validation

macOS 26.5 (arm64, CPU), Python 3.9.6, torch 2.8.0, DS_ACCELERATOR=cpu LOCAL_SIZE=4.

• test_affine_shard_map.py, test_affine_transfer_plan.py, test_zero_affine.py, test_autoep_affine.py and test_autoep_affine_universal.py: 280 passed.
• test_universal_checkpoint.py, test_autotp_uc_checkpoint.py, test_autoep_checkpoint.py, test_autoep_unit.py and test_autoep_affine_universal.py: 344 passed.
  • I ran master (ba0d63ce2) side by side, and it fails exactly the same 57 tests.
  • 53 of them build the CPU Adam op, which fails on macOS arm64 (cpuinfo reports no flags).
  • The other 4 are test_storage_maps_match_runtime_ownership_and_plan_into_a_new_placement.
• End to end, to confirm no valid map is refused: AutoTP training, save, ds_to_universal, then a universal load at another TP degree, ZeRO 1, gloo.
  • Models: an MLP block, CodeGen (fused QKV), GPTBigCode (multi-query, replicated K/V), a Yuan-style shared-QK attention (replicated o_proj bias), and Llama with a tied embedding and vocab_parallel_lm_head (vocabulary of 67, uneven at every degree).
  • TP changes: 2→4, 4→2 and 1→2.
  • Parameters, exp_avg and exp_avg_sq are restored bit for bit, and the resumed losses match the uninterrupted run to within 4.8e-7.
• On GPU and XPU, @jinyouzhi ran the first commit on 8× NVIDIA L20 (torch 2.14.1+cu130) and on 8× Intel Arc Pro B70.
  • The same affine files: 283 passed on each.
  • On CUDA, rebuilding a 4096×4096 fp32 column split at TP 8 takes 1.2 ms, and 8 extracts take 0.9 ms.
  • The second commit adds only metadata checks and a one-byte mask in place of a bool one. It has not been rerun on GPU.
• yapf 0.43.0 and flake8 are clean on both files.

Refs #8252, #8922.

@delock

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

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

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-11-01T13:15:22.271018Z	cd35be3	PR opened

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

[jinyouzhi]

Thanks, the mask approach is a clean way to make conversion fail closed without per-element Python. Since this was not run on GPU, I ran it on two machines:

• 8x NVIDIA L20 (torch 2.14.1+cu130): test_affine_shard_map / test_affine_transfer_plan / test_zero_affine / test_autoep_affine / test_autoep_affine_universal, 283 passed. On CUDA, rebuild of a 4096x4096 fp32 TP8 column split takes 1.2 ms and 8 extracts take 0.9 ms, so the mask syncs are negligible.
• 8x Intel Arc Pro B70: the same files, 283 passed. extract on a 4096-piece rank takes ~0.04 s, so the added _validate_rank is cheap.

Two gaps against the "every element once" contract:

1. A source view that overlaps itself is still silently wrong. covered only records that an element was written, so a single piece whose source strides alias elements writes one element twice and loses a shard value, while the mask ends up all-True:

m = ParamAffineMap((6,), {0: (4,), 1: (3,)}, {
    0: [AffinePiece((2, 2), 0, (1, 1), 0, (2, 1), [0])],   # source addresses 0,1,1,2
    1: [AffinePiece((3,), 3, (1,), 0, (1,), [1])]})
m.validate()  # passes
m.rebuild({0: torch.tensor([9., 8, 7, 6]), 1: torch.tensor([1., 2, 3])})
# CPU:  [9, 7, 6, 1, 2, 3], shard value 8 is dropped, no error
# CUDA: [9, 8, 6, 1, 2, 3], shard value 7 is dropped instead

Which value survives the aliased write depends on the device, so the same checkpoint converts to different tensors on CPU and GPU.
With stride 0 torch raises its own RuntimeError (internal overlap), but (1, 1) is "too hard" for torch's overlap check and passes. A metadata check that the source view does not overlap itself (sort non-degenerate axes by stride, require each stride >= the span of the inner axes) would close this and give a ValueError for the stride-0 case too; the in-tree builders are row-major, so it should not refuse any valid map.

2. Single-holder vs replica overlap depends on piece order. The single-holder check only looks at the piece being written now, not at who wrote first:

p = lambda loc: AffinePiece((4,), 0, (1,), 0, (1,), loc)
shapes = {0: (4,), 1: (4,)}
ParamAffineMap((4,), shapes, {0: [p([0, 1])], 1: [p([1])]}).rebuild(...)   # refused
ParamAffineMap((4,), shapes, {1: [p([1])], 0: [p([0, 1])]}).rebuild(...)   # accepted

The same map is refused or accepted depending on pieces_by_rank order. Keeping a second mask of elements written by multi-holder pieces, and requiring an overlap to be replica-on-replica (len(piece.locations) > 1 and shared[seen].all()), would make it order-independent.

[0z5a]

Read at cd35be35. Making conversion refuse instead of returning whatever torch.empty left is the right default.

I re-ran the corruption fuzz with a stricter oracle. The test extracts and rebuilds through the same corrupted map, so an error that cancels between the two goes unnoticed: the self-overlapping source view from jinyouzhi's comment round-trips exactly that way. A real checkpoint holds shards written by the correct layout, so I generated the shards from the uncorrupted map and rebuilt them through the corrupted one. Same generator and seed, 2,000 cases:

	refused	exact	silently wrong
master	868	711	421
this PR	1,303	697	0

So the claim holds under that oracle too. Switching the test to it would be cheap, a deepcopy of the map before corrupting, and it would catch cancelling errors if the generator ever produces them.

On jinyouzhi's first point: extract has the same gap. With his map, extract(torch.tensor([10., 11, 12, 13, 14, 15]), 0) returns [10, 11, 11, 12] without an error, and that is the AutoTP restore path. A source-overlap check in _validate_rank would cover rebuild and extract together.

On CPU, test_affine_shard_map.py, test_affine_transfer_plan.py, test_zero_affine.py and test_autoep_affine.py pass at this head (263).

[Achyuthan-S]

Thanks both, all three points are in the new commit:

1. Views that overlap themselves (@jinyouzhi's first point, and @0z5a's note that extract has it too). _validate_rank now refuses a piece whose source or destination view reaches one address twice, which covers rebuild and extract together.
   • It is the stride check you described: taking the axes that move in order of stride, each stride has to clear what the axes before it reach.
   • A zero stride now raises a ValueError instead of torch's RuntimeError.
   • No in-tree map is refused. Every map that passes validate_coverage() in the fuzz still round-trips (801 of 801), and the end-to-end convert-and-resume runs are unchanged.
2. Order dependence. The mask now records who wrote each element (nobody, a single holder, or a replica), still one byte per element, and an overlap is accepted only when both writers are replicas. Both orders of your example are refused.
3. Oracle. The property test now builds the shards from the map before corrupting it, as @0z5a suggested.

New tests:

• test_a_view_that_overlaps_itself_is_refused: your source map, a zero stride, and a destination that writes one address twice.
• test_a_replica_overlapping_a_single_holder_is_refused_in_either_order.

All of the new tests fail on master.

One limit worth stating. A corruption that turns the map into a different valid layout, for example swapping the destinations of two equal-shaped pieces, passes every map check: the map is self-consistent and disagrees only with the memory it describes. That case needs a runtime digest of the data (SyncCheck-style), not a check on the map.

[0z5a]

Re-checked at 4cabdb96, and all three points look right to me. The writer record makes the replica rule independent of visiting order, and the stride check refuses self-overlapping views for both rebuild and extract. It's conservative: an interleaved but non-overlapping view such as shape (3, 2) with strides (2, 3) is refused too, which the docstring already states and which nothing in-tree builds. On CPU, test_affine_shard_map.py, test_affine_transfer_plan.py and test_zero_affine.py pass here (236), and the updated test_affine_shard_map.py run against master fails 10 cases, the new ones among them.

[delock]

Hi @Achyuthan-S thanks for hardening affine map and thanks @jinyouzhi @0z5a for reviewing. This PR also looks good to me.