Describe the GPTBigCode and Yuan shared-QK layouts - #8575
Achyuthan-S wants to merge 3 commits into
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
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/. |
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