Carry the affine scale on the replicated map, not the split - #8477
Merged
delock merged 2 commits intoSep 11, 2026
Merged
Conversation
The layouts that pre-divide a value hold it whole on every rank: a row-parallel layer replicates its bias divided by the world size so the all-reduced sum adds it once. The weight beside it is split and unscaled, so no in-tree layout scales a split and the split constructors no longer take the argument. Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
Achyuthan-S
requested review from
loadams,
tjruwase and
tohtana
as code owners
September 10, 2026 17:40
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1fc72e8d33
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
|
|
||
| def replicated_map(shape, tp_degree): | ||
| def replicated_map(shape, tp_degree, scale=1.0): |
There was a problem hiding this comment.
This non-merge commit has no Signed-off-by trailer (git show -s --format='%(trailers:key=Signed-off-by,valueonly)' returns nothing), so it violates the repository's mandatory commit policy; recreate the commit with --signoff before merging.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
🟢 Approval recommended
The change is small, consistent with the stated layout semantics, and includes a focused unit test validating the new replicated-scale behavior.
Pull request overview
This PR adjusts the affine-map constructors so that elementwise scaling is associated with replicated layouts (where every rank holds the whole value) rather than with split layouts, matching how DeepSpeed’s TP partitioning treats pre-divided replicated biases. It also adds a unit test to lock in this behavior.
Changes:
- Extend
replicated_map(...)to accept ascaleparameter and carry that scale on its single replicatedAffinePiece. - Remove the
scaleparameter fromcontiguous_split_map(...), reflecting that split weights are unscaled in the supported in-tree layouts. - Add a unit test verifying that
replicated_mapproperly extracts scaled shards and rebuilds the unscaled full tensor.
File summaries
| File | Description |
|---|---|
deepspeed/checkpoint/affine.py |
Moves scale support to replicated_map and removes scale from contiguous_split_map to align constructor semantics with actual TP layout behavior. |
tests/unit/checkpoint/test_affine_shard_map.py |
Adds coverage ensuring replicated maps carry scaling and round-trip correctly. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 0
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
delock
approved these changes
Sep 11, 2026
delock
enabled auto-merge
September 11, 2026 00:11
delock
disabled auto-merge
September 11, 2026 00:12
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Sep 11, 2026
Collaborator
|
@Achyuthan-S can you fix formatting error? Thanks! |
Contributor
Author
|
@delock I don't think there's a formatting issue to fix — yapf crashed on the runner rather than reporting a diff: That's yapf's lib2to3 grammar cache being truncated in the pre-commit env. There's no I re-checked locally with the pinned v0.40.0 against the branch including your merge of master — no diff on either file, and those are the only two Happy to push an empty commit if that's easier than a re-queue, but I don't believe there's anything to correct. |
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Sep 11, 2026
pull Bot
pushed a commit
to davidsolomon21cn/DeepSpeed
that referenced
this pull request
Sep 16, 2026
Step 4 of the staging plan in deepspeedai#8252: AutoTP layers now emit an affine map, and the converter reads it on a path a real job takes rather than one a test supplies. Follows deepspeedai#8385 (the IR, the lowering constructors, the converter branch) and deepspeedai#8477 (the scale fix). Scope here is contiguous splits and replicated parameters; the fused QKV and Yuan shared-QK layouts are next and still carry `unsupported_reason`, so they emit no map and are unaffected. ### Where a map is built, and why it is not where you would expect The producer cannot build a map from what `collect_autotp_universal_checkpoint_info` sees. The conversion metadata carries eight fields and per-rank partition sizes are not among them — those were deliberately restore-only. But the layer does have them: `_freeze_partition_sizes` resolves them via `get_shard_size_list` while the layer is built, and they are not recoverable later from a shape alone. So the map is derived at mark time, in `_set_param_uc_meta`, which already receives `partition_sizes`, `sub_param_shard_widths`, `logical_shape` and `replicated`, and has `tp_world_size` in scope. That puts the derivation in one place rather than in each of the seven `_mark_uc_metadata` implementations, and `collect_...` then gathers what the layers produced. A layout that is not describable yet returns `None`, so conversion falls back to the categories. ### The map and the categories have to coexist A checkpoint carrying an affine map also carries the existing pattern lists, so a converter predating the map still reads it — that is the additivity §6.3 promises. The consequence is that the converter which prefers the map never consults those branches, which leaves their patterns looking unused and fails a strict conversion. They are superseded, not unused, so the affine branch now marks them consumed. This is the compatibility question raised as "one boundary for review" in deepspeedai#8385; it turns out both have to be present and the reader has to account for the other. This only surfaced end to end. The producer's output was correct in isolation, and the emitted maps matched a layout already verified by a passing resume — the failure was a strict-mode assert on rank 0 during conversion, which presented as rank 1 blocking on the following barrier. ### Tests `TestAffineMapProducer` builds a real AutoTP engine at TP2 and requires the producer to emit exactly the four maps the resume fixture previously supplied by hand. Those maps are not a guess: a full train → save → convert → resume cycle reproduces uninterrupted training through them, so matching them is evidence rather than self-consistency. `affine_resume_checkpoint` now takes its layout from the producer instead of injecting one, so the four save-convert-resume cases exercise the metadata a real job writes. All four pass — TP2 → TP1 and TP2 → TP2, through both the legacy and affine paths, compared against uninterrupted training on logits, losses, gradients, FP32 weights, both Adam moments and step counters. Validated on CPU/gloo (DS_ACCELERATOR=cpu LOCAL_SIZE=4): 109 passed across the resume matrix, the producer and coverage tests, the vocab cross-TP test, the affine unit suite and tests/unit/runtime/tensor_parallel/. The five remaining failures in that file are FusedAdam JIT-compile errors on this machine and are identical on upstream master. ### Review findings Three things came out of review, all fixed here: - affine_map was added to every conversion metadata dict, changing the schema even for layouts with no map. Now emitted only where there is one, so an existing layer publishes exactly what it did before. - The producer test checked that the expected maps were present but not that nothing extra was. It now asserts the pattern set exactly. - Parameters the TP machinery leaves untouched — embeddings, norm weights — were classified TP_REPLICATED with no map, so conversion still fell back to the category. TestAffineMapCoverage now asserts the general invariant: every pattern in any category must carry a map, with only AutoTP-unsupported layouts exempt. It fails without the fix. - TestUnevenVocabCrossTpResume covers a 101-row vocabulary head saved at TP2 and restored at TP1 and TP2. The 51/50 split is the point — restoring at TP1 merges two unequal shards, where an even-split assumption would surface. The fixture asserts the emitted map records that split before saving. Adapted from @jinyouzhi's cross-tp test in deepspeedai#8309, using the column partition already in tree so it carries no dependency on that PR. affine_ir_spec.md gains §6.4 on where a map is built, and a §6.3 rule on the map and categories coexisting. Related: deepspeedai#8252, deepspeedai#8230. cc @delock --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to #8385. The scale move answering @delock's question was pushed after the merge queue had already snapshotted the branch, so it did not land with the rest.
The layouts that pre-divide a value hold it whole on every rank: Yuan's o_proj and the last conv layer both replicate the bias divided by the world size, so the all-reduced sum adds it exactly once. The weight beside them is what gets split, and it is unscaled — so no in-tree layout scales a split, and the split constructors no longer take the argument.
Adds a test covering the replicated case.