Repository navigation
[UC] Plan a scale-1 affine shard-to-shard transfer without rebuilding the parameter - #8623
Conversation
|
@0z5a Both numbers are reachable — §8.2's 7 is not conditional on choosing the replica per target. I recomposed it from If you are getting 8, I would guess the merge stops at the query/kv boundary. That rule is On the policy: agreed it should be a parameter. Lowest holder is not topology-neutral either |
Thanks, that clears up the distinction. The current merge is indeed stopping at the source-piece boundary, so the TP2 -> TP4 BigCode case stays at 8. That is carrying the map-level homogeneity constraint from §5 into the resolved transfer plan, where locations is no longer part of the descriptor. The update will separate those two rules:
With the existing lowest-holder policy, BigCode TP2 -> TP4 should therefore regress to 7 segments. I’ll also make replica-holder selection an explicit planner policy rather than baking lowest-holder into the geometry. A per-target policy can then exercise the 6-segment case separately, without treating that as the canonical/default plan. This also seems worth clarifying next to the composer rules in §8.2: P5 constrains the IR pieces, not already-resolved transfer segments. |
Conversion and restore both go through the full logical tensor: rebuild it from the shards, then extract the target's. When the two topologies are known, that round trip is unnecessary -- the two descriptions compose, and the copies they imply go straight from one shard to another. plan_transfer composes a source map and a target map into strided copies between named ranks of each, and refuses anything it cannot express rather than approximating it. A piece is accepted only when both of its stride vectors are the row-major strides of the tensor it addresses, which is what every builder in affine.py already emits and what makes each piece a box; a permutation, a scale, or a shard described with the wrong number of axes is refused by name. Two ranks are two namespaces, so a source rank and a target rank that share a number are not the same process, and the plan keeps both. Coverage is argued over pieces, never over elements, so nothing here is sized by the parameter: validate_coverage walks every element by design and is not called. Each rank's pieces must partition its shard by address -- being disjoint in the parameter does not make two pieces disjoint in the shard, and their element counts agree either way -- and each target region must be tiled by the intersections that reach it. Volume totals are checked last for exactly that reason: one gap and one overlap of equal size cancel in any sum, including a sum over the whole plan. Nothing moves. There is no transport here, no scale other than 1, no judgment about whether two parameters are the same parameter, and no claim of a shortest plan. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
f4e1e16 to
b8f0f3b
Compare
|
Verified against the new head. I rebuilt the maps with segmented_map at TP2 and TP4 rather than using your fixture, planned with plan_transfer, and executed the segments element for element against extract: lowest holder gives 7, a per-target selector gives 6, both write 768 with nothing uncovered. Your committed suite also passes here — 97, CPU. The spec text is right, and separating "resolved copies are not another homogeneous map" from P5 is the part I'd have got wrong if I'd written it alone. A folded copy having no single logical origin is the consequence I hadn't followed through. One thing to tighten: "a specified per-target holder choice yields 6" depends on which per-target policy. A selector that takes the highest holder still gives 7 — I tried it. Only a selector that picks the rank holding the target's own private rows reaches 6, because that is what puts the query read adjacent to the kv read in the source shard. Worth saying, since The selector hook is the right shape. It is the seam a cost model plugs into, which is the part I took on #8252. |
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
@Achyuthan-S Thanks, good catch. §8.2 and both draft PR descriptions now specify that the 6-segment fold requires choosing the holder of each target's private Q rows; choosing the highest KV holder still gives 7. I added planner and two-GPU execution regressions for 7/6/7. B1 passes 98 CPU tests, and B2 passes 12 cases per rank. |
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The planner can accept incomplete or invalid maps and reject some valid replica layouts.
Review effort: Balanced
Findings: 5
Open (8)
What changed in this PR
This PR adds a metadata-only planner for copying a parameter directly between affine shard layouts, supporting the broader goal of transfer without rebuilding the full parameter.
Changes:
- Adds replica selection and folding of adjacent shard-to-shard copies.
- Adds transfer-plan regression tests and clarifies the distinction between map pieces and resolved copies.
| File | Description |
|---|---|
tests/unit/checkpoint/test_affine_transfer_plan.py |
Tests copy addresses, replica choices, refusals, and coverage. |
deepspeed/checkpoint/affine_transfer.py |
Implements transfer planning and copy descriptors. |
deepspeed/checkpoint/affine_ir_spec.md |
Clarifies when resolved copies may be folded. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
|
||
| logical_strides = _row_major_strides(logical_shape) | ||
| shard_strides = _row_major_strides(shard_shape) | ||
| for axis, (got, want, size) in enumerate(zip(piece.source_strides, logical_strides, piece.shape)): |
| f'{side}: a shard is described with a different number of axes than the logical parameter, so ' | ||
| 'no axis correspondence exists to plan against') | ||
|
|
||
| frames = [] |
| if len(frames) > limits.max_pieces_per_map: | ||
| raise AffineTransferError(f'{side}: {len(frames)} pieces exceeds max_pieces_per_map ' | ||
| f'{limits.max_pieces_per_map}') | ||
| _shards_tile(frames, side, limits) |
| if len({frozenset(holder.locations) for holder in holders}) > 1: | ||
| raise InvalidMapError(f'source: {box!r} is claimed by ranks {[h.rank for h in holders]} ' | ||
| 'with different locations') |
| f'{budget.max_pair_checks}') | ||
|
|
||
| pending: List[Tuple[_Frame, LogicalBox, Tuple[_Frame, ...]]] = [] | ||
| for frame in sorted(targets, key=lambda frame: (frame.rank, frame.box.origin, frame.box.extent)): |
| if not box.intersection(other).is_empty: | ||
| raise InvalidMapError(f'source: {box!r} and {other!r} overlap without agreeing, so part of ' | ||
| 'the parameter has two owners') |
| plan_bytes = _plan_bytes(segments) | ||
| if plan_bytes > budget.max_plan_bytes: |
| for segment in merged: | ||
| for index, held in enumerate(remaining): | ||
| candidate = _merge(held, segment) |
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Hi @0z5a can you take a look at the CI failure? Thanks! |
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
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>
Summary
Compose source and target affine maps into shard-to-shard copy descriptors without rebuilding the logical parameter. This PR remains metadata only and based on
master.Resolved copies now fold when both physical shard address mappings continue along an axis. Their descriptor no longer claims a single
logical_origin: a copy can span two distinct logical regions. The planner keeps every validated replica candidate until a target intersection is resolved. It uses the lowest holder by default and accepts a caller-suppliedreplica_selectorfor other choices. Invalid selections fail explicitly.The spec distinguishes homogeneous input map pieces from resolved copies. The regression suite is committed in
tests/unit/checkpoint/test_affine_transfer_plan.py.Review result
The 7/6/7 counts apply to this layout and these three holder policies. For targets 0 and 1, the private Q rows are on source 0; for targets 2 and 3, they are on source 1. Selecting that same source for the replicated KV block makes the query and KV reads adjacent, allowing the sixth segment. Selecting the highest-ranked KV holder for every target does not. The previous 8-segment plan was valid; this change compresses its copy schedule. The selector does not change the input map's
locationsor relax its homogeneity rule.On the 8× RTX 5090 host, Python 3.12.13 and torch 2.13.0+cu130, the initial B1 CPU suite passed: 98 passed (
PYTHONPATH=$PWD CUDA_VISIBLE_DEVICES="" DS_ACCELERATOR=cpu python -m pytest tests/unit/checkpoint/test_affine_transfer_plan.py -q). Tests cover exact address oracles, nonzero storage offsets and guard regions, replica offsets, invalid selectors, serialization, budgets, determinism and the no-materialization contract.Follow-up at
c2ac3a4: the targeted suite passed 119 tests in the existing 0z5a Python environment using task-local PR modules. New cases cover stride rank, raw-piece and merge budgets, missing shard pieces, declared holders, complete target coverage, replica boundary mismatch, and serialized plan size. No environment packages were changed.The stacked B2 executor PR has two-GPU copy and model E2E results, including the direct-versus-rebuild speed table. B1 itself does not move data, so it makes no transport speed claim.
Scope
Scale 1 and matching logical shapes only. No checkpoint format change, topology search, CUDA graph, online weight update or optimizer coordinate transform. Refs #8230 and #8252.
CI follow-up (October 10)
The failed
cpu-torch-latest / unit testsrun 36671037164 had three rollout test failures because its Transformers main revision requires the configuration'sper_layer_configcontract. The three test models now use the realPretrainedConfigrather than a partialSimpleNamespacesubclass.Merged current master
967803b0fe1a2b4a2738fac676ded4a85d7ce491and resolved conflicts with its concurrent cache-config fix, retaining all upstream rollout tests and behavior. The affine transfer planner and its regressions remain intact.Local CPU validation on macOS arm64, Python 3.12.13 and PyTorch 2.14.1:
5c9f079ca0651631cd8c9c10ebf8d6ac2d760ea1: reproduced all three original failures before the fix. After the master merge, both affected rollout test files plus the affine transfer-plan suite report 157 passed, 1 skipped. The skip requires native CUDA shared-prefill support.The refreshed PR is mergeable, and DCO and test collection have passed. The full Linux CPU workflow is currently
action_requiredwith no jobs started (run 38042624340); maintainer approval is needed before execution. These local results do not claim a full CI or GPU run, or Transformers 4.x rollout compatibility.CI fixture fix:
dc765cac01e67a14cc501f00b9d6fef52b796533; conflict refresh:a2e60af442dd2207e92a8b6a0f90fc6d6f7af9a4.