[pengdurice]

Part of: #8672 (Q4 2026 roadmap), "Enable training of large scale models".

Summary

vocab_parallel_lm_head (#8309) shards an input embedding only when it is tied to the output layer. An untied input
embedding stays replicated on every tensor-parallel rank: on Qwen2.5-72B at TP 8 it is 1.25B of each GPU's 10.18B
parameters (12%), with its gradient, master copy and Adam states, and its gradient is the last one produced in
backward. This PR adds tensor_parallel.vocab_parallel_embedding (default false): AutoTP replaces the model's
untied nn.Embedding by the existing VocabParallelEmbedding (1/tp of the rows per rank, masked lookup, sum over
the TP group).

"tensor_parallel": {
    "autotp_size": 8,
    "vocab_parallel_embedding": true
}

Results

Qwen2.5-72B, 4 nodes x 8 H200 (EFA), AutoTP 8 x ZeRO-1 over 4, micro-batch 2, 64 x 8,192 tokens per step, bf16 with
fp32 master weights and Adam states, no recompute, Liger kernels for RMSNorm, rotary embedding and SwiGLU (applied by
our benchmark script); the two runs in one allocation, with this code in our study tree (together with the AutoTP
sequence-sharding changes, which are separate PRs):

	Without	With
Step time	15.88 s	15.77 s
Peak memory (rank 0)	114.1 GiB	107.0 GiB
Embedding parameters per GPU	1.25B	0.16B
Gradient reduce after the last backward (wall_clock_breakdown)	89 ms	44 ms

Duplicate check

Searched issues and pull requests on 2026-10-18 for vocab_parallel_embedding, "untied embedding AutoTP" and
VocabParallelEmbedding untied: #8309 and #8680 handle tied embeddings, #8173 (open RFC) discusses sharding output
layers; none shards an untied input embedding.

The change

File	Change
deepspeed/module_inject/layers.py (+26)	shard_untied_input_embedding(model, mp_group, tp_meta): replaces model.get_input_embeddings() by a VocabParallelEmbedding; returns None when it already is one (a tied embedding sharded by vocab_parallel_lm_head); raises NotImplementedError when the embedding's weight is shared with another module, so a tied embedding is never untied
deepspeed/runtime/engine.py (+6)	calls it in finalize_autotp, before the replicated-gradient hooks are registered, when the option is set
deepspeed/runtime/tensor_parallel/config.py (+7)	the option
docs/_pages/config-json.md, docs/_tutorials/autotp-training.md	the option, and a short "Untied Input Embedding" section
tests/unit/v1/autotp/test_autotp_training.py (+83)	TestVocabParallelEmbedding

Tests

TestVocabParallelEmbedding (2 ranks, AutoTP 2, a 2-layer untied Llama built from a config):

• test_matches_single_device_training (ZeRO 0 and 1): the input embedding becomes a VocabParallelEmbedding with
  half of the rows; with tokens from both halves of the vocabulary, the losses of 2 SGD steps match the unsharded
  model, and the gathered embedding weight matches the unsharded model's after them.
• test_rejects_tied_embedding_without_vocab_parallel_lm_head: a tied embedding with vocab_parallel_lm_head: false
  raises instead of being untied from the output layer.

Hardware: 1 node, 4 x H200 (driver 580.126.09), torch 2.8.0+cu126, transformers 5.15.1, pytest 9.1.1, pytest-forked
1.7.5.

cd tests && pytest --forked unit/v1/autotp/test_autotp_training.py unit/v1/autotp/test_tp_plan_real_models.py
# 52 passed, 2 skipped (the existing BF16-optimizer overflow skips), 59 deselected (sequential) in 12:50
pytest --forked -m sequential unit/v1/autotp/test_autotp_training.py
# 59 passed in 12:48

pre-commit run --files passes on every changed file.

Compatibility

Cell	Validated
AutoTP 2, ZeRO 0 and 1, fp32 (unit test)	yes
AutoTP 8 + ZeRO-1, bf16, Qwen2.5-72B, 4 nodes	yes (study tree)
ZeRO 2 and 3	not run
Checkpoint save / load and universal checkpoint conversion with the sharded embedding	not run; the replacement passes AutoTP's tp_meta to VocabParallelEmbedding
Tied embeddings	unchanged: left to vocab_parallel_lm_head; with it disabled, an error
Models whose get_input_embeddings() is not an nn.Embedding	rejected with ValueError
Accelerators other than CUDA	not run; uses the existing VocabParallelEmbedding

AI assistance

This change, its tests and this description were prepared with AI assistance (Claude).

🤖 Generated with Claude Code

[jinyouzhi]

Thanks for picking this up. This fills the untied half of the tied/untied unification in RFC #8173 (#8582 handled tied + vocab-parallel head), so it fits well. A few points from tracing the paths on master:

1. Tied + row-parallel head gives a misleading error. With a tied model and a ROW lm_head spec ([AutoTP] Support eager row-parallel output head training #8538), the embedding is already a HiddenParallelEmbedding, so shard_untied_input_embedding raises needs ... nn.Embedding, got HiddenParallelEmbedding. Could any embedding already replaced by a TP layer (e.g. isinstance(embedding, TensorParallel_Layer)) be treated like the VocabParallelEmbedding case (return None), or get an explicit message?

2. Double partitioning with partition_config. If a user spec matches the embedding (non-SKIP), _slice_embedding has already returned a new nn.Embedding sliced on the hidden dim ([V, H/tp], requires_grad=False). That passes the nn.Embedding check and is then sharded on the vocab dim as well. Suggest raising when the weight already carries DS_AUTOTP_UC_META (or otherwise detecting a prior slice) and the user also set vocab_parallel_embedding.

3. DeepCompile autotp. The HF tp_plan path explicitly keeps embeddings replicated when the DeepCompile autotp pass is on ("does not support vocabulary-parallel embeddings"), but finalize_autotp shards unconditionally here. A guard (error or warning + skip) would keep the two consistent.

4. HF embedding_rowwise for untied models. HF transformers adds embed_tokens: embedding_rowwise whenever a model has a _tp_plan, and torchtitan / Megatron also shard the input embedding along vocab by default. Today that request is honoured only for tied heads (auto-enabling vocab_parallel_lm_head); for untied models it is converted to SKIP. Would you consider making the flag tri-state like vocab_parallel_lm_head (None = follow the HF plan), either here or as a follow-up?

5. Tests. The Llama test model probably goes through the HF tp_plan branch, so the heuristic and partition_config branches of finalize_autotp are not exercised. Cases that would be useful:

   • tied + vocab_parallel_lm_head: true + vocab_parallel_embedding: true -> no-op path (return None), still matches the reference;
   • untied + vocab_parallel_lm_head: true (both head and embedding sharded);
   • one of the non-HF branches (e.g. partition_config).