[Core][LoRA] Support fp32 lm_head (head_dtype) on the LoRA path (#48525)
Signed-off-by: Karthik Kothuri <karthikkothuri2009@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Karthik Kothuri committed
Jul 14, 2026 at 18:51 UTC
4f7fffb92f6e50fbbee47a89f6ebb4b662aba135
3 files changed
+11
-29
tests/lora/test_layers.py
+4
-1
@@ -368,8 +368,9 @@ def test_embeddings(
368
@pytest.mark.parametrize("device", DEVICES)
369
@pytest.mark.parametrize("vocab_size", [64000, 256512, 258048])
370
@pytest.mark.parametrize("stage", STAGES)
371
+@pytest.mark.parametrize("head_dtype", [None, torch.float32])
372
def test_lm_head_logits_processor(
372
- default_vllm_config, dist_init, num_loras, device, vocab_size, stage
373
+ default_vllm_config, dist_init, num_loras, device, vocab_size, stage, head_dtype
374
) -> None:
375
if current_platform.is_cuda_alike() or current_platform.is_xpu():
376
torch.accelerator.set_device_index(device)
@@ -391,6 +392,8 @@ def test_lm_head_logits_processor(
392
linear.weight.data = torch.rand_like(linear.weight.data)
393
linear.weight.data[:, vocab_size:] = 0
394
logits_processor = LogitsProcessor(vocab_size)
395
+ # Exercise an fp32 lm_head (head_dtype) on the LoRA path.
396
+ logits_processor.head_dtype = head_dtype
397
lora_logits_processor = LogitsProcessorWithLoRA(
398
logits_processor, 1024, linear.weight.dtype, linear.weight.device, None
399
)
tests/v1/sample/test_head_dtype.py
-16
@@ -136,22 +136,6 @@ def test_get_top_tokens_honors_head_dtype(default_vllm_config):
136
assert torch.equal(top, expected)
137
138
139
-def test_fp32_head_rejected_with_lora(default_vllm_config):
140
- from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
141
-
142
- base = _build_processor(64)
143
- base.head_dtype = torch.float32
144
-
145
- with pytest.raises(ValueError, match="not yet supported with LoRA"):
146
- LogitsProcessorWithLoRA(
147
- base,
148
- hidden_size=16,
149
- dtype=torch.bfloat16,
150
- device=torch.device("cpu"),
151
- sharded_to_full_mapping=None,
152
- )
153
-
154
-
139
@pytest.mark.core_model
140
def test_fp32_head_e2e_no_nan():
141
"""An fp32 head produces finite logprobs end-to-end.
vllm/lora/layers/logits_processor.py
+7
-12
@@ -45,15 +45,6 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
45
self.hidden_size = hidden_size
46
self.dtype = dtype
47
self.device = device
48
- # The fp32 lm_head path lives in the base LogitsProcessor._get_logits,
49
- # which this wrapper bypasses. Rather than silently emit model-dtype
50
- # logits, reject the combination until the LoRA path supports it.
51
- head_dtype = getattr(base_layer, "head_dtype", None)
52
- if head_dtype is not None and head_dtype != dtype:
53
- raise ValueError(
54
- "A head_dtype different from the model dtype (e.g. an fp32 "
55
- "lm_head) is not yet supported with LoRA."
56
- )
48
self.tp_size = get_tensor_model_parallel_world_size()
49
self.tp_rank = get_tensor_model_parallel_rank()
50
self.sharded_to_full_mapping = sharded_to_full_mapping
@@ -158,9 +149,13 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
149
actual_lm_head = lm_head.base_layer
150
else:
151
actual_lm_head = lm_head
161
- logits = actual_lm_head.quant_method.apply(actual_lm_head, hidden_states)
162
- if embedding_bias is not None:
163
- logits += embedding_bias
152
+ # Run the base projection through the LogitsProcessor so head_dtype
153
+ # (e.g. an fp32 lm_head) is honored on the LoRA path too. The LoRA
154
+ # delta is accumulated into these logits by add_lora_logits below,
155
+ # whose internal buffer is already fp32.
156
+ logits = self.base_layer._apply_head(
157
+ actual_lm_head, hidden_states, embedding_bias
158
+ )
159
160
# Gather logits for TP
161
logits = self.base_layer._gather_logits(logits)