Repository navigation
[checkpoint] Refuse a map whose pieces write the same shard address - #8624
Conversation
validate() compared per-rank element counts, so a map can pass it while two of a rank's pieces write the same destination and another range is never written at all. Rebuild, restore and any direct transfer all read such a shard as if it were complete. Walk the destination addresses of the pieces that pack densely instead, which bounds the check by the piece count rather than the shard size, and leave a rank holding strided pieces alone, since intervals cannot describe where they land. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
5624904 to
8fb5e28
Compare
|
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 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. |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Singleton-axis layouts can bypass overlap detection, and the described regression tests are not committed.
Review effort: Balanced
Findings: 1
Open (3)
What changed in this PR
Adds destination-overlap validation for affine checkpoint shard maps.
Changes:
- Detects overlapping or out-of-bounds dense destination intervals.
- Replaces assertion-based size validation with
ValueError.
| File | Description |
|---|---|
deepspeed/checkpoint/affine.py |
Adds dense destination interval validation. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
Confirmed, and the reachability is worth stating precisely. A (1, 2) piece with dest strides (4, 1) occupies two adjacent addresses, but _dest_interval It is not reachable from the constructors. I scanned every piece replicated_map, Your suggested fix holds. Dropping size-1 axes before comparing catches the case, still reports |
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Hi @0z5a, thank you for your PR, I have left my comments, can you take a look? Thanks! |
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
…ai#8706) CodeGen was the last layout in `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS` whose geometry had no description. Its weight now publishes an affine map, so conversion stops refusing it. Follows deepspeedai#8519 (layers emit the map), deepspeedai#8575 (GPTBigCode and Yuan describe themselves) and deepspeedai#8622 (restore reads it). ### What the layout actually is `_codegen_type_transpose` views the weight as `codegen_mp_num` blocks, cuts each into q/k/v, subdivides every third across the ranks, then reassembles the pieces q, k, v at a time. The result is **one fixed permutation of the whole weight** — rank-independent — from which each rank takes a contiguous slice. So the outer split is a plain even row split, and the interleaving lives entirely inside each rank's region, where the rows are a permutation of equal-sized blocks. That is why one `partition_dim` cannot say it and a block selection can. `block_gather_map` already expresses this, so there is **no new IR primitive here**. I said when deepspeedai#8575 landed that codegen "needs more than a segment list or a block selection"; that was wrong, and this is the correction. The work is deriving the block order, not extending the language. ### How the description is kept honest `codegen_block_layout` derives the order independently, because the partition reaches it through reshapes and never names a block. That departs from deepspeedai#8575, where `shared_qk_value_head_ids` is shared between the partition and its description. Sharing here would mean refactoring `_codegen_type_transpose` itself, which is riskier than the description it would protect, so the drift guard is instead a test requiring the derived map to equal one **probed from the real partition function**. That checks behaviour rather than structure. @delock — if you would rather have the refactor so there is one definition, say so and I will do it that way. The helper refuses every case the partition does not produce: no head count, a head count it cannot divide (the same condition `_codegen_type_transpose` asserts), fewer than two ranks, and rows that do not divide into equal blocks. An earlier version of this change published a map for `AutoTPMeta()` with no head count — a config whose own assert cannot be evaluated — and `test_collect_records_fused_layouts_without_a_sub_param_split` caught it. ### The tests discriminate | corruption | result | |---|---| | q/k/v loop order swapped with the block loop | 6 fail | | block size missing its `world_size` factor | 6 fail | | hook renamed so it never fires | 1 fail — the wiring test, which is its only job | The piece-count test is the one that matters beyond correctness. A compressor that missed the permutation would emit one piece per row, still cover the tensor and still round-trip, so only the count reveals it. Across rows 24 → 192 at TP 2 the count holds at 11. ### Validation The derived map equals the probed map for 24x8 TP2, 48x16 TP2, 96x32 TP2, 96x8 TP4 and 192x8 TP8. In every case `extract` is bit-exact against the real partition, `rebuild` round-trips, the map covers `F`, and both `validate()` and `validate_coverage()` pass. 11 pieces per rank at TP 2, 12 at TP 4 and TP 8 — matching the count in §5 of `affine_ir_spec.md`. 126 passed on this rebase across the affine suite, the resume matrix, the coverage test, `tests/unit/runtime/tensor_parallel/` and the AutoEP affine suite, on CPU/gloo (`DS_ACCELERATOR=cpu LOCAL_SIZE=4`) — now on top of deepspeedai#8624's stricter `validate()`, which every codegen piece passes with its destination checked. With deepspeedai#8623 merged, codegen maps also compose through `plan_transfer` exactly: TP 2↔4, 2↔8 and 4↔8, with source shards taken from the real partition function and every target shard compared element for element. deepspeedai#8623's own suite passes here too (119). Earlier, at base `581ac41bd`, I ran `tests/unit/checkpoint/` and `tests/unit/runtime/tensor_parallel/` on this branch and on a pristine worktree of that commit: 244 passed here against 236 on master — the 8 this PR adds — with 74 failed and 106 errors on both, and the failing set identical **by name** rather than merely in count, 180 names either way with none on one side only. All of them are environmental: 126 FusedAdam JIT (no ninja), 54 fp16-on-CPU, 4 cpu_adam JIT. Master has moved since, so that pairing is against `581ac41bd` rather than today's tip; I can redo it if the numbers matter for review. ### What this does not close Yuan's o_proj bias is still refused, so the unsupported set is not empty yet. That one is not a description gap: its partition pre-divides by the world size and that path returns an empty bias rather than a usable one, so there is no layout to describe until the partition is fixed. Separate change. Related: deepspeedai#8252, deepspeedai#8230. cc @delock Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
ParamAffineMap.validate()previously checked element counts without checking where dense pieces write. Two two-element pieces both writingdest[0:2]satisfy the count for a four-element shard while overwriting the head and leaving the tail unwritten.Check dense destination intervals for bounds and overlap using piece metadata. Ignore strides on singleton axes when identifying contiguous footprints, so
(1, 2)pieces with destination strides(4, 1)receive the same validation. Genuinely strided layouts retain their existing checks. Invalid shard element counts raiseValueErrorconsistently with other structural errors.The code docstrings describe the validation contract; the motivation and counterexample are kept here, as requested in review.
Validation:
tests/unit/checkpoint/test_affine_shard_map.py: 59 passed, using the real repository modules on macOS CPU, Python 3.12.13, PyTorch 2.14.1.10^9 x 8map.deepspeed/checkpoint/affine.pyandtests/unit/checkpoint/test_affine_shard_map.py; repeated for the revised docstrings.Full upstream CPU and formatting workflows require maintainer approval (
action_required, no jobs executed). The passing local checks do not establish a full CI pass.Refs #8230, #8252.