[0z5a]

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:

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.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5453799d02

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Avoid treating the default world as the TP group

When a multi-rank data-parallel job calls the public API with backend="liger" and leaves tp_group=None, the preceding metadata logic treats TP as disabled, but this call passes None to Liger's distributed group argument, where it denotes the default world process group. Liger then interprets the DP replicas' full-vocabulary logits as separate vocabulary shards, producing an incorrect global loss and gradients. Fall back to the PyTorch implementation in this case or pass an explicit singleton TP group.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Add the mandatory commit sign-off

Commit 5e07eda8cb904aad0dce2b9a9df136255b432c78 is a non-merge commit but its message has no Signed-off-by trailer, so it does not satisfy this repository's commit/CI requirements and will fail DCO-style validation. Recreate the commit with the configured author identity and the required sign-off.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[jinyouzhi]

@0z5a Could you have a rebase? Thanks

[delock]

liger kernel is written on trion, does it mean it could run on not only cuda device but any accelerator supports triton? May need something cuda independent here.

[delock]

just tested liger kernel works for xpu device as well (with triton and liger kernel installed), so guard it by cuda should be too restrict.

[0z5a]

Agree. Can our CI test xpu or other device?

[delock]

current XPU does not have regular nightly run. I think allow the test run in v1 (install liger-kernel in workflow if necessary) and guard by whether triton is supported (accelerator.is_triton_supported()) should be good enough. This is a contract between accelerator and deepspeed whether triton kernels should be supported.

[delock]

Would this test be skipped on current Modal runner?

[jinyouzhi]

Need to check whether VocabParallelCausalLMLoss accepts the backend argument

[0z5a]

@delock @jinyouzhi Thanks for review. I will check it tomorrow. Zzz...

[0z5a]

Addressed all four review points in 1152579.

tp_group=None must not select Liger (@chatgpt-codex-connector): _liger_can_run now returns False without an explicit tensor-parallel group, so a data-parallel job that leaves tp_group=None uses the PyTorch reference. Before the fix an accelerator tensor with local_vocab_size == global_vocab_size returned True and forwarded that None to Liger's group argument. Liger 0.8.1/0.8.2 happen to read None as "not distributed", so the kernel would have been numerically correct there, but the guard no longer depends on that version-specific reading. TestLigerBackendSelection::test_missing_tp_group_uses_pytorch_reference pins the dispatch (it asserts the Liger entry point is not called) and test_tp_group_selects_liger_only_for_equal_shards pins the complementary half: an explicit group still selects Liger for equal shards and still falls back for uneven ones.

CUDA-only guard (@delock): the gate is now get_accelerator().is_triton_supported(), the interface every DeepSpeed accelerator already implements (CUDA reports it from compute capability 8.0), so a Triton-capable XPU takes the same path. The status message and the tutorial text follow it.

Test skipped on the Modal runner (@delock): yes — it was skipped there, because the sandbox had no liger-kernel and importorskip turned that into a silent skip. liger-kernel>=0.8.1 is now in requirements/requirements-dev.txt, which ci/torch_latest.py installs from the candidate tree inside the sandbox, so the file executes in v1. The CUDA-only skip is replaced by the Triton-support check, and the two dispatch tests that do not need Liger stay runnable without it.

VocabParallelCausalLMLoss did not accept backend (@jinyouzhi): confirmed, and it was a hard break — configure_vocab_parallel_loss(..., backend=...) raised TypeError: VocabParallelCausalLMLoss.__init__() got an unexpected keyword argument 'backend' for any vocab_parallel_ce_backend value. The constructor now takes backend="torch", stores it, and forwards it on each call; test_vocab_parallel_lm_loss_forwards_backend covers the round trip and test_five_optimizer_steps exercises it through deepspeed.initialize. The _liger_vocab_cross_entropy call site also needed one space of reindentation for the repo's yapf config.

DCO (@chatgpt-codex-connector): the flagged commit was dropped when the branch was rebased onto the merged #8309; both commits now carry Signed-off-by, and the DCO check passes.

Validation on the current head (1152579)

Two RTX 5090 (32607 MiB each), Python 3.12.13, torch 2.13.0+cu130, Triton 3.7.1, transformers 5.10.4, liger-kernel 0.8.2.

• tests/unit/v1/sequence_parallelism/test_liger_vocab_parallel_cross_entropy.py: 7 passed.
• tests/unit/v1/sequence_parallelism/ (whole directory): 56 passed, 43 skipped, 0 failed.
• tests/unit/v1/autotp/test_autotp_training.py -m sequential: 3 passed, 24 skipped, 0 failed — the three vocab_parallel_lm_head paths, including the heuristic head replacement and the checkpoint-parity round trip. The skips are world-size-4 and optional-dependency cases; only two GPUs were available.

Stage 1: the CE op alone (forward + backward, TP=2, BF16, local vocab 16000)

Tokens	Backend	ms	peak allocated GiB	loss
8192	torch	6.530	2.197	33.670269
8192	liger	5.112	1.465	33.670269
16384	torch	11.635	4.395	33.731960
16384	liger	6.084	2.930	33.731960

The fused op is 1.28x faster at 8192 tokens and 1.91x at 16384, with 33% less peak allocation at both sizes, and the losses agree exactly.

