[checkpoint] Refuse a map whose pieces write the same shard address

[0z5a]

Incremental review: Files changed against master.

Summary

• ParamAffineMap.validate() compared element counts, which cannot see where those elements land. A map can pass it while two of one rank's pieces write the same shard address and a third range is never written at all.
• Walk the destination addresses of the pieces that pack densely instead. Bounded by the piece count, not the shard size, and silent about ranks whose pieces stride -- an interval cannot describe those.
• One method changed in affine.py, nine tests added.

The counterexample

Four elements [10, 11, 12, 13] into a (4,) shard, with two pieces both writing the front:

piece A: source[0:2] -> dest[0:2]
piece B: source[2:4] -> dest[0:2]

source coverage      [0, 1, 2, 3]      complete
piece numels         2 + 2 = 4         equals the shard
destination writes   [2, 2, 0, 0]      two addresses twice, two never

Run against the real module at 1eb56d29 (deepspeed/checkpoint/affine.py, sha256 21701805c23abc303028e95049f1a423fad4d50fe27ac3f43f9356229d218892), both validate() and validate_coverage() accept this map. validate_coverage is the thorough one, and it is not enough either: it checks that every element of the parameter is covered by someone and that each piece's holders are honest. Neither question is about the shard's own address space, which is what a reader or a writer actually indexes.

Downstream that means rebuild returns a tensor whose tail came from nowhere, and restore or a direct transfer reads a shard as though it were complete.

What the guard does

For each rank, every piece whose destination strides are the row-major strides of its own shape occupies one contiguous run of the shard. Those runs must lie inside the shard and must not overlap.

It stops there deliberately. Where some piece of a rank strides through the shard -- a column split, Yuan's o_proj -- a gap between the runs is not yet a gap in the shard, so the check says nothing rather than guessing. And where every piece is dense, no separate hole check is needed: runs that are in bounds, non-overlapping and sum to the shard size cannot leave a gap. That is why there is no hole assertion here.

Validation

	test_affine_shard_map.py at this PR's head	with the guard
pre-existing tests	60 passed	62 passed
test_double_written_shard_is_refused	fails (map accepted)	passes
test_destination_outside_the_shard_is_refused	fails (map accepted)	passes
6 builders asserted still valid (replicated, uneven row, column, sub-params, GPTBigCode segments, Yuan gather)	n/a	pass
test_validate_never_walks_elements (per-element helpers bound to raisers, then a 10^9 x 8 map validated)	n/a	passes

The two refusals failing at base is the point: the same commit that adds them fails without the change, so the guard is what makes them pass. Run on CPU (Python 3.12.13, torch 2.13.0+cu130). The wider tests/unit/checkpoint directory was run on a pristine 1eb56d29 worktree and on this branch; failing and erroring test names are identical, 241 either way, none introduced. Two environmental causes on both trees: multi-rank cases cannot run under CUDA_VISIBLE_DEVICES="", and this environment ships pytest 9.1.0 while dev requirements pin <8.4.0, so the repository's distributed harness never registers its fixtures.

One behavior change to be aware of

validate() raised AssertionError on a numel mismatch. It now raises ValueError, because a map that overwrites part of its own shard is not a condition worth tying to an interpreter flag that python -O clears. Every call site is inside the checkpoint path, where a silent pass is the failure mode being fixed. No test in the repository relied on the assertion type.

Scope

• No general overlap decision for arbitrary strided views. That needs interval algebra over each rank's address set, and the layouts in tree do not need it yet.
• Not a validation framework, and not a dependency for anything else: affine_transfer.py in [UC] Plan a scale-1 affine shard-to-shard transfer without rebuilding the parameter #8623 proves coverage over boxes on its own.
• No change to the restore path's fallback policy -- what a loader should do when a map is present but invalid is a separate question with its own call sites, and mixing it in here would hide which of the two a failure came from.
• No GPU run: the guard is pure Python over integers.

Refs #8230, #8252

Test and benchmark sources are retained locally; the validation results below refer to those local files.

[Achyuthan-S]

Confirmed against master, and the downstream consequence is sharper than "a tail that came from nowhere" — it is wrong values, not absent ones.

I built your counterexample: validate() and validate_coverage() both accept it, extract returns [12, 13, 0, 0] (piece B overwrites A; the tail is never written), and rebuild returns [12, 13, 12, 13] from an original of [10, 11, 12, 13]. So a map like this round-trips a checkpoint to different values with every check green. Your sha256 for affine.py at 1eb56d2
matches here, so we are looking at the same file.

Your reading of why validate_coverage does not catch it is right: source coverage is complete, and the collision is in the shard's address space, which it never looks at.

The pigeonhole argument in _validate_destinations holds — intervals in bounds, pairwise disjoint, lengths summing to the shard size must tile it, so a dense rank needs no separate hole check. Stopping at strided pieces rather than guessing is the right call.

On AssertionError -> ValueError: agreed, and it makes the module consistent rather than just safer. Every other guard in affine.py already raises ValueError; the assert was the outlier.

One gap, and it is distinct from the fallback-policy question you scoped out. extract() never calls validate() at all — only rebuild() does. So on the restore path the guard does not fire, invalid or not: on your patched tree, extract still returns [12, 13, 0, 0] without complaint. That is not "what should a loader do when a map is invalid" but "the check never runs there." validate() is O(pieces) rather than O(elements), so extract could afford it. Happy to add that side in #8622 if you would rather keep this PR to validate().

Last thing, on reproducibility rather than correctness: with the tests taken back out in the second commit, the validation table cannot be checked from the diff. I get 53 pre-existing tests in test_affine_shard_map.py at base where you report 60 — the file is byte-identical to 1eb56d2 here and nothing is skipped, so it is probably pytest 9.1.0 versus 7.4.3, but it is worth reconciling since the 62 is the number that shows the guard is load-bearing. Committing them the way you did on #8623 would settle it.