Carry the affine scale on the replicated map, not the split

[Achyuthan-S]

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.

[chatgpt-codex-connector]

💡 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".

[chatgpt-codex-connector]

Add the required sign-off trailer

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 👍 / 👎.

[Copilot]

🟢 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 a scale parameter and carry that scale on its single replicated AffinePiece.
• Remove the scale parameter from contiguous_split_map(...), reflecting that split weights are unscaled in the supported in-tree layouts.
• Add a unit test verifying that replicated_map properly 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]

@Achyuthan-S can you fix formatting error? Thanks!

[Achyuthan-S]

@delock I don't think there's a formatting issue to fix — yapf crashed on the runner rather than reporting a diff:

from yapf_third_party._ylib2to3 import pygram
  ...pgen2/grammar.py
EOFError: Ran out of input

That's yapf's lib2to3 grammar cache being truncated in the pre-commit env. There's no would reformat line anywhere in the run, and every other hook passed.

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 .py files this PR touches. The PR's own formatting check is green too.

Happy to push an empty commit if that's easier than a re-queue, but I don't believe there's anything to correct.