Stage 2/3: E2E AutoTP training, real checkpoints and data

vocab_parallel_lm_head: true, TP=2, BF16, AdamW, 8192 tokens/step, one fresh two-rank process per arm, A/P/P/A order. Data is a fixed 8M-token stream from llm-wizard/alpaca-gpt4-data, tokenized per model; the timed window excludes weight loading, partitioning and 8 warmup steps. No arm fell back: every Liger arm reports backend_observed=liger, both ranks hold an equal 16000-of-32000 (or 25152-of-50304) shard, and no arm logged the fallback warning.

Block	Position	Backend	s/step	global tokens/s	peak allocated GiB	mean loss	observed backend
TinyLlama seq2048 x mb4	0	torch	0.4441	18,447	19.809	4.4471	torch
TinyLlama seq2048 x mb4	3	torch	0.4420	18,533	19.809	4.4470	torch
TinyLlama seq2048 x mb4	1	liger	0.4441	18,447	19.099	4.4471	liger
TinyLlama seq2048 x mb4	2	liger	0.4453	18,398	19.099	4.4471	liger
TinyLlama seq4096 x mb2	0	torch	0.4586	17,864	19.810	5.8857	torch
TinyLlama seq4096 x mb2	3	torch	0.4579	17,888	19.810	5.8856	torch
TinyLlama seq4096 x mb2	1	liger	0.4576	17,900	19.099	5.8856	liger
TinyLlama seq4096 x mb2	2	liger	0.4578	17,896	19.099	5.8857	liger
pythia-1.4b seq2048 x mb4	0	torch	0.4100	19,982	20.227	4.4588	torch
pythia-1.4b seq2048 x mb4	3	torch	0.4135	19,812	20.227	4.4538	torch
pythia-1.4b seq2048 x mb4	1	liger	0.4059	20,181	18.693	4.4648	liger
pythia-1.4b seq2048 x mb4	2	liger	0.4049	20,233	18.693	4.4579	liger
Llama-3.1-8B seq2048 x mb1	0	torch	13.8762	148	21.823	5.0370	torch
Llama-3.1-8B seq2048 x mb1	3	torch	13.5853	151	21.823	5.0348	torch
Llama-3.1-8B seq2048 x mb1	1	liger	14.1002	145	21.823	5.0537	liger
Llama-3.1-8B seq2048 x mb1	2	liger	13.5516	151	21.823	5.0347	liger

Block	Model	Vocab	torch tokens/s	liger tokens/s	throughput gain	torch s/step	liger s/step	step time change	torch peak GiB	liger peak GiB	peak change	loss rel. diff
TinyLlama seq2048 x mb4	TinyLlama-1.1B-Chat-v1.0	32,000	18,490	18,423	-0.36%	0.4431	0.4447	+0.37%	19.809	19.099	-3.58%	2.0e-05
TinyLlama seq4096 x mb2	TinyLlama-1.1B-Chat-v1.0	32,000	17,876	17,898	+0.12%	0.4583	0.4577	-0.12%	19.810	19.099	-3.59%	1.1e-05
pythia seq2048 x mb4	pythia-1.4b-deduped	50,304	19,897	20,207	+1.56%	0.4117	0.4054	-1.54%	20.227	18.693	-7.59%	1.1e-03
Llama-3.1-8B seq2048 x mb1	Meta-Llama-3.1-8B-Instruct	128,256	149	148	-0.67%	13.7300	13.8232	+0.68%	21.823	21.823	+0.00%	1.7e-03

pythia-1.4b is the GPTNeoX embed_out head, so the run also covers the second name in the head-pattern list.

Reading the numbers

• TinyLlama (vocab 32k) shows no throughput difference: both blocks sit inside +/-0.4%, which is the spread between two arms of the same backend in the same block. CE is roughly 6% of the forward FLOPs there and it is bandwidth-bound, so a 0.44 s/step contains only a few milliseconds of it. The reproducible E2E effect is memory: 0.71 GiB per rank (3.6%) lower peak allocation in both blocks.
• pythia-1.4b (vocab 50k) is the one block where the direction is consistent: +1.56% throughput with a -1.54% step time, and 1.53 GiB (7.6%) lower peak allocation. The within-backend spread is 0.9% for torch and 0.26% for Liger there, so the gain is small but not obviously noise. It is a single block, so treat it as PROMISING, not confirmed.
• Llama-3.1-8B is a scale smoke test only. AdamW state has to be host-offloaded to fit 32 GiB, so the step is PCIe-bound: peak allocation is byte-identical across all four arms and the loss spread between two identical Liger arms (5.0347 vs 5.0537) is larger than the torch/Liger difference. It shows the delegation works at 8B with a 128,256-row vocabulary, and nothing about throughput.
• A TinyLlama seq4096 x mb4 configuration (16,384 tokens/step) OOMs on this 32 GiB card for both backends, so the 0.71 GiB saving does not move the feasible batch-size envelope at that model size. The fused CE is a memory-headroom option, not a throughput claim.
• Losses agree to float noise wherever the run is long enough to compare (2e-5 relative on TinyLlama); exact gradient and parameter parity is established by the unit tests, which compare against an unsharded reference layer by layer.

[chatgpt-codex-connector]

To use Codex here, create a Codex account and connect to github.

[jinyouzhi]

LGTM