DeepSpeed

468a8ad1 - [AutoTP] Support eager row-parallel output head training (#8538)

[AutoTP] Support eager row-parallel output head training (#8538) ## Summary Implements PR-E from #8173: eager training for explicitly row-parallel, untied `lm_head` / `embed_out` output heads. This PR is independent of the vocabulary-parallel CE work in PR-B #8309 and contains no PR-D changes. - Slice replicated activations along the last dimension using the frozen weight-shard geometry, including uneven hidden dimensions. - Reuse the existing row-parallel forward output reduction and replicated bias behavior. - Reconstruct and all-reduce the complete input gradient in backward so preceding replicated layers receive the correct gradient. - Preserve the original parameter identity, existing inference specialization, and default gathered-column training behavior. - Reject tied output weights at TP > 1, reshaped/non-input-dimension partitions, and deferred compiler collectives for this path. ## Validation Executed on jk01 with two NVIDIA L20 GPUs (46068 MiB each), in an isolated container: Python 3.12.3, PyTorch 2.13.0+cu130, CUDA 13.0, Transformers 5.14.1. - Partition/config regression: **17 passed**. - Real two-rank DeepSpeed-engine integration: **three configurations each completed five optimizer updates**, with two accumulated micro-batches per update (SGD LR 0.05, ZeRO stage 0, FP32, clipping explicitly disabled in both paths). - Configurations cover `lm_head` / `embed_out`, bias / no bias, hidden dimensions 32 / 35, and flattened / sequence-shaped inputs. - Compared logits, CE loss, input gradients, replicated projection gradients, local head gradients, and updated weights against an independent unsharded PyTorch model after every update. Also checked the exact SGD update and that weights actually changed. - Changed-file pre-commit checks and `git diff --check` passed. The integration method `TestRowParallelOutputHeadTraining.test_five_optimizer_steps_match_unsharded_reference` was executed through two-rank torchrun on the real devices. It can also be run through pytest: ```sh PYTHONPATH=.:tests python3 -m pytest tests/unit/module_inject/test_tp_partition_config_path.py -q PYTHONPATH=.:tests python3 -m pytest tests/unit/v1/autotp/test_autotp_training.py::TestRowParallelOutputHeadTraining -v -s ``` This validates the minimal eager training contract, not large-model convergence, mixed-precision/ZeRO checkpoint round trips, or DeepCompile. Coupled tied-embedding sharding remains outside this PR's scope. Signed-off-by: 0z5a <dezhen.lu@uni-tuebingen.de> Co-authored-by: 0z5a <dezhen.lu@uni-tuebingen.de> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>