[AutoTP] Fix training lm_head routing

[gaoxiaomo]

Summary

This implements the independently mergeable PR-A described in #8173.

• Add an explicit training_mode to AutoTP and propagate it from the training and inference entry points. Layer-type routing no longer depends on the sticky process-global AutoTP mode.
• During training, route an untied legacy lm_head / embed_out to the column-parallel LinearLayer with gather_output=True. Every TP rank therefore receives full-vocabulary logits and can continue to use standard cross entropy with autograd.
• Preserve explicit row-parallel output-head plans during training.
• Keep the inference-only LmHeadLinearAllreduce routing unchanged.
• Keep tied embedding/output-head weights replicated when the legacy path cannot shard both sides of the tie consistently.

Here, the legacy path means AutoTP model injection without a converted HF/custom partition-plan rule for the output head. Vocabulary-parallel cross entropy remains a separate follow-up: it can later switch the gathered training output to sharded logits together with the matching loss implementation.

Why gather_output=True

Column-parallelizing the vocabulary dimension without gathering leaves each rank with only its local vocabulary shard. Standard cross entropy then either rejects labels outside that shard or computes an incorrect denominator over a partial vocabulary. Gathering restores full logits on every rank, while GatherFromTensorParallelRegion preserves the backward path to each local weight shard.

Tests

• pytest -q tests/unit/module_inject/test_tp_partition_config_path.py: 14 passed
• TP=2 distributed regression test with an uneven vocabulary of 269 tokens: 1 passed
  • complete logits on both ranks: [4, 269]
  • TP loss equals the unsharded reference loss: 5.735173225402832
  • maximum logits error: 0
  • maximum input-gradient error: 1.1175870895385742e-08
  • local weight and bias gradients match the corresponding reference shards
• Baseline reproduction before the gather fix:
  • rank-local logits were [4, 135] and [4, 134]
  • label 268 failed with IndexError: Target 268 is out of bounds
• pre-commit run --files ...: all applicable hooks passed
• git diff --check: passed

The branch is rebased on current master at 6e3bd087.

Duplicate check

No open PR was found for #8173 PR-A or these two lm_head routing cases. #8241 touches the same AutoTP area but addresses per-model metadata isolation and does not change this routing behavior.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 986497b54f

ℹ️ 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".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[jinyouzhi]

Thank you for contributing this great PR! @gaoxiaomo Could you check whether the unit test coverage for tp_size > 1 is sufficient and behaves as expected?

I also wanted to better understand how the column-parallel and legacy paths are selected, particularly which scenarios still require the legacy path—for example, tied embeddings.

CC @delock for visibility. Please feel free to share any thoughts or comments.

[delock]

TP > 1 legacy training is unusable after PR-A alone (Hunk 1)

The training-path LinearLayer in _replace is created without gather_output, so it defaults to False. With TP > 1 each rank emits only its vocab/TP logits shard, and any standard cross-entropy breaks:

• labels outside the local shard → IndexError (or device-side assert)
• labels inside → silently wrong loss (logsumexp over half the vocab)

The old path was equally unusable for training (inference_all_reduce has no autograd), so this is not a regression — but it does mean TP > 1 legacy training stays broken until PR-B lands, so "independently mergeable" doesn't hold as-is.

Suggestion: pass gather_output=True here, mirroring the #8146 transition step (column-parallel + gather + standard CE — GatherFromTensorParallelRegion already supports autograd). PR-B can then flip both the config path and the legacy path to gather_output=False + vocab-parallel CE together. This also makes legacy-path training behavior identical to the config path (colwise_gather_output).

@jinyouzhi could you confirm whether this matches your RFC #8173 roadmap intent for PR-A?

[delock]

The new routing reads a sticky process-global mode — this breaks the OPD/OPSD same-process rollout case

set_autotp_mode is never reset (the only two call sites both set training=True, nothing ever sets it back), so once training is initialized the process stays in TRAINING forever. Before this PR the global only affected behavior knobs (add_bias in-place, all_reduce flavor), where leakage is mostly harmless. This PR makes it decide layer types for the first time:

training_column_lm_head = is_autotp_training_mode() and legacy_lm_head   # _replace
... and not is_autotp_training_mode()                                    # _create_row_parallel_layer

In OPD/OPSD the teacher rollout engine and training live in one process (the RFC #8173 motivation scenario): after deepspeed.initialize(..., autotp_size>0) sets TRAINING, any later teacher-side AutoTP construction (re-sharding after weight sync, CUDA-graph recapture, or teacher built after student) hits is_autotp_training_mode() == True and gets lm_head → LinearLayer(gather_output=False). Rollout then emits vocab/TP sharded logits → HF sampler crashes, or silently samples wrong tokens — hard-to-debug data corruption for distillation/RL.

Suggestion: these two new routing decisions should not read the global. Pass an explicit training_mode into AutoTP.__init__ (callers know: init_inference → False, training paths → True) and route on that. The sticky global can stay for the behavior knobs it already serves — only the layer-type selection introduced here needs to be explicit.

[gaoxiaomo]

Thanks @jinyouzhi and @delock. I updated the PR to address both review points.

• Legacy training now uses LinearLayer(..., gather_output=True), so TP>1 produces complete vocabulary logits on every rank and standard cross entropy remains correct and differentiable.
• Layer-type routing now uses an explicit AutoTP.training_mode. Training callers pass True; inference callers pass False. The existing process-global mode is no longer consulted for the new lm_head routing decisions.
• Legacy tied embedding/output-head weights remain replicated when both sides of the tie cannot be partitioned consistently.
• Explicit row-parallel plans are still respected during training, while the inference-only LmHeadLinearAllreduce path remains unchanged.

I also added a real TP=2 distributed regression test with an uneven vocabulary size of 269. Before the fix, the two ranks emitted [4, 135] and [4, 134] logits and label 268 raised IndexError. After the fix, both ranks emit [4, 269]; logits and loss match the unsharded reference, and input, weight-shard, and bias-shard gradients match as well.

Validation on the rebased branch:

• routing tests: 14 passed
• TP=2 distributed test: 1 passed
• all applicable pre-commit hooks passed
• git diff --check passed

The branch is now rebased on current master (6e3bd087).

[delock]

Thanks for the quick turnaround — the gather_output=True fix and the explicit training_mode plumbing both look good, and TestLegacyLmHeadGatheredTraining is exactly the kind of coverage I was hoping for.

The row-parallel point isn't closed yet, though. In _create_row_parallel_layer, training now falls through to plain LinearAllreduce, whose forward does matmul(input, weight.T) with no input slicing. lm_head receives a replicated full-width hidden state, so at TP > 1 this is H vs H/TP and fails. LmHeadLinearAllreduce avoids it by slicing the input first; the training path has no equivalent. The new tests set mp_size = 1 and never run forward, so they can't catch this.

Two ways to close it — I'd prefer (b) for this PR:

(a) Implement the LmHeadRowParallel autograd function I sketched. Note the dx all_reduce in backward is required; without it the input gradient only covers this rank's shard.

(b) Raise NotImplementedError for training + explicit ROW + mp_size > 1, and leave the real implementation to a follow-up. That keeps PR-A's scope tight, and (a) is a self-contained feature that deserves its own review. @jinyouzhi do you think we should list this as a separate item in the #8173 plan?

Either way the test needs mp_size = 2: a real forward/backward against an unsharded reference for (a) — same shape as TestLegacyLmHeadGatheredTraining — or a pytest.raises(NotImplementedError) for (b).

One related gap: the tied fallback in _configure_gathered_column_tie_fallbacks only matches COLUMN + gather_output, so a tied ROW lm_head still reaches LinearAllreduce and shards the shared Parameter along the hidden dim, taking the input embedding with it. LmHeadLinearAllreduce used to decouple this via weight.clone().detach(). Option (b) would block this case too.

[jinyouzhi]

Thanks, @delock. I agree with your suggestions.

The column-parallel direction matches the RFC #8173 plan. The legacy training path now uses column-parallel + gather_output=True + standard cross entropy, and the HF colwise_gather_output plan is converted to the same behavior. The training callers, including the explicit config and HF tp_plan paths, pass training_mode=True, while inference passes training_mode=False.

For a generic explicit partition_config, gather_output remains controlled by spec.gather_output and defaults to False. Thus, an LM head intended to use the PR-A standard-CE transition must explicitly configure gather_output=True. This gives it the same behavior as the legacy and HF gathered-column paths. PR-B can then switch both paths together to gather_output=False with vocabulary-parallel cross entropy.

I also agree that training support for explicit ROW-parallel output heads should be tracked as a separate follow-up item in RFC #8173. For PR-A, training + explicit ROW + mp_size > 1 should be rejected with NotImplementedError, since the current LinearAllreduce assumes an already-sharded hidden state and does not implement the required LM-head input slicing and dx all-reduce in backward. The tied ROW case should be addressed as part of the same follow-up.

[gaoxiaomo]

Thanks @delock and @jinyouzhi. I implemented option (b) in f42b1ad.

• Training with an explicit ROW-parallel lm_head / embed_out now raises NotImplementedError when mp_size > 1, before LinearAllreduce can shard the output head (or a tied embedding parameter) incorrectly.
• TP=1 training behavior is unchanged.
• The inference path continues to use LmHeadLinearAllreduce.
• The regression test covers both untied and tied output heads at mp_size = 2.

Validation:

• tests/unit/module_inject/test_tp_partition_config_path.py: 16 passed
• Focused TP>1 rejection cases: 2 passed
• git diff --check: passed

The row-parallel autograd implementation, including input slicing and the required dx all-reduce, remains a separate follow-up as discussed.