Check where column-split pieces write in a shard (#8922)
`ParamAffineMap.validate()` checks where a piece writes in its shard
only when the piece's destination is one dense interval. That check came
from #8624. A column split has no such interval: each piece writes a few
columns of every shard row. So for fused QKV split along columns, Yuan's
o_proj, or any `sub_param_map` / `segmented_map` / `block_gather_map`
with `partition_dim=1`, two pieces can write the same shard columns and
pass every check.
Reproduction on `master` (fused QKV `(4, 12)` split along columns at TP
2):
1. Move rank 0's value piece onto its key columns.
2. `validate()` and `validate_coverage()` both pass.
3. `rebuild(extract(F))` differs from `F` in 8 of 48 elements.
4. Two columns of rank 0's shard are never written by `extract`, so they
hold uninitialised memory.
## Changes
- **`validate()`** now also checks strided pieces, still from metadata
only, without walking elements.
- Bounds: each strided piece's lowest and highest shard address must lie
inside the shard.
- Overlap: a column-split piece is a box in its shard's row-major grid,
so overlap between such pieces follows from their corners. A strided
view that wraps past the end of a shard row is valid but is not a box;
the full check below covers it.
- The box pass runs only on ranks that hold a strided piece. Ranks with
only dense pieces cost the same as before.
- **`validate_coverage()`**, the full O(n) check, now walks every shard
address and requires each one to be written exactly once. Together with
the element-count check, that makes "every shard address is written by
exactly one element" a complete check, for any stride, including strided
pieces that are not boxes.
## Tests
`tests/unit/checkpoint/test_affine_shard_map.py`:
| Test | Contract | Fails on `master` |
|---|---|---|
| `test_strided_destination_overlap_is_refused` | column-split pieces
that collide are refused | yes (did not raise) |
| `test_strided_destination_outside_the_shard_is_refused` | a
column-split piece past the shard's width is refused | yes (did not
raise) |
| `test_full_check_refuses_a_destination_collision_of_any_stride` |
`validate_coverage()` refuses a collision that no box can describe | yes
(did not raise) |
| `test_strided_destination_validation_does_not_enumerate_elements` |
`validate()` stays metadata-only for a `10^9 × 8` column split | no; it
guards cost |
| `test_strided_destination_that_wraps_a_shard_row_round_trips` | a
strided view that wraps a shard row but stays inside the shard is
accepted and round-trips | no; it guards against refusing a valid map |
## 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`: **76 passed**.
- `test_affine_shard_map.py`, `test_affine_transfer_plan.py`,
`test_autoep_affine.py`, `test_autoep_affine_universal.py` and
`test_autotp_uc_checkpoint.py`: **276 passed, 5 failed**. The failures
are identical on `master`. The 5 tests build the CPU Adam op, which
fails on macOS arm64 (`cpuinfo` reports no `flags`). They were not run
here.
- In those suites the new box pass checks 28 strided pieces from the
existing tests' layouts and rejects none of them.
- End to end, to confirm the converter path is unaffected: AutoTP
training, save, `ds_to_universal`, then a universal load at another TP
degree. Runs: TP 2→4 under ZeRO 1 and 2, including an uneven split at TP
4; TP 4→2 under ZeRO 1. Parameters, `exp_avg` and `exp_avg_sq` are
restored bit for bit, and the resumed losses match the uninterrupted run
to within 1.2e-7.
- `yapf` 0.43.0 and `flake8` are clean on both files.
Refs #8252, #8624.
@delock
---------
Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
