[xsxszab]

Overview

Qwen3-MoE (including the Qwen3-VL-MoE and Qwen3-Omni language models), Qwen3-Next, Qwen3.5 and gpt-oss build the input embedding, and some of them also the final norm and lm_head, on every pipeline stage. They now build these modules only on the stages that use them, as Llama does, so each stage loads 0.58 to 1.08 GiB less at PP=2.

Claims

• The input embedding is built on the first stage, and on the last stage when lm_head is tied to it or a drafter reuses it; the final norm and lm_head are built only on the last stage. This is the condition Llama and Qwen2 use.
• With PP=2, each stage loads 0.58 to 1.08 GiB less and gets a correspondingly larger KV cache; GSM8K stays within 0.007 of main.
• test_pp_tied_word_embeddings covers Qwen3.5-0.8B, whose lm_head is tied to the input embedding.

Validation

B300 GPUs, Model Runner V2.

.venv/bin/python -m pytest -v -s \
  tests/distributed/test_pipeline_parallel.py::test_pp_tied_word_embeddings -k qwen3_5

The qwen3_5 case compares PP=2 with PP=1 logprobs on Qwen3.5-0.8B and passes on 2× B300. CI runs this file in the Pipeline + Context Parallelism step: nightly on NVIDIA L4 and on PRs through its AMD MI250 mirror.

Greedy outputs with PP=2 (eager, one request at a time, 4 prompts) are identical to main token for token for Qwen3-30B-A3B (text) and Qwen3-VL-30B-A3B-Instruct-FP8 (3 images, 1 text prompt).

GSM8K (1319 questions, 5-shot, one run each) and model loading per stage, -pp 2 --max-model-len 4096 with torch.compile and CUDA graphs:

Model	GSM8K main	GSM8K this PR	Model loading PP0 / PP1, main	Model loading PP0 / PP1, this PR	KV cache main (tokens)	KV cache this PR (tokens)
gpt-oss-20b	0.313	0.318	8.01 / 8.01 GiB	6.93 / 6.93 GiB	10,814,845	10,861,857
Qwen3-30B-A3B	0.898	0.901	29.03 / 29.03 GiB	28.45 / 28.45 GiB	4,983,808	4,997,120
Qwen3-30B-A3B-FP8	0.904	0.897	15.11 / 15.11 GiB	14.53 / 14.53 GiB	5,283,840	5,297,152
Qwen3.5-35B-A3B	0.836	0.842	33.17 / 33.28 GiB	33.17 / 32.33 GiB	6,669,548	6,698,535
Qwen3-Next-80B-A3B-Instruct-FP8	0.859	0.861	38.06 / 38.06 GiB	37.48 / 37.48 GiB	7,764,811	7,785,773

gpt-oss-20b is evaluated through completions without its chat format, as in its GSM8K config (threshold 0.30). On main, Qwen3.5 already skips lm_head on PP0, so only PP1 changes; the KV cache is sized by the stage with less free memory, which is now PP0.

Other parallel layouts and speculative decoding

Same GSM8K setup:

Setting	GSM8K main	GSM8K this PR	Model loading per stage, main	Model loading per stage, this PR
Qwen3-30B-A3B, PP=4	0.892	0.895	15.1 GiB each	14.52 / 13.94 / 13.94 / 14.52 GiB
Qwen3-30B-A3B, TP=2 x PP=2	0.886	0.884	14.53 GiB each	14.24 GiB each
gpt-oss-20b + EAGLE3 (RedHatAI/gpt-oss-20b-speculator.eagle3, 3 tokens), PP=2	0.296, acceptance length 2.07	0.323, acceptance length 2.08	8.01 / 8.54 GiB	6.93 / 8.54 GiB
Qwen3.5-35B-A3B + MTP (2 tokens), PP=2	0.857, acceptance length 2.78	0.854, acceptance length 2.78	33.17 / 35.86 GiB	33.17 / 34.91 GiB

The EAGLE3 head reuses the target input embedding, so the last stage keeps it and only PP0 changes; the MTP drafter has its own embedding, so the target's is skipped on PP1.

Details

What each model built on every stage before this change:

Model	Modules
Qwen3-MoE, and the Qwen3-VL-MoE and Qwen3-Omni language models	input embedding, final norm, lm_head
gpt-oss	input embedding, final norm, lm_head
Qwen3-Next	input embedding, lm_head
Qwen3.5 (models/qwen3_5.py and models/qwen3_5/nvidia/model.py)	input embedding

The code follows Llama: the other stages get a PPMissingLayer, and tie_weights moves into the last-stage branch. Qwen3MoeModel reads tie_word_embeddings from hf_config.get_text_config(), the config its causal LMs tie by: inside Qwen3-VL-MoE and Qwen3-Omni, with_hf_config (#38035) updates hf_config but leaves hf_text_config with the text config's own value, which defaults to True for Qwen3-VL-MoE.

Limits:

• PP shards of these models saved with load_format="sharded_state" before this change hold the skipped modules and fail to load with a KeyError; they must be saved again.
• Converting a tied Qwen3.5 checkpoint such as Qwen3.5-0.8B to a classifier with classifier_from_token at PP>=3 fails on the middle stage, which has no input embedding to take the classifier rows from. Other tied models such as Qwen3-Reranker-0.6B already fail this way on main.

No open PR changes where these models build these modules.

AI assistance (Claude Code) was used for this PR.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.