DeepSpeed

1a956270 - [VocabTP] vocab_parallel_ce support fused LM-head projection + loss for the HF forward (#8945)

[VocabTP] vocab_parallel_ce support fused LM-head projection + loss for the HF forward (#8945) ## Motivation `vocab_parallel_lm_head` (#8173 follow-ups) keeps the output vocabulary sharded, but each rank still materializes `[tokens, vocab / tp]` logits (plus their FP32 copies in the loss). #8792 added `vocab_parallel_linear_cross_entropy`, which fuses the projection into the loss, but its docs state it is "not yet wired into the HuggingFace forward or the `vocab_parallel_lm_head` configuration". This PR does that wiring, so HF `Trainer` / plain `deepspeed.initialize` users get it via config only. Tracked in the #8790 plan. **Not a duplicate:** searched open PRs/issues for `fused_linear`, `linear_cross_entropy`, `vocab_parallel_ce_backend`; the only related work is #8540 (Liger CE on logits, does not fuse the projection) and #8792 (the kernel itself). Liger's own FLCE has no TP support. ## What changes - `tensor_parallel.vocab_parallel_ce_backend` accepts `"fused_linear"`. - `VocabParallelLinear.forward`: when the installed fused loss is armed (training forward with `labels`), the head returns a `[..., 0]` placeholder and hands its input to the loss instead of projecting it. - `VocabParallelLinearCausalLMLoss` (installed as `model.loss_function`): HF label shift, `ignore_index`, `num_items_in_batch`; falls back to the logits-based vocab-parallel loss when nothing was captured (eval mode, no labels). - Explicit rejections: ZeRO-3 (head params freed before the loss runs), DeepCompile `autotp` pass, and models that soft-cap / rescale logits after the head (`final_logit_softcapping`, `logit_scale`, `logits_scaling`). - `vocab_parallel_linear_cross_entropy` now drops `ignore_index` tokens before the projection (prompt/padding tokens cost no GEMMs or collectives); the range check and the compacted length share one host sync. - Default chunk size keeps the FP32 logits chunk near the hidden-state size **but at least 256 tokens**. Without the floor, 4,096 valid tokens were split into 85 chunks of ~49 rows (narrow GEMMs, 170 tiny all-reduces): on L20 the op took 128 ms vs 47 ms at 256 rows (+150 MiB); larger chunks gave no further gain. - Docs: `docs/_tutorials/autotp-training.md` — new backend, HF `Trainer` usage, measurements, limitations, and the incompatibility of any vocab-parallel backend with callers that compute their own loss from logits / `lm_head.weight` (`Trainer(label_smoothing_factor=...)`, TRL ≥ 1.15 `SFTTrainer`/`GRPOTrainer` fused LM head). ## Downstream HF usage No HF code change needed. `transformers` `Trainer` reads `tensor_parallel.autotp_size` (`trainer.py` `get_tp_size`) and `integrations/deepspeed.py` calls `deepspeed.tp_model_init(model, tp_size, dtype, config=...)` before `deepspeed.initialize`; Accelerate (`accelerator.py`, deepspeed ≥ 0.16.4) does the same. Verified end-to-end with `Trainer` + the config below: ```json "tensor_parallel": {"autotp_size": 2, "vocab_parallel_lm_head": true, "vocab_parallel_ce_backend": "fused_linear"} ``` ## Results Qwen2.5-0.5B (151,936 vocab, tied), HF `Trainer` (transformers 5.17.0, accelerate 1.15.0), TP=2, ZeRO-1, BF16, gradient checkpointing, micro-batch 4 × 2,048 tokens, 50 % prompt-masked, 20 steps (mean of steps 6–20). 2× NVIDIA L20 (torch 2.14.1+cu130, NCCL): | backend | peak / rank | step | grad-norm steps 1–3 | | --- | --- | --- | --- | | `vocab_parallel_lm_head: false` | 18,860 MiB | 558 ms | 56.6 / 37.9 / 29.6 | | `"torch"` | 13,201 MiB | 514 ms | 56.6 / 37.9 / 29.6 | | `"liger"` (#8540) | 8,456 MiB | 477 ms | 56.6 / 37.9 / 29.6 | | `"fused_linear"` | **5,200 MiB (−61 % vs torch, −39 % vs liger)** | **467 ms** | 56.6 / 37.9 / 29.6 | 2× Intel Arc Pro B70 (torch 2.14.0.dev20260804+xpu, XCCL): | backend | peak / rank | step | grad-norm steps 1–3 | | --- | --- | --- | --- | | `vocab_parallel_lm_head: false` | 18,844 MiB | 715 ms | 56.63 / 37.95 / 29.58 | | `"torch"` | 13,185 MiB | 653 ms | 56.63 / 37.92 / 29.45 | | `"fused_linear"` | **5,184 MiB** | **590 ms (−9.6 %)** | 56.63 / 37.92 / 29.55 | Op only (TP=2, V=151,936, H=896, BF16, fwd+bwd, 8,192 tokens with 50 % ignored): | | L20 | B70 | | --- | --- | --- | | torch | — | 104 ms | | liger | 48 ms (+4.6 GiB) | 104 ms | | fused_linear, old default (~49 rows/chunk) | 128 ms | ~100 ms | | fused_linear, new default (256 rows/chunk) | **47 ms (+193 MiB)** | **64.6 ms (+193 MiB)** | Tried and **not** included: merging the per-chunk MAX + SUM all-reduces into one all-gather — no measurable change on 2× B70 (148.0 vs 148.4 ms), so the simpler code is kept. ## Tests Hardware: 2× Intel Arc Pro B70 (XCCL, torch 2.14.0.dev20260804+xpu) and 2× NVIDIA L20 (NCCL, torch 2.14.1+cu130); transformers 5.17.0. All tests pass on both (L20: 53 CE / linear-CE + 40 AutoTP). ```bash export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 # test harness on XPU without clinfo pytest tests/unit/v1/autotp/test_autotp_training.py -m sequential -k "VocabParallel or Tied" pytest tests/unit/v1/sequence_parallelism/test_vocab_parallel_cross_entropy.py \ tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py \ tests/unit/v1/sequence_parallelism/test_vocab_parallel_linear_cross_entropy.py ``` New tests: tied-embedding training parity for `fused_linear`; HF-forward contract (empty logits only when the loss is computed, `num_items_in_batch`, positional labels, eval/no-label fallbacks); rejection of Gemma2 softcap (any cap, including 1.0) and ZeRO-3; non-finite activations at ignored positions affect neither loss nor gradients; many-token BF16 parity on the accelerator. Compatibility combos verified on TP=2: unit tests with ZeRO-0, FP32, Llama/Gemma3, tied and untied heads, eval and no-label forwards; HF `Trainer` end-to-end with ZeRO-1, BF16 and FP32, gradient checkpointing on/off, 512–8,192 tokens per micro-batch (B70), and the four-backend `Trainer` comparison above on L20. ZeRO-3 and DeepCompile `autotp` are rejected. ## AI assistance Implementation, tests and this description were drafted with GitHub Copilot CLI and reviewed/validated by the author on the hardware above. --------- Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>