Repository navigation
Describe the GPTBigCode and Yuan shared-QK layouts - #8575
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
Critical partition and map-width mismatches can cause incorrect or lost data.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds affine checkpoint maps for GPTBigCode and Yuan shared-QK layouts, enabling universal checkpoint conversion for previously unsupported weights.
Changes:
- Adds segmented and block-gather affine map constructors.
- Integrates GPTBigCode/Yuan metadata generation.
- Adds shared head-selection logic and validation tests.
File summaries
| File | Description |
|---|---|
tests/unit/checkpoint/test_affine_shard_map.py |
Tests map correctness and fallback behavior. |
deepspeed/module_inject/layers.py |
Publishes GPTBigCode and Yuan maps. |
deepspeed/module_inject/fusedqkv_utils.py |
Shares Yuan head-selection logic. |
deepspeed/checkpoint/affine.py |
Adds affine map constructors. |
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 2
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| head_per_rank = num_heads // world_size | ||
| q_head_start = rank * head_per_rank | ||
| v_head_ids = [] | ||
| index = 0 | ||
| while index < head_per_rank: | ||
| v_head_ids.append(q_head_start // 2) | ||
| q_head_start += 2 | ||
| index += 2 | ||
| v_head_ids.extend([head + num_heads // 2 for head in v_head_ids]) |
| n_embd = self.tp_meta.n_embd | ||
| if not shape or n_embd is None or n_embd >= shape[0]: | ||
| return None | ||
| return segmented_map(shape, [(n_embd, False), (shape[0] - n_embd, True)], 0, self.tp_world_size) |
5f84b03 to
8a0e033
Compare
delock
left a comment
There was a problem hiding this comment.
Thanks for the PR — the shared-source design (shared_qk_value_head_ids used by both the partition and the map) and the marker-tensor verification methodology are excellent.
One blocking item (missing GPTBigCode wiring test) and three non-blocking suggestions below.
| if self._subparam_shard_widths is not None: | ||
| set_fused_qkv_shard_state(self.fused_module.module, self._subparam_shard_widths, self.tp_index) | ||
|
|
||
| def _segmented_affine_map(self, shape): |
There was a problem hiding this comment.
Missing wiring test for GPTBigCode
test_yuan_weight_is_no_longer_refused_by_conversion verifies the full wiring for Yuan (the layer actually publishes AFFINE_MAP_PARAMS and is removed from AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS), but there is no equivalent for bigcodetype. The geometry tests validate segmented_map itself, but nothing fails if fused_LinearLayer._segmented_affine_map never gets called (e.g. a hook-name typo or get_fused_qkv_type returning something else in that context) — the layout would silently stay refused while all tests stay green.
Could you mirror the Yuan test for a fused_LinearLayer with a GPTBigCodeBlock module? The pattern should transfer directly.
| _unsupported_uc_reason = ("Yuan shared-QK tensor parallelism selects noncontiguous head groups that universal " | ||
| "checkpoint conversion cannot currently describe") | ||
|
|
||
| def _shared_qk_affine_map(self, shape): |
There was a problem hiding this comment.
Duplicated _shared_qk_affine_map in the two Yuan classes (non-blocking)
Yuan_LinearAllreduce and Yuan_LinearLayer carry near-identical copies of _shared_qk_affine_map (guards, completeness check, and its comment included); the only real difference is the partition axis. This is essential rather than incidental duplication — both copies encode the same invariant (the head selection must partition the head set) and will need to change together. The len(shape) <= 0 leftover in the other copy shows how they can drift. A shared helper taking partition_dim would collapse them to one-liners at each call site.
| scale=entry.get('scale', 1.0)) | ||
|
|
||
|
|
||
| def segmented_map(shape, segments, partition_dim, tp_degree, split_widths=None): |
There was a problem hiding this comment.
Consider making split_widths required (non-blocking)
The only production caller (fused_LinearLayer._segmented_affine_map) already passes explicit widths; the split_widths=None fallback is only exercised by one test. Since an inferred split can read the wrong rows while still covering the tensor (the exact bug the second commit fixed), removing the fallback would close that door for future callers. _even_split_sizes has no other callers, so it can be deleted along with the fallback — the one test call site just passes explicit widths.
| _unsupported_uc_reason = ("Yuan shared-QK tensor parallelism selects noncontiguous head groups that universal " | ||
| "checkpoint conversion cannot currently describe") | ||
|
|
||
| def _shared_qk_affine_map(self, shape): |
There was a problem hiding this comment.
Dead condition in Yuan_LinearLayer._shared_qk_affine_map (non-blocking)
len(shape) <= 0 is always false for a tensor shape (ndim >= 1). It looks like a copy-paste leftover from the len(shape) <= 1 guard in Yuan_LinearAllreduce (which protects the shape[1] access). Harmless here since shape[0] is valid for 1-D, but the check as written does nothing — either drop it or restore the intended bound.
|
Hi @Achyuthan-S , thanks for the PR. Being able to handle GPTBigCode and Yuan really show the power of this work. The comments is left above. Thanks! |
These were refused by conversion because a single partition dimension cannot say that one block of a parameter is sharded while the next is held whole, nor that a rank takes head blocks which are not adjacent. Both are describable as pieces. GPTBigCode splits its query rows and replicates the key/value block, so its pieces differ in which ranks hold them. Yuan selects the value heads pairing with a rank's query heads, which fall into two runs. The layers build these from head counts and the tp degree, and the tests require that to agree with what the partition functions actually produce. A layout without the head counts to describe itself stays unsupported rather than guessing. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
The query rows of a GPTBigCode weight are not divided evenly: the widths follow the head count and the grain size, so a map built by dividing reads the wrong rows while still covering the tensor. Segments now carry the widths the partition itself computed, and a width list that does not add up is refused. Yuan's head pairing only partitions the heads when each rank takes an even number of them. Otherwise two ranks hold the same head, and a piece naming a single owner contradicts the one beside it, so publish no map instead. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Add the GPTBigCode equivalent of the Yuan wiring test. The geometry tests cover the constructor, so a hook that never fires would leave the layout refused with everything still green; renaming the hook now fails this test. Collapse the two copies of the shared-QK map builder into one helper taking the partition axis. They encoded the same invariant and had already drifted -- one copy guarded len(shape) <= 0, which is never true. Require the split widths rather than falling back to an even division, since an inferred split can read the wrong rows while still covering the tensor. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
8a0e033 to
f8b9022
Compare
|
@delock All four addressed in f8b9022. The wiring test — added and confirmed it earns its place: renaming _segmented_affine_map makes it fail while the geometry tests stay green, which is the scenario you described. The duplication — collapsed into one module-level helper taking partition_dim. Your point about it being essential rather than incidental was borne out: the len(shape) <= 0 leftover disappeared as a consequence of merging them, rather than needing its own fix. split_widths required — done, and _even_split_sizes deleted with it. Agreed that leaving the fallback open invites exactly the bug the second commit fixed. 90 passed across the affine suite, the producer and coverage tests, the resume matrix and tests/unit/runtime/tensor_parallel/. |
Restore now takes this rank's shard from the map the layer published, which is the same piece list conversion uses run in the opposite direction. That completes the symmetry §7 of `affine_ir_spec.md` describes and that conversion has had since deepspeedai#8385: until now the two sides described the same layout in two different ways, and could disagree. Follows deepspeedai#8519 (layers emit the map) and deepspeedai#8575 (GPTBigCode and Yuan describe themselves). ### This one is not additive The map lives on the parameter, built by this job's layer. It is `M_t` — it describes the topology being restored into, regardless of how the source checkpoint was converted. So the new path is taken for **every** AutoTP restore, not only for checkpoints that were converted through a map. A parameter whose layer published no map falls back to the existing category keys unchanged. That is the correct reading of the spec rather than a shortcut: §6.1 keeps `M_t` out of the file precisely because the restoring job derives it from its own layers. But it does mean the blast radius here is wider than the three PRs before it, so the equivalence is worth checking rather than assuming. I ran `tests/unit/checkpoint/` and `tests/unit/runtime/tensor_parallel/` (385 tests) against this branch and against upstream master in a separate worktree. The failing set is **identical by name**, not merely the same count — 75 failed, 187 passed either way, all of them FusedAdam JIT, fp16-on-CPU, or the pre-existing `st_st_size` typo in `test_convert_checkpoint.py`. ### The tests are built to discriminate `test_restore_prefers_the_map_over_the_category_keys` gives one parameter deliberately contradictory geometry: the category keys describe a plain dim-0 split, so rank 0 would take the first half, while the map hands rank 0 the second half. Which half rank 0 receives says which side of the metadata restore actually consulted. This matters because the old path covers the same layouts correctly. Without a discriminating case the map branch could be deleted and every existing test would stay green. I checked: removing the branch fails that test, while `test_restore_falls_back_when_no_map_was_published` keeps passing. ### Validation 112 passed on CPU/gloo (`DS_ACCELERATOR=cpu LOCAL_SIZE=4`): the affine suite, the resume matrix, the coverage test, `tests/unit/runtime/tensor_parallel/`, and deepspeedai#8544's AutoEP affine suite, which this now sits on top of after the rebase. ### Scale powers Restore resolves a partition once per state file, so each state takes its own power of a piece's scale. The power table lives in affine.py as SCALE_POWER_BY_STATE — ds_to_universal.py had its own copy since deepspeedai#8519, so conversion and restore now read one definition rather than two that could drift. This also reaches the refusal built for scaled optimizer states, which only triggers on a non-parameter power. No layer emits a scaled map today, so it is ahead of the path rather than behind it. Related: deepspeedai#8252, deepspeedai#8230. cc @delock --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
…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>
…ated (deepspeedai#8857) Partitioning Yuan's output projection crashes on a bias. The obvious fix for that would make it silently wrong, so this keeps the bias whole on every rank — and once it is, the bias is an ordinary replicated parameter, which removes the last parameter universal checkpoint conversion refused unconditionally. ### The crash `shard_value_with_share_qk` handles both Yuan projections. On the output-projection path it computes the divided bias into `bias` but returns `sharded_bias`, which that path never fills: ```python else: bias = bias / float(world_size) return torch.nn.Parameter(sharded_weight), torch.nn.Parameter(sharded_bias) # sharded_bias is still [] ``` So any Yuan configuration with a bias on `o_proj` fails at TP ≥ 2 during partitioning, before training starts — `AttributeError: 'list' object has no attribute 'detach'` on torch 2.8.0. It has been there since `8ea995ee1f` (June 2024); no test partitioned that projection with a bias. ### Why the obvious fix is wrong Assigning the divided bias to `sharded_bias` makes the crash go away and the numbers wrong. `Yuan_LinearAllreduce` inherits `LinearAllreduce.forward`, which all-reduces the partial products **first** and adds the bias **once afterwards** — and `LinearAllreduce.gather_params` already says so: the bias is replicated across ranks. A pre-divided bias would come out of the layer as `bias / world_size`. Checked against the unsharded layer at 8 heads TP 2, 8 heads TP 4 and 16 heads TP 4, with the real head selection: | bias treatment | matches the unsharded output | |---|---| | kept whole on every rank (this PR) | exact | | divided by the world size | wrong | | current code | raises | ### What the description becomes With the bias kept whole, `Yuan_LinearAllreduce` publishes it exactly as base `LinearAllreduce` does: `replicated=True`, one piece per rank held by every rank, scale 1. Only the weight follows the shared-QK head selection. That corrects two things I said earlier, in deepspeedai#8706 and in the comment this PR replaces: that this path returned an empty bias — it raises — and that the bias needed a scaled piece. That reasoning came from the divide on the broken line. Once the bias is right, no scale is involved. This was the last unconditional `unsupported_reason` in `layers.py`. The refusals that remain are all conditional on a configuration that cannot be described: no kv head count, a head count the shared-QK pairing cannot divide without two ranks sharing a head, and codegen configurations its own partition asserts against. ### The tests discriminate | corruption | result | |---|---| | original line restored | partition test fails on all 3 configurations | | bias divided by the world size | partition test fails on all 3 configurations | | bias published as unsupported again | description test and the Yuan collection test fail | | original line restored, through the real layer | distributed test raises the original `AttributeError` at TP 2 and 4 | | bias divided by the world size, through the real layer | distributed test fails 33 of 36 output elements | `TestYuanOutputProjectionBias` is the integration test: it builds the injected `Yuan_LinearAllreduce` on a real process group, feeds each rank the activations for its own value heads, and checks the output and both gradients against the unsharded `nn.Linear` at TP 2 and 4 — so it exercises the wrapper's choice of the output-projection branch and the real all-reduce-then-bias order, not a reconstruction of it. `test_yuan_output_projection_keeps_its_bias_whole` is the cheaper unit check over more head counts, using the unsharded layer as its reference. `test_yuan_output_projection_bias_is_described_as_replicated` checks the published map, and that with head counts set the layer has nothing left in `AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS`. `test_collect_records_yuan_layout_as_unsupported` used to require both Yuan biases to be refused; it now separates them — the value projection's bias follows its weight's head selection and is still refused without head counts, while the output projection's is not. ### Validation 250 passed across the affine suite, the resume matrix, the coverage test, the AutoEP affine suite, the transfer-plan suite, `tests/unit/runtime/tensor_parallel/` and `TestYuanOutputProjectionBias`, on CPU/gloo (`DS_ACCELERATOR=cpu LOCAL_SIZE=4`; Apple M4 Pro, macOS 26.5, Python 3.9.6, torch 2.8.0). The distributed test runs as real 2- and 4-process groups — confirmed by making it fail at each world size in turn. Not run on GPU. yapf 0.43.0 and flake8 are clean on the changed files. ### Scope The conv layer at `layers.py` that also divides its bias by the world size is untouched. Whether that is correct depends on where that layer adds the bias relative to its reduction — the case raised in deepspeedai#8252 — and I have not verified it. Related: deepspeedai#8252, deepspeedai#8230, deepspeedai#8575, deepspeedai#8706. cc @delock --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
Three of the four layouts in
AUTOTP_UNSUPPORTED_PARAMETER_PATTERNSnow describe themselves, so universal checkpoint conversion stops refusing them. Follows #8519, which made AutoTP layers emit affine maps for the layouts a generic rule already covered.codegentypeis not included — its reshape-and-interleave needs more than a segment list or a block selection, and it is the next piece of work.What was actually missing
Both refusals came from the same limit: a parameter carries one partition dimension, which cannot say that one block of it is sharded while the next is held whole, nor that a rank takes blocks which are not adjacent.
GPTBigCode splits the query rows across ranks and hands every rank the whole key/value block.
segmented_mapdescribes a parameter as an ordered list of segments, each split or replicated, so the two end up as pieces that differ in which ranks hold them rather than in kind. That is the case the per-piece location set was introduced for in #8385.Yuan shared-QK gives a rank the value heads that pair with its query heads, plus the matching second-half heads. Those land in two runs rather than one span, so
block_gather_maptakes a selection of equal-sized blocks and merges consecutive ids into single pieces.Why the tests are worth reading
A layer builds its map analytically, from head counts and the tp degree. The test harness derives one independently by running the real partition function on a marker tensor and reading back where each slice came from. The tests require the two to agree, so what ships is checked against what the partition code actually does rather than against a reading of it.
I confirmed they fail when the map is wrong: corrupting
segmented_mapso the replicated segment claims a single owner produces two failures rather than passing quietly.shared_qk_value_head_idsis extracted fromshard_value_with_share_qkand used by both the partition and the map, so the description cannot drift from the layout it describes.What stays unsupported, deliberately
num_kv_heads; without it the layer publishes the refusal rather than guessing, and there is a test for that.Validation
Validated on CPU/gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4): 90 passed across the affine suite, the producer and coverage tests, the resume matrix and tests/unit/runtime/tensor_parallel/.
Related: #8252, #8230.
cc @delock