[AutoTP] Fix training lm_head routing (#8302)
## 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.
---------
Signed-off-by: gaoxiaomo <165135449+gaoxiaomo@users.noreply.github.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>
