A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 17:45 UTC
[Bugfix][Model] Support LoRA on Qwen3 Output Embedding (#29816)
Signed-off-by: kurt <kurt@thinkingmachines.ai>
Kurt Shuster committed
Feb 6, 2026 at 04:25 UTC
2991dd3d2241cb3b188a047282e58111684e0201
6 files changed
+132
-13
tests/lora/test_qwen3_unembed.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for Qwen3 unembed LoRA support.
5
+
6
+This test creates synthetic LoRA weights that include lm_head (output embedding)
7
+to verify that Qwen3 properly supports LoRA on the unembed/lm_head layer.
8
+"""
9
+
10
+import json
11
+import os
12
+import tempfile
13
+
14
+import numpy as np
15
+import torch
16
+from safetensors.torch import save_file
17
+
18
+from vllm import LLM, SamplingParams
19
+from vllm.lora.request import LoRARequest
20
+
21
+MODEL_PATH = "Qwen/Qwen3-0.6B"
22
+HIDDEN_SIZE = 1024
23
+VOCAB_SIZE = 151936
24
+
25
+
26
+def create_qwen3_lora_with_lm_head(save_dir: str, rank: int = 8) -> None:
27
+ """Create synthetic Qwen3 LoRA weights with lm_head."""
28
+ lora_weights = {}
29
+ for module in ["q_proj", "v_proj"]:
30
+ lora_A = torch.from_numpy(
31
+ np.random.randn(rank, HIDDEN_SIZE).astype(np.float16) * 0.01
32
+ )
33
+ lora_B = torch.zeros(HIDDEN_SIZE, rank, dtype=torch.float16)
34
+ key_prefix = f"base_model.model.model.layers.0.self_attn.{module}"
35
+ lora_weights[f"{key_prefix}.lora_A.weight"] = lora_A
36
+ lora_weights[f"{key_prefix}.lora_B.weight"] = lora_B
37
+
38
+ # lm_head LoRA weights
39
+ lora_weights["base_model.model.lm_head.lora_A.weight"] = torch.from_numpy(
40
+ np.random.randn(rank, HIDDEN_SIZE).astype(np.float16) * 0.01
41
+ )
42
+ lora_weights["base_model.model.lm_head.lora_B.weight"] = torch.zeros(
43
+ VOCAB_SIZE, rank, dtype=torch.float16
44
+ )
45
+
46
+ adapter_config = {
47
+ "peft_type": "LORA",
48
+ "base_model_name_or_path": MODEL_PATH,
49
+ "task_type": "CAUSAL_LM",
50
+ "inference_mode": True,
51
+ "r": rank,
52
+ "lora_alpha": rank * 2,
53
+ "lora_dropout": 0.0,
54
+ "bias": "none",
55
+ "target_modules": ["q_proj", "v_proj", "lm_head"],
56
+ }
57
+
58
+ os.makedirs(save_dir, exist_ok=True)
59
+ with open(os.path.join(save_dir, "adapter_config.json"), "w") as f:
60
+ json.dump(adapter_config, f)
61
+ save_file(lora_weights, os.path.join(save_dir, "adapter_model.safetensors"))
62
+
63
+
64
+def test_qwen3_unembed_lora():
65
+ """Verify Qwen3 can load and generate with LoRA adapters with lm_head."""
66
+ with tempfile.TemporaryDirectory() as tmpdir:
67
+ # Initialize engine first (before creating torch tensors)
68
+ llm = LLM(
69
+ model=MODEL_PATH,
70
+ enable_lora=True,
71
+ max_loras=4,
72
+ max_lora_rank=8,
73
+ max_model_len=128,
74
+ gpu_memory_utilization=0.8,
75
+ enforce_eager=True,
76
+ )
77
+
78
+ # Create LoRA weights after engine init
79
+ create_qwen3_lora_with_lm_head(tmpdir, rank=8)
80
+
81
+ lora_request = LoRARequest("lm_head_lora", 1, tmpdir)
82
+ llm.llm_engine.add_lora(lora_request)
83
+
84
+ assert 1 in llm.llm_engine.list_loras(), "lm_head LoRA should be loaded"
85
+
86
+ # Test generation
87
+ sampling_params = SamplingParams(temperature=0, max_tokens=32)
88
+ prompts = ["Hello, my name is"]
89
+
90
+ # Generate with base model (no LoRA)
91
+ base_outputs = llm.generate(prompts, sampling_params, use_tqdm=False)
92
+ assert len(base_outputs) == 1
93
+ assert len(base_outputs[0].outputs[0].text) > 0
94
+
95
+ # Generate with lm_head LoRA
96
+ lora_outputs = llm.generate(
97
+ prompts, sampling_params, lora_request=lora_request, use_tqdm=False
98
+ )
99
+ assert len(lora_outputs) == 1
100
+ assert len(lora_outputs[0].outputs[0].text) > 0
vllm/lora/layers/logits_processor.py
+5
-2
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-
4
import torch
5
import torch.nn as nn
6
from transformers import PretrainedConfig
@@ -148,7 +147,11 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
147
embedding_bias: torch.Tensor | None = None,
148
) -> torch.Tensor | None:
149
# Get the logits for the next tokens.
151
- logits = lm_head.quant_method.apply(lm_head, hidden_states)
150
+ if hasattr(lm_head, "base_layer"):
151
+ actual_lm_head = lm_head.base_layer
152
+ else:
153
+ actual_lm_head = lm_head
154
+ logits = actual_lm_head.quant_method.apply(actual_lm_head, hidden_states)
155
if embedding_bias is not None:
156
logits += embedding_bias
157
vllm/lora/layers/vocal_parallel_embedding.py
-1
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-
4
import torch
5
import torch.nn as nn
6
import torch.nn.functional as F
vllm/lora/model_manager.py
+17
-10
@@ -452,16 +452,23 @@ class LoRAModelManager:
452
if module_name not in self.packed_modules:
453
assert embedding_modules is not None
454
if parts[-1] in embedding_modules:
455
- input_dim = (
456
- module.base_layer.org_vocab_size
457
- if hasattr(module.base_layer, "org_vocab_size")
458
- else module.base_layer.weight.shape[1]
459
- )
460
- output_dim = (
461
- module.base_layer.embedding_dim
462
- if hasattr(module.base_layer, "embedding_dim")
463
- else module.base_layer.weight.shape[0]
464
- )
455
+ # Special-case lm_head: wrapped by LogitsProcessorWithLoRA.
456
+ # LoRA input dim is hidden_size, output dim is vocab size.
457
+ # LogitsProcessorWithLoRA handles extra vocab size directly.
458
+ if parts[-1] == "lm_head":
459
+ input_dim = module.lora_a_stacked[0].shape[-1]
460
+ output_dim = module.lora_b_stacked[0].shape[-2]
461
+ else:
462
+ input_dim = (
463
+ module.base_layer.org_vocab_size
464
+ if hasattr(module.base_layer, "org_vocab_size")
465
+ else module.base_layer.weight.shape[1]
466
+ )
467
+ output_dim = (
468
+ module.base_layer.embedding_dim
469
+ if hasattr(module.base_layer, "embedding_dim")
470
+ else module.base_layer.weight.shape[0]
471
+ )
472
lora = LoRALayerWeights.create_dummy_lora_weights(
473
module_name,
474
input_dim,
vllm/model_executor/models/qwen3.py
+5
@@ -271,6 +271,11 @@ class Qwen3ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
271
],
272
}
273
274
+ embedding_modules = {
275
+ "embed_tokens": "input_embeddings",
276
+ "lm_head": "output_embeddings",
277
+ }
278
+
279
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
280
super().__init__()
281
config = vllm_config.model_config.hf_config
vllm/model_executor/models/qwen3_moe.py
+5
@@ -689,6 +689,11 @@ class Qwen3MoeForCausalLM(
689
]
690
}
691
692
+ embedding_modules = {
693
+ "embed_tokens": "input_embeddings",
694
+ "lm_head": "output_embeddings",
695
+ }
696
+
697
fall_back_to_pt_during_load = False
698
699
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):