Skip to content

Describe the GPTBigCode and Yuan shared-QK layouts - #8575

Merged
delock merged 4 commits into
deepspeedai:masterfrom
Achyuthan-S:step4b-fused-yuan
Sep 20, 2026
Merged

delock merged 4 commits into
deepspeedai:masterfrom
Achyuthan-S:step4b-fused-yuan

Conversation

@Achyuthan-S

@Achyuthan-S Achyuthan-S commented Sep 17, 2026 •

Copy link
Copy Markdown
Contributor

Three of the four layouts in AUTOTP_UNSUPPORTED_PARAMETER_PATTERNS now 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.

codegentype is 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_map describes 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_map takes 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_map so the replicated segment claims a single owner produces two failures rather than passing quietly.

shared_qk_value_head_ids is extracted from shard_value_with_share_qk and used by both the partition and the map, so the description cannot drift from the layout it describes.

What stays unsupported, deliberately

  • A layout with no head counts. Describing Yuan needs num_kv_heads; without it the layer publishes the refusal rather than guessing, and there is a test for that.
  • The Yuan o_proj bias. Its partition pre-divides by the world size and replicates, which needs a scaled piece — and that code path does not currently produce a usable bias, so there is no working layout to describe.

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

Copilot AI lite review requested due to automatic review settings September 17, 2026 15:25

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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.

Comment on lines +227 to +235
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])
Comment thread deepspeed/module_inject/layers.py Outdated
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)
@delock
delock self-requested a review September 18, 2026 01:05

@delock delock left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread deepspeed/module_inject/layers.py Outdated
_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):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread deepspeed/checkpoint/affine.py Outdated
scale=entry.get('scale', 1.0))


def segmented_map(shape, segments, partition_dim, tp_degree, split_widths=None):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread deepspeed/module_inject/layers.py Outdated
_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):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@delock

delock commented Sep 18, 2026

Copy link
Copy Markdown
Collaborator

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>
@Achyuthan-S

Achyuthan-S commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor Author

@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/.

@delock
delock enabled auto-merge September 20, 2026 02:50
@delock
delock added this pull request to the merge queue Sep 20, 2026
Merged via the queue into deepspeedai:master with commit b2370ad Sep 20, 2026
13 checks passed
pull Bot pushed a commit to AmirulAndalib/DeepSpeed that referenced this pull request Sep 29, 2026
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>
banxingmjj pushed a commit to openanolis/DeepSpeed that referenced this pull request Oct 16, 2026
…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>
banxingmjj pushed a commit to openanolis/DeepSpeed that referenced this pull request Oct 21, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants