DeepSpeed

7e58feb6 - Check that converting through an affine map writes every element once (#8975)

Check that converting through an affine map writes every element once (#8975) 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. #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 --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>