[Achyuthan-S]

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:

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 #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 #8252 — and I have not verified it.

Related: #8252, #8230, #8575, #8706.

cc @delock

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing.

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-10-19T00:06:54.220084Z	bed66b0	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to 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" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: bed66b0cc6

ℹ️ 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 a distributed Yuan bias regression test

When a biased Yuan o_proj is used at TP ≥ 2, this test bypasses the injected Yuan_LinearAllreduce and its process group by calling shard_value_with_share_qk directly and manually summing the products. It therefore cannot catch regressions in AutoTP routing, the wrapper's choice of the output-projection branch, or the runtime ordering of the all-reduce and bias addition—the integration behavior this fix changes. Add a distributed test that runs the injected layer against the unsharded output on an actual device.

AGENTS.md reference: AGENTS.md:L81-L82

Useful? React with 👍 / 👎.

[jinyouzhi]

looks good to me.