DeepSpeed

b39bf260 - Keep Yuan's o_proj bias whole on every rank and describe it as replicated (#8857)

Keep Yuan's o_proj bias whole on every rank and describe it as replicated (#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 #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 --------- Signed-off-by: Achyuthan Sivasankar <achyuthan.sivasankar@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>