[Core] Support fp32 lm_head for generation models via head_dtype (RFC #48305 §3.6) (#48390)
Signed-off-by: Karthik Kothuri <karthikkothuri2009@gmail.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Karthik Kothuri committed
Jul 13, 2026 at 01:43 UTC
107a03ba63e005ff03424fed9c4e6cf551b98bb2
6 files changed
+249
-15
tests/models/language/generation_ppl_test/ppl_utils.py
+1
@@ -126,3 +126,4 @@ def wikitext_ppl_test(
126
# We are not concerned that the vllm PPL is less than Transformers,
127
# so we only perform one-sided testing.
128
assert differ < atol
129
+ return vllm_ppl
tests/models/language/generation_ppl_test/test_gpt.py
+11
-1
@@ -11,4 +11,14 @@ MODELS = [GenerateModelInfo("openai-community/gpt2-large", hf_ppl=19.45705604553
11
12
@pytest.mark.parametrize("model_info", MODELS)
13
def test_ppl(hf_runner, vllm_runner, model_info: GenerateModelInfo):
14
- wikitext_ppl_test(hf_runner, vllm_runner, model_info)
14
+ bf16_ppl = wikitext_ppl_test(hf_runner, vllm_runner, model_info)
15
+ fp32_ppl = wikitext_ppl_test(
16
+ hf_runner,
17
+ vllm_runner,
18
+ model_info,
19
+ vllm_extra_kwargs={"hf_overrides": {"head_dtype": "float32"}},
20
+ )
21
+
22
+ differ = ((fp32_ppl - bf16_ppl) / bf16_ppl) * 100
23
+ print("fp32 head difference (%):", differ)
24
+ assert differ < 0
tests/v1/sample/test_head_dtype.py
new
+172
@@ -0,0 +1,172 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for running the generation lm_head in fp32 via ``head_dtype``.
4
+
5
+An fp32 head lets rollout logits match a trainer that computes the lm_head in
6
+fp32, which is required for RL training-inference consistency.
7
+"""
8
+
9
+import math
10
+
11
+import pytest
12
+import torch
13
+
14
+from vllm import LLM, SamplingParams
15
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
+from vllm.model_executor.layers.vocab_parallel_embedding import (
17
+ UnquantizedEmbeddingMethod,
18
+)
19
+
20
+
21
+class _FakeLmHead:
22
+ def __init__(
23
+ self,
24
+ weight: torch.Tensor,
25
+ quantized: bool = False,
26
+ shard_indices: object | None = None,
27
+ ):
28
+ self.weight = weight
29
+ self.quant_method = object() if quantized else UnquantizedEmbeddingMethod()
30
+ self.shard_indices = shard_indices
31
+
32
+
33
+def _build_processor(vocab_size: int) -> LogitsProcessor:
34
+ lp = LogitsProcessor(vocab_size)
35
+ # The TP gather is orthogonal to the dtype behavior under test.
36
+ lp._gather_logits = lambda logits: logits
37
+ return lp
38
+
39
+
40
+def test_fp32_head_runs_projection_in_fp32(default_vllm_config):
41
+ vocab_size, hidden_size, num_tokens = 64, 16, 4
42
+ lp = _build_processor(vocab_size)
43
+ lp.head_dtype = torch.float32
44
+
45
+ hidden_states = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16)
46
+ weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
47
+
48
+ logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)
49
+
50
+ assert logits.dtype == torch.float32
51
+ assert torch.isfinite(logits).all()
52
+ expected = torch.nn.functional.linear(hidden_states.float(), weight.float())
53
+ torch.testing.assert_close(logits, expected)
54
+
55
+
56
+def test_non_fp32_head_dtype_uses_cast_path(default_vllm_config):
57
+ # head_dtype != fp32 must not hit the CUDA out_dtype-mm fast path
58
+ # (torch.mm only supports fp32 out for fp16/bf16 inputs); the cast path
59
+ # handles any dtype.
60
+ vocab_size, hidden_size = 64, 16
61
+ lp = _build_processor(vocab_size)
62
+ lp.head_dtype = torch.float16
63
+
64
+ hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
65
+ weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
66
+
67
+ logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)
68
+
69
+ assert logits.dtype == torch.float16
70
+ expected = torch.nn.functional.linear(hidden_states.half(), weight.half())
71
+ torch.testing.assert_close(logits, expected)
72
+
73
+
74
+def test_head_dtype_equal_to_model_dtype_uses_quant_method(default_vllm_config):
75
+ vocab_size, hidden_size = 64, 16
76
+ lp = _build_processor(vocab_size)
77
+ lp.head_dtype = torch.bfloat16
78
+
79
+ hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
80
+ weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
81
+
82
+ logits = lp._get_logits(hidden_states, _FakeLmHead(weight), None)
83
+ assert logits.dtype == torch.bfloat16
84
+
85
+
86
+def test_fp32_head_rejects_quantized_lm_head(default_vllm_config):
87
+ lp = _build_processor(64)
88
+ lp.head_dtype = torch.float32
89
+ lm_head = _FakeLmHead(torch.randn(64, 16, dtype=torch.bfloat16), quantized=True)
90
+
91
+ with pytest.raises(ValueError, match="unquantized"):
92
+ lp._get_logits(torch.randn(4, 16, dtype=torch.bfloat16), lm_head, None)
93
+
94
+
95
+def test_get_top_tokens_honors_head_dtype(default_vllm_config):
96
+ # The spec-decode local-argmax path (get_top_tokens) must run the lm_head
97
+ # in head_dtype too, not just _get_logits.
98
+ import types
99
+ from unittest import mock
100
+
101
+ vocab_size, hidden_size = 64, 16
102
+ lp = _build_processor(vocab_size)
103
+ lp.head_dtype = torch.float32
104
+
105
+ hidden_states = torch.randn(4, hidden_size, dtype=torch.bfloat16)
106
+ weight = torch.randn(vocab_size, hidden_size, dtype=torch.bfloat16)
107
+ lm_head = _FakeLmHead(
108
+ weight,
109
+ shard_indices=types.SimpleNamespace(
110
+ num_org_vocab_padding=0, org_vocab_start_index=0
111
+ ),
112
+ )
113
+
114
+ with mock.patch(
115
+ "vllm.model_executor.layers.logits_processor."
116
+ "get_tensor_model_parallel_world_size",
117
+ return_value=1,
118
+ ):
119
+ top = lp.get_top_tokens(lm_head, hidden_states, None)
120
+
121
+ expected = torch.nn.functional.linear(hidden_states.float(), weight.float()).argmax(
122
+ dim=-1
123
+ )
124
+ assert torch.equal(top, expected)
125
+
126
+
127
+def test_fp32_head_rejected_with_lora(default_vllm_config):
128
+ from vllm.lora.layers.logits_processor import LogitsProcessorWithLoRA
129
+
130
+ base = _build_processor(64)
131
+ base.head_dtype = torch.float32
132
+
133
+ with pytest.raises(ValueError, match="not yet supported with LoRA"):
134
+ LogitsProcessorWithLoRA(
135
+ base,
136
+ hidden_size=16,
137
+ dtype=torch.bfloat16,
138
+ device=torch.device("cpu"),
139
+ sharded_to_full_mapping=None,
140
+ )
141
+
142
+
143
+@pytest.mark.core_model
144
+def test_fp32_head_e2e_no_nan():
145
+ """An fp32 head produces finite logprobs end-to-end.
146
+
147
+ Runs on the default (v2) model runner and exercises the
148
+ processed_logprobs path, which forces the native sampler and is where a
149
+ non-contiguous fp32 logits row previously produced NaN.
150
+ """
151
+ llm = LLM(
152
+ model="facebook/opt-125m",
153
+ hf_overrides={"head_dtype": "float32"},
154
+ logprobs_mode="processed_logprobs",
155
+ enforce_eager=True,
156
+ gpu_memory_utilization=0.5,
157
+ max_model_len=256,
158
+ )
159
+ sampling_params = SamplingParams(
160
+ temperature=1.0, top_p=0.95, top_k=50, max_tokens=32, logprobs=5, seed=0
161
+ )
162
+ outputs = llm.generate(
163
+ ["The capital of France is", "Once upon a time,"], sampling_params
164
+ )
165
+
166
+ for output in outputs:
167
+ for completion in output.outputs:
168
+ for token_id, position in zip(completion.token_ids, completion.logprobs):
169
+ # The sampled token survived filtering, so its logprob is finite.
170
+ assert math.isfinite(position[token_id].logprob)
171
+ # No returned logprob is NaN.
172
+ assert not any(math.isnan(lp.logprob) for lp in position.values())
vllm/config/model.py
+5
-11
@@ -1709,24 +1709,18 @@ class ModelConfig:
1709
such as the lm_head in a generation model,
1710
or the score or classifier in a classification model.
1711
1712
- `head_dtype` currently only supports pooling models.
1713
-
1714
- - The pooling model defaults to using fp32 head, you can use
1712
+ - Pooling models default to an fp32 head; use
1713
--hf-overrides '{"head_dtype": "model"}' to disable it.
1714
+ - Generation models default to the model dtype; set
1715
+ --hf-overrides '{"head_dtype": "float32"}' to run the lm_head in
1716
+ fp32, which is required for RL training-inference consistency
1717
+ (the trainer computes logits in fp32).
1718
"""
1719
1720
head_dtype = _get_head_dtype(
1721
config=self.hf_config, dtype=self.dtype, runner_type=self.runner_type
1722
)
1723
1722
- if self.runner_type != "pooling" and head_dtype != self.dtype:
1723
- logger.warning_once(
1724
- "`head_dtype` currently only supports pooling models, "
1725
- "fallback to model dtype [%s].",
1726
- self.dtype,
1727
- )
1728
- return self.dtype
1729
-
1724
if head_dtype not in current_platform.supported_dtypes:
1725
logger.warning_once(
1726
"The current platform does not support [%s] head dtype, "
vllm/lora/layers/logits_processor.py
+9
@@ -45,6 +45,15 @@ 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
+ )
57
self.tp_size = get_tensor_model_parallel_world_size()
58
self.tp_rank = get_tensor_model_parallel_rank()
59
self.sharded_to_full_mapping = sharded_to_full_mapping
vllm/model_executor/layers/logits_processor.py
+51
-3
@@ -3,14 +3,19 @@
3
"""A layer that compute logits from hidden_stats."""
4
5
import torch
6
+import torch.nn.functional as F
7
8
+from vllm.config import get_current_vllm_config
9
from vllm.distributed import (
10
get_tensor_model_parallel_world_size,
11
tensor_model_parallel_all_gather,
12
tensor_model_parallel_gather,
13
)
14
from vllm.model_executor.custom_op import PluggableLayer
13
-from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
15
+from vllm.model_executor.layers.vocab_parallel_embedding import (
16
+ UnquantizedEmbeddingMethod,
17
+ VocabParallelEmbedding,
18
+)
19
from vllm.platforms import current_platform
20
21
@@ -50,6 +55,11 @@ class LogitsProcessor(PluggableLayer):
55
self.soft_cap = soft_cap
56
# Whether to use gather or all-gather to gather the logits.
57
self.use_all_gather = current_platform.use_all_gather()
58
+ # Dtype of the lm_head projection. Defaults to the model dtype; an
59
+ # fp32 head (via `--hf-overrides '{"head_dtype": "float32"}'`) is
60
+ # required for RL training-inference consistency.
61
+ model_config = get_current_vllm_config().model_config
62
+ self.head_dtype = model_config.head_dtype if model_config is not None else None
63
64
def forward(
65
self,
@@ -86,6 +96,44 @@ class LogitsProcessor(PluggableLayer):
96
logits = tensor_model_parallel_gather(logits)
97
return logits
98
99
+ def _apply_head(
100
+ self,
101
+ lm_head: VocabParallelEmbedding,
102
+ hidden_states: torch.Tensor,
103
+ embedding_bias: torch.Tensor | None,
104
+ ) -> torch.Tensor:
105
+ """Project hidden states through the lm_head, honoring head_dtype."""
106
+ if self.head_dtype is None or self.head_dtype == hidden_states.dtype:
107
+ return lm_head.quant_method.apply(
108
+ lm_head, hidden_states, bias=embedding_bias
109
+ )
110
+
111
+ if not isinstance(lm_head.quant_method, UnquantizedEmbeddingMethod):
112
+ raise ValueError(
113
+ "A head_dtype different from the model dtype is only "
114
+ "supported for an unquantized lm_head."
115
+ )
116
+ if (
117
+ self.head_dtype == torch.float32
118
+ and current_platform.is_cuda()
119
+ and hidden_states.is_cuda
120
+ ):
121
+ # Accumulate the projection directly into fp32. This avoids
122
+ # materializing an fp32 copy of the lm_head weight on every step,
123
+ # unlike casting both operands. `torch.mm(out_dtype=...)` is
124
+ # CUDA-only and only supports fp32 output for fp16/bf16 inputs, so
125
+ # other cases fall back to the cast path below.
126
+ flat = hidden_states.reshape(-1, hidden_states.shape[-1])
127
+ logits = torch.mm(flat, lm_head.weight.t(), out_dtype=self.head_dtype)
128
+ if embedding_bias is not None:
129
+ logits = logits + embedding_bias.to(self.head_dtype)
130
+ return logits.reshape(*hidden_states.shape[:-1], -1)
131
+ return F.linear(
132
+ hidden_states.to(self.head_dtype),
133
+ lm_head.weight.to(self.head_dtype),
134
+ embedding_bias.to(self.head_dtype) if embedding_bias is not None else None,
135
+ )
136
+
137
def _get_logits(
138
self,
139
hidden_states: torch.Tensor,
@@ -93,7 +141,7 @@ class LogitsProcessor(PluggableLayer):
141
embedding_bias: torch.Tensor | None,
142
) -> torch.Tensor | None:
143
# Get the logits for the next tokens.
96
- logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)
144
+ logits = self._apply_head(lm_head, hidden_states, embedding_bias)
145
146
# Gather logits for TP
147
logits = self._gather_logits(logits)
@@ -122,7 +170,7 @@ class LogitsProcessor(PluggableLayer):
170
)
171
tp_size = get_tensor_model_parallel_world_size()
172
125
- logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)
173
+ logits = self._apply_head(lm_head, hidden_states, embedding_bias)
174
if self.soft_cap is not None:
175
logits = torch.tanh(logits / self.soft_cap) * self.soft_cap
176
if self.scale != 1.0: