DeepSpeed

4bd251b6 - [AutoTP] Add optional Liger vocabulary-parallel cross entropy (#8540)

[AutoTP] Add optional Liger vocabulary-parallel cross entropy (#8540) ## Dependency and review scope Depends on #8309 (PR-B from #8173), currently open. This branch is based on its head `1b4ba6b162e87b10bedbb7316466bd877e2cb278`; inherited PR-B commits retain their original authorship. The final commit is the optional Liger backend work. Please review that commit for this PR's independent changes. This is a draft until PR-B lands; the master comparison currently includes the dependency changes. This addresses the optional fused CE and memory-profiling direction in #8173. It does not include row-parallel output-head training (#8538) or coupled tied-embedding sharding. ## Changes - Add opt-in `tensor_parallel.vocab_parallel_ce_backend: "liger"` alongside `vocab_parallel_lm_head: true`; default `"torch"` remains dependency-free. - Lazy-load Liger's public vocabulary-parallel CE function (`liger-kernel>=0.8.1`) and retain the existing reduction / causal-label-shifting contract. - Use identical gathered shard metadata to select Liger on every TP rank only for equal vocabulary shards, nonempty CUDA tokens, and FP32/FP16/BF16 logits. Uneven shards or unsupported layouts uniformly fall back to the PyTorch reference. - Reject repeated retained-graph backward and higher-order gradients for Liger's destructive saved-buffer backward instead of silently producing incorrect gradients. - Add optional-backend numerical and real AutoTP training integration coverage, and document the compatibility boundaries. This fuses CE only, **not the linear projection**. Local vocabulary logits are still materialized. ## Real-device validation jk01: two NVIDIA L20 GPUs (46068 MiB each), isolated container, Python 3.12.3, PyTorch 2.13.0+cu130, CUDA 13.0, Transformers 5.14.1, Triton 3.7.1, liger-kernel 0.8.2. - Partition/config and non-distributed CE regressions: **42 passed**, 7 distributed cases deselected in that regression invocation. - Two-rank numerical tests: FP32/FP16/BF16, none/sum/mean, even/uneven vocab, ignored targets, unchanged input logits, retained-graph rejection, and higher-order rejection passed. - Real DeepSpeed-engine training: vocab 34 (Liger) and vocab 35 (uneven fallback), each **five optimizer updates**, with two accumulated micro-batches per update. SGD LR 0.05, FP32, ZeRO 0, clipping explicitly disabled. Logits, loss, gradients, updated sharded head weights/biases, and replicated projection weights/biases agree with an independent unsharded PyTorch model. - Changed-file pre-commit checks and `git diff --check` passed. The distributed integration methods in `TestLigerVocabParallelCE` were executed via two-rank torchrun on the actual devices. Reproduce with pytest: ```sh PYTHONPATH=.:tests python3 -m pytest tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py -v -s ``` ## Memory profile Compared local linear + vocabulary-parallel CE forward/backward + five SGD updates, two micro-batches each, after kernel warm-up and resetting weights/allocated peaks. BF16, TP=2, 1024 tokens per rank, hidden=128, global vocabulary=65536, local vocabulary=32768. Statistics are the maximum across ranks. | Backend | Peak allocated | Incremental peak over baseline | Short-run elapsed | | --- | ---: | ---: | ---: | | PyTorch | 512247296 bytes (488.52 MiB) | 478160384 bytes | 0.0606 s | | Liger | 310922240 bytes (296.52 MiB) | 276834304 bytes | 0.0779 s | Peak allocated decreases approximately 39.3%; incremental peak decreases approximately 42.1%. Losses and final local weights agree. This is a short single-sample memory comparison, **not evidence of throughput speedup**; allocated memory is not total device or reserved memory. Four-GPU TP x SP integration, mixed-precision/ZeRO checkpoint round trips, and large-model convergence are not validated by this two-GPU run. --------- Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>