diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 13534230bad9..31b6d06581da 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -242,6 +242,10 @@ class ModelConfig(Generic[TConfig]): # draft models so greedy draft sampling can do a lighter TP gather. Defaults # to True to preserve behavior for every non-draft model. lm_head_gather_output: bool = True + # Dtype of the logits written by the LM head GEMM (LlmArgs.lm_head_output_dtype). + # None keeps the model dtype; torch.float32 writes the float32 accumulator + # directly instead of rounding the logits to bf16/fp16 first. + lm_head_output_dtype: Optional[torch.dtype] = None lora_config: Optional["LoraConfig"] = None sparse_attention_config: Optional["SparseAttentionConfig"] = None kv_cache_compression_config: Optional["KvCacheCompressionConfig"] = None diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 0a6060b47288..9078559f8b61 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -580,6 +580,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], hidden_size, dtype=config.pretrained_config.torch_dtype, quant_config=lm_head_quant_config, + output_dtype=config.lm_head_output_dtype, ) else: if (hasattr(config, 'lora_config') @@ -609,6 +610,7 @@ def __init__(self, model: TModel, *, config: ModelConfig[TConfig], use_custom_cublas_mm=getattr(model, 'use_custom_cublas_mm', False), quant_config=lm_head_quant_config, + output_dtype=config.lm_head_output_dtype, ) if self.has_custom_lm_head: diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py index a707372590a2..e527f8aab629 100644 --- a/tensorrt_llm/_torch/modules/embedding.py +++ b/tensorrt_llm/_torch/modules/embedding.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import math from typing import Dict, List, Optional, Tuple @@ -14,6 +17,23 @@ from .linear import Linear, TensorParallelMode +def _matmul_with_output_dtype(input: torch.Tensor, weight: torch.Tensor, + output_dtype: torch.dtype | None) -> torch.Tensor: + """Compute ``input @ weight.T`` with an optional output dtype override. + + None keeps the model dtype. An FP32 override stores the result directly + in FP32 without first rounding it to the BF16/FP16 model dtype. + """ + out_features = weight.shape[0] + output_shape = (*input.shape[:-1], out_features) + input_2d = input.reshape(-1, input.shape[-1]) + mm_kwargs = {} + if output_dtype is not None and output_dtype != weight.dtype: + mm_kwargs["out_dtype"] = output_dtype + output = torch.mm(input_2d, weight.t(), **mm_kwargs) + return output.view(output_shape) + + class LMHead(Linear): """LM head layer. @@ -23,6 +43,11 @@ class LMHead(Linear): dtype (Optional[torch.dtype]): type of the parameters. mapping (Optional[Mapping]): parallelism configuration. If not provided, the embedding is not parallelized. + output_dtype (Optional[torch.dtype]): dtype of the logits written by + the GEMM. None keeps the parameter dtype. torch.float32 keeps the + weights and inputs in the parameter dtype but writes the float32 + accumulator directly, so logits are not rounded to bf16/fp16. + Only supported for unquantized, non-ROW-parallel heads. """ def __init__( @@ -36,6 +61,7 @@ def __init__( reduce_output: bool = True, use_custom_cublas_mm: bool = False, quant_config=None, + output_dtype: Optional[torch.dtype] = None, ): local_in_features = embedding_dim local_out_features = num_embeddings @@ -119,6 +145,22 @@ def __init__( self.weight = Parameter(torch.empty(weight_shape, dtype=dtype)) self.register_parameter("bias", None) + if output_dtype is not None: + if self.has_any_quant: + raise NotImplementedError( + "LMHead output_dtype is only supported for unquantized " + f"LM heads (quant_algo={quant_config.quant_algo})") + if self.tp_mode == TensorParallelMode.ROW: + raise NotImplementedError( + "LMHead output_dtype does not support ROW tensor-parallel " + "mode") + self.output_dtype = output_dtype + + def _writes_output_dtype(self) -> bool: + """Whether the GEMM must store a dtype different from its weights.""" + return (self.output_dtype is not None + and self.output_dtype != self.weight.dtype) + @property def vocab_size_padded(self) -> int: if self.tp_mode == TensorParallelMode.COLUMN and self.gather_output: @@ -141,8 +183,8 @@ def forward( slice_width = ceil_div(self.out_features, tp_size) slice_start = tp_rank * slice_width slice_end = min((tp_rank + 1) * slice_width, self.out_features) - output = F.linear(input, self.weight[slice_start:slice_end, :], - None) + weight = self.weight[slice_start:slice_end, :] + output = _matmul_with_output_dtype(input, weight, self.output_dtype) else: output = super().forward(input, all_reduce_params=all_reduce_params) if (self.tp_mode == TensorParallelMode.COLUMN and self.gather_output @@ -151,6 +193,21 @@ def forward( return output + def apply_linear(self, + input: torch.Tensor, + bias: Optional[torch.Tensor], + lora_params: Optional[dict] = None, + layer_idx: Optional[int] = None) -> torch.Tensor: + """Project hidden states in the configured GEMM output dtype.""" + # Linear.forward runs the GEMM through apply_linear for the COLUMN and + # non-parallel modes, so the TP all-gather and the vocab-padding slice + # in forward() apply unchanged to the output_dtype logits. + if self._writes_output_dtype(): + assert bias is None, "LMHead has no bias" + return _matmul_with_output_dtype(input, self.weight, + self.output_dtype) + return super().apply_linear(input, bias, lora_params, layer_idx) + def skip_forward( self, input: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 9f2f4e12a855..78c98838c96a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -51,6 +51,7 @@ from ..models.checkpoints.weight_load_plan import WeightLoadPlan from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode, get_registered_model_class, timing_metric) +from ..modules.embedding import Embedding, LMHead from ..modules.low_m_gemm import LOW_M_GEMM_ACTIVE, prepare_low_m_gemm from ..moe.fused_moe.moe_load_balancer import (MoeLoadBalancer, maybe_create_moe_load_balancer) @@ -296,6 +297,31 @@ def validate_encoder_decoder_tp_scope(model_config: ModelConfig) -> None: f"attn_backend='{model_config.attn_backend}'.") +def validate_lm_head_output_dtype(model: torch.nn.Module, + requested_dtype: str) -> None: + """Reject an explicit output dtype that the constructed LM heads ignore.""" + if requested_dtype == "auto": + return + + expected_dtype = str_dtype_to_torch(requested_dtype) + found_lm_head = False + for name, module in model.named_modules(): + # Embedding inherits LMHead but does not project hidden states to logits. + if not isinstance(module, LMHead) or isinstance(module, Embedding): + continue + found_lm_head = True + if module.output_dtype != expected_dtype: + raise ValueError( + f"lm_head_output_dtype={requested_dtype!r} is not supported " + f"by {type(model).__name__}: {name or ''}.output_dtype " + f"is {module.output_dtype}, expected {expected_dtype}.") + + if not found_lm_head: + raise ValueError( + f"lm_head_output_dtype={requested_dtype!r} is not supported by " + f"{type(model).__name__}: no LMHead was constructed.") + + def initialize_dummy_weights( model: torch.nn.Module, low: float = -1e-3, @@ -974,6 +1000,9 @@ def load( model = AutoModelForCausalLM.from_config(config) is_meta_init = False + validate_lm_head_output_dtype(model, + self.llm_args.lm_head_output_dtype) + requires_standard_hf_loading = any( getattr(module, "_requires_standard_hf_loading", False) for module in model.modules()) @@ -1946,6 +1975,9 @@ def _load_and_validate_config( multimodal_config=self.llm_args.multimodal_config, use_cute_dsl_bf16_bmm=self.llm_args.use_cute_dsl_bf16_bmm, use_cute_dsl_bf16_gemm=self.llm_args.use_cute_dsl_bf16_gemm, + lm_head_output_dtype=(torch.float32 + if self.llm_args.lm_head_output_dtype + == "float32" else None), locality_domain_policy=LocalityDomainPolicy( enabled=self.llm_args.enable_locality_domains), ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 01f76b2886f9..5515ae8e0a3f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6280,6 +6280,18 @@ def init_multimodal_config(cls, v): "Separated-routing MoE backends only.", status="beta") + lm_head_output_dtype: Literal["auto", "float32"] = Field( + default="auto", + description= + "Output dtype written by the LM head GEMM. 'auto' uses the model " + "dtype before logits are upcast to float32. 'float32' writes the " + "float32 accumulator directly, avoiding intermediate bf16/fp16 " + "rounding of logits before logprob computation. Weights and inputs " + "remain in the model dtype. Only supported for unquantized LM heads " + "of decoder-only models.", + status="prototype", + json_schema_extra={"type": "Literal['auto', 'float32']"}) + moe_config: MoeConfig = Field(default_factory=MoeConfig, description="MoE config.", status="beta") diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f6c872897b8f..65fee0aebc51 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -866,6 +866,15 @@ "kind": "categorical", "path": "layer_wise_benchmarks_config.calibration_mode" }, + { + "allowed_values": [ + "auto", + "float32" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "lm_head_output_dtype" + }, { "allowed_values": [ "auto", @@ -1379,6 +1388,11 @@ "kind": "value", "path": "sparse_attention_config.fuse_qkv_index_projection" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "sparse_attention_config.gvr_deterministic_topk" + }, { "allowed_values": [ "triton", diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 3376a66186c3..f22eec9bb718 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -41,6 +41,7 @@ l0_h100: - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules - unittest/_torch/modules/test_w4a16_nvfp4_linear.py + - unittest/_torch/modules/test_lm_head_output_dtype.py - unittest/_torch/modules/test_multi_stream_utils.py # ------------- MoE components tests --------------- - unittest/_torch/moe/test_moe_load_balancer.py @@ -148,6 +149,7 @@ l0_h100: - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[PYTHON-mla-ctx_fp16_gen_fp16] - unittest/llmapi/test_llm_telemetry.py - unittest/llmapi/test_return_routed_experts.py + - unittest/llmapi/test_lm_head_dtype.py - unittest/scaffolding - unittest/usage/test_transport.py - unittest/usage/test_e2e_capture.py diff --git a/tests/unittest/_torch/executor/test_model_loader_gms.py b/tests/unittest/_torch/executor/test_model_loader_gms.py index 551edab14244..5e16c03a904c 100644 --- a/tests/unittest/_torch/executor/test_model_loader_gms.py +++ b/tests/unittest/_torch/executor/test_model_loader_gms.py @@ -94,6 +94,7 @@ def _pool_scope(events): def _make_loader(monkeypatch, *, events, spec_config=None): llm_args = SimpleNamespace( load_format=LoadFormat.GMS, + lm_head_output_dtype="auto", gms_config=SimpleNamespace(socket_path="/tmp/gms.sock", mode="auto", tag="weights"), ) loader = ModelLoader( diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 630f4598f814..f86e4601c8eb 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -547,7 +547,7 @@ def _tiny_profile_registry(*, speculative_mode: str | None = None) -> PostTransf def _make_loader(monkeypatch, *, events, spec_config=None): - llm_args = SimpleNamespace(load_format=LoadFormat.AUTO) + llm_args = SimpleNamespace(load_format=LoadFormat.AUTO, lm_head_output_dtype="auto") loader = ModelLoader( llm_args=llm_args, mapping=MagicMock(name="mapping"), diff --git a/tests/unittest/_torch/modules/test_lm_head_output_dtype.py b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py new file mode 100644 index 000000000000..16fefac631ff --- /dev/null +++ b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py @@ -0,0 +1,300 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""LMHead(output_dtype=torch.float32), which backs LlmArgs.lm_head_output_dtype="float32". + +The weights and inputs stay bf16; the GEMM writes its float32 accumulator instead of +rounding the logits to bf16. The float64 reference uses the same bf16 values, so the +float32 path should match it up to accumulation order, while the default path is off +by the bf16 rounding of each logit. +""" + +from contextlib import nullcontext +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +import torch +from transformers import LlamaConfig + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.embedding import Embedding, LMHead +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_module +from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader, validate_lm_head_output_dtype +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +VOCAB_SIZE = 32000 +HIDDEN_SIZE = 1024 +DTYPE = torch.bfloat16 + + +def _make_inputs(num_tokens: int) -> tuple[torch.Tensor, torch.Tensor]: + """Generate BF16 inputs that expose final-output rounding.""" + torch.manual_seed(0) + hidden = torch.randn(num_tokens, HIDDEN_SIZE, dtype=DTYPE) + # Logits with a std of about 4, where bf16 rounding is on the order of 1e-2. + weight = (torch.randn(VOCAB_SIZE, HIDDEN_SIZE) * 4 / HIDDEN_SIZE**0.5).to(DTYPE) + return hidden.cuda(), weight + + +def _make_lm_head(weight: torch.Tensor, output_dtype: torch.dtype | None, **kwargs) -> LMHead: + """Load the same unquantized weights into either output-dtype path.""" + lm_head = LMHead(VOCAB_SIZE, HIDDEN_SIZE, dtype=DTYPE, output_dtype=output_dtype, **kwargs) + lm_head.load_weights([{"weight": weight}]) + return lm_head.cuda() + + +def _reference_logits(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Compute an independent FP64 reference from the same BF16 inputs.""" + return hidden.double() @ weight.cuda().double().t() + + +@pytest.mark.parametrize( + "tp_mode", [None, TensorParallelMode.COLUMN], ids=["no_tp", "column_gather"] +) +@pytest.mark.parametrize("num_tokens", [0, 1, 7, 128]) +def test_float32_output_matches_reference( + num_tokens: int, tp_mode: TensorParallelMode | None +) -> None: + """Avoid BF16 output rounding for empty, single-token, and batched inputs.""" + hidden, weight = _make_inputs(num_tokens) + kwargs = dict( + mapping=Mapping(), tensor_parallel_mode=tp_mode, gather_output=True, reduce_output=False + ) + logits_bf16 = _make_lm_head(weight, None, **kwargs)(hidden) + logits_fp32 = _make_lm_head(weight, torch.float32, **kwargs)(hidden) + reference = _reference_logits(hidden, weight) + + assert logits_bf16.dtype == DTYPE + assert logits_fp32.dtype == torch.float32 + assert logits_fp32.shape == logits_bf16.shape == reference.shape + if num_tokens == 0: + return + + torch.testing.assert_close(logits_fp32.double(), reference, rtol=1e-4, atol=1e-3) + # The logits are no longer bf16 values ... + assert not torch.equal(logits_fp32, logits_fp32.to(DTYPE).float()) + # ... and the error left is far below the bf16 rounding error, for logits and logprobs. + err_bf16 = (logits_bf16.double() - reference).abs().max() + err_fp32 = (logits_fp32.double() - reference).abs().max() + assert err_fp32 * 50 < err_bf16 + logprobs_ref = torch.log_softmax(reference, dim=-1) + lp_err_bf16 = (torch.log_softmax(logits_bf16.double(), dim=-1) - logprobs_ref).abs().max() + lp_err_fp32 = (torch.log_softmax(logits_fp32.double(), dim=-1) - logprobs_ref).abs().max() + assert lp_err_fp32 * 10 < lp_err_bf16 + + +def test_float32_output_one_dim_input() -> None: + """Accept the one-dimensional hidden state used without attention metadata.""" + # LogitsProcessor passes hidden_states[-1] (1-D) when there is no attention metadata. + hidden, weight = _make_inputs(1) + logits = _make_lm_head(weight, torch.float32)(hidden[0]) + assert logits.shape == (VOCAB_SIZE,) + assert logits.dtype == torch.float32 + torch.testing.assert_close( + logits.double(), _reference_logits(hidden, weight)[0], rtol=1e-4, atol=1e-3 + ) + + +@pytest.mark.parametrize("output_dtype", [None, DTYPE, torch.float32], ids=["auto", "bf16", "fp32"]) +@pytest.mark.parametrize("num_tokens", [0, 1, 16]) +def test_output_dtype_spec_decoding_head_slice( + num_tokens: int, output_dtype: torch.dtype | None +) -> None: + """Preserve default and FP32 output precision for a sliced speculative head.""" + # With lm_head TP in attention DP, the spec-decoding head slices the vocab at forward time. + hidden, weight = _make_inputs(num_tokens) + mapping = Mapping(enable_attention_dp=True, enable_lm_head_tp_in_adp=True) + lm_head = _make_lm_head( + weight, + output_dtype, + mapping=mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + reduce_output=False, + ) + lm_head_tp = Mapping(world_size=2, tp_size=2, rank=1) + logits = lm_head(hidden, mapping_lm_head_tp=lm_head_tp, is_spec_decoding_head=True) + reference = _reference_logits(hidden, weight)[:, VOCAB_SIZE // 2 :] + assert logits.dtype == (output_dtype or DTYPE) + assert logits.shape == reference.shape + if output_dtype == torch.float32: + torch.testing.assert_close(logits.double(), reference, rtol=1e-4, atol=1e-3) + else: + torch.testing.assert_close(logits, reference.to(DTYPE)) + + +def test_output_dtype_rejects_quantized_head() -> None: + """Reject the FP32-output option for unsupported quantized weights.""" + with pytest.raises(NotImplementedError, match="unquantized"): + LMHead( + VOCAB_SIZE, + HIDDEN_SIZE, + dtype=DTYPE, + quant_config=QuantConfig(quant_algo=QuantAlgo.FP8), + output_dtype=torch.float32, + ) + + +def test_output_dtype_rejects_row_parallel() -> None: + """Reject ROW parallelism rather than silently changing reduction precision.""" + with pytest.raises(NotImplementedError, match="ROW"): + LMHead( + VOCAB_SIZE, + HIDDEN_SIZE, + dtype=DTYPE, + tensor_parallel_mode=TensorParallelMode.ROW, + output_dtype=torch.float32, + ) + + +def test_llm_args_lm_head_output_dtype(tmp_path: Path) -> None: + """Validate the public option's default, FP32 value, and unsupported values.""" + assert TorchLlmArgs(model=str(tmp_path)).lm_head_output_dtype == "auto" + assert ( + TorchLlmArgs(model=str(tmp_path), lm_head_output_dtype="float32").lm_head_output_dtype + == "float32" + ) + with pytest.raises(ValueError): + TorchLlmArgs(model=str(tmp_path), lm_head_output_dtype="float16") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +@torch.inference_mode() +def test_float32_output_strided_batched_input(dtype: torch.dtype) -> None: + """Preserve leading dimensions and FP32 precision for strided hidden states.""" + generator = torch.Generator(device="cuda").manual_seed(42) + hidden = torch.randn(2, 3, 128, device="cuda", dtype=dtype, generator=generator)[..., ::2] + head = LMHead(127, 64, dtype=dtype, output_dtype=torch.float32).cuda() + head.weight.copy_(torch.randn(127, 64, device="cuda", dtype=dtype, generator=generator)) + assert not hidden.is_contiguous() + + logits = head(hidden) + reference = hidden.double() @ head.weight.double().t() + assert logits.shape == (2, 3, 127) + assert logits.dtype == torch.float32 + torch.testing.assert_close(logits.double(), reference, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +@torch.inference_mode() +def test_float32_output_cuda_graph(dtype: torch.dtype) -> None: + """Replay the native FP32-output GEMM with updated hidden states.""" + generator = torch.Generator(device="cuda").manual_seed(42) + hidden = torch.randn(7, 64, device="cuda", dtype=dtype, generator=generator) + head = LMHead(127, 64, dtype=dtype, output_dtype=torch.float32).cuda() + head.weight.copy_(torch.randn(127, 64, device="cuda", dtype=dtype, generator=generator)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + head(hidden) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + logits = head(hidden) + + for _ in range(2): + hidden.copy_(torch.randn(hidden.shape, device="cuda", dtype=dtype, generator=generator)) + graph.replay() + reference = hidden.double() @ head.weight.double().t() + assert logits.dtype == torch.float32 + torch.testing.assert_close(logits.double(), reference, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("nested", [False, True], ids=["direct", "nested"]) +@pytest.mark.parametrize("output_dtype", [None, torch.bfloat16, torch.float32]) +def test_built_model_lm_head_output_dtype(nested: bool, output_dtype: torch.dtype | None) -> None: + """Validate real LM heads, including nested ones, without mistaking embeddings for heads.""" + model = torch.nn.Module() + model.embed_tokens = Embedding(16, 8, dtype=DTYPE) + owner = model + if nested: + model.language_model = torch.nn.Module() + owner = model.language_model + owner.lm_head = LMHead(16, 8, dtype=DTYPE, output_dtype=output_dtype) + + validate_lm_head_output_dtype(model, "auto") + if output_dtype == torch.float32: + validate_lm_head_output_dtype(model, "float32") + else: + with pytest.raises(ValueError, match=r"lm_head_output_dtype=.*lm_head.output_dtype"): + validate_lm_head_output_dtype(model, "float32") + + +def test_built_model_lm_head_output_dtype_checks_all_heads() -> None: + """Reject an ignored option in a nested head even if another head honours it.""" + model = torch.nn.Module() + model.lm_head = LMHead(16, 8, dtype=DTYPE, output_dtype=torch.float32) + model.draft = torch.nn.Module() + model.draft.lm_head = LMHead(16, 8, dtype=DTYPE) + with pytest.raises(ValueError, match=r"draft.lm_head.output_dtype"): + validate_lm_head_output_dtype(model, "float32") + + +def test_built_model_lm_head_output_dtype_requires_a_head() -> None: + """Reject an explicit output dtype when only input embeddings were constructed.""" + model = torch.nn.Module() + model.embed_tokens = Embedding(16, 8, dtype=DTYPE) + validate_lm_head_output_dtype(model, "auto") + with pytest.raises(ValueError, match="no LMHead was constructed"): + validate_lm_head_output_dtype(model, "float32") + + +@pytest.mark.parametrize("nested", [False, True], ids=["direct", "nested"]) +@pytest.mark.parametrize("fallback", [False, True], ids=["meta_init", "regular_init"]) +def test_model_loader_rejects_ignored_lm_head_output_dtype( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, nested: bool, fallback: bool +) -> None: + """Reject unsupported constructed heads before any checkpoint weights are loaded.""" + model = torch.nn.Module() + owner = model + if nested: + model.language_model = torch.nn.Module() + owner = model.language_model + owner.lm_head = LMHead(16, 8, dtype=DTYPE) + config = ModelConfig( + pretrained_config=LlamaConfig(architectures=["LlamaForCausalLM"], dtype=DTYPE), + lm_head_output_dtype=torch.float32, + ) + loader = ModelLoader( + llm_args=TorchLlmArgs(model=str(tmp_path), lm_head_output_dtype="float32"), + mapping=Mapping(), + spec_config=None, + sparse_attention_config=None, + max_num_tokens=8, + max_seq_len=8, + ) + monkeypatch.setattr(loader, "_load_and_validate_config", lambda *_args: config) + monkeypatch.setattr(model_loader_module, "MetaInitMode", nullcontext) + monkeypatch.setattr( + model_loader_module, "maybe_create_moe_load_balancer", lambda *_args: nullcontext() + ) + monkeypatch.setattr(model_loader_module, "_checkpoint_startup_metadata", lambda *_args: {}) + factory = MagicMock( + side_effect=[RuntimeError("meta initialization unavailable"), model] + if fallback + else [model] + ) + monkeypatch.setattr(model_loader_module.AutoModelForCausalLM, "from_config", factory) + checkpoint_loader = MagicMock() + + with pytest.raises(ValueError, match=r"lm_head_output_dtype=.*lm_head.output_dtype"): + loader.load(str(tmp_path), checkpoint_loader) + + assert factory.call_count == (2 if fallback else 1) + checkpoint_loader.load_weights.assert_not_called() diff --git a/tests/unittest/_torch/multi_gpu/test_lm_head_output_dtype_tp.py b/tests/unittest/_torch/multi_gpu/test_lm_head_output_dtype_tp.py new file mode 100644 index 000000000000..eb0bc72a6420 --- /dev/null +++ b/tests/unittest/_torch/multi_gpu/test_lm_head_output_dtype_tp.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Verify FP32 LM-head logits through a real tensor-parallel all-gather.""" + +import pickle +import sys + +import cloudpickle +import pytest +import torch +from mpi4py import MPI +from mpi4py.futures import MPIPoolExecutor + +import tensorrt_llm +from tensorrt_llm._torch.modules.embedding import LMHead +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm.mapping import Mapping + +cloudpickle.register_pickle_by_value(sys.modules[__name__]) +MPI.pickle.__init__(cloudpickle.dumps, cloudpickle.loads, pickle.HIGHEST_PROTOCOL) + +# The shared MPI executor keeps its worker thread alive between tests. +pytestmark = pytest.mark.threadleak(enabled=False) + + +@torch.inference_mode() +def _run_column_lm_head( + tensor_parallel_size: int, hidden: torch.Tensor, weight: torch.Tensor +) -> bool: + """Check the gathered vocabulary and FP32 precision on every MPI rank.""" + assert tensorrt_llm.mpi_world_size() == tensor_parallel_size + rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(rank) + vocab_size, hidden_size = weight.shape + lm_head = LMHead( + num_embeddings=vocab_size, + embedding_dim=hidden_size, + dtype=weight.dtype, + mapping=Mapping(world_size=tensor_parallel_size, tp_size=tensor_parallel_size, rank=rank), + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + output_dtype=torch.float32, + ) + lm_head.load_weights([{"weight": weight}]) + lm_head.cuda() + hidden = hidden.cuda() + + logits = lm_head(hidden) + reference = hidden.double() @ weight.cuda().double().t() + + assert lm_head.weight.dtype == hidden.dtype == weight.dtype + assert logits.dtype == torch.float32 + assert logits.shape == (hidden.shape[0], vocab_size) + torch.testing.assert_close(logits.double(), reference, rtol=1e-4, atol=1e-4) + assert not torch.equal(logits, logits.to(weight.dtype).float()) + return True + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two GPUs") +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +@pytest.mark.parametrize("vocab_size", [128, 127], ids=["balanced", "padded"]) +@pytest.mark.parametrize("num_tokens", [1, 7]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +def test_lm_head_float32_column_gather( + mpi_pool_executor: MPIPoolExecutor, + vocab_size: int, + num_tokens: int, + dtype: torch.dtype, +) -> None: + """Preserve FP32 logits and vocabulary order across two padded or balanced shards.""" + hidden_size = 128 + generator = torch.Generator().manual_seed(42) + hidden = torch.randn(num_tokens, hidden_size, generator=generator, dtype=dtype) + weight = (torch.randn(vocab_size, hidden_size, generator=generator) * 4 / hidden_size**0.5).to( + dtype + ) + tensor_parallel_size = mpi_pool_executor.num_workers + assert tensor_parallel_size == 2 + arguments = [(tensor_parallel_size, hidden, weight)] * tensor_parallel_size + results = mpi_pool_executor.map(_run_column_lm_head, *zip(*arguments)) + assert list(results) == [True] * tensor_parallel_size diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index d2d1739473ff..e3e0e82ddcaa 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -182,6 +182,10 @@ methods: annotation: bool default: False status: beta + lm_head_output_dtype: + annotation: Literal['auto', 'float32'] + default: auto + status: prototype moe_config: annotation: tensorrt_llm.llmapi.llm_args.MoeConfig status: beta diff --git a/tests/unittest/llmapi/test_lm_head_dtype.py b/tests/unittest/llmapi/test_lm_head_dtype.py new file mode 100644 index 000000000000..e63b647d6ab0 --- /dev/null +++ b/tests/unittest/llmapi/test_lm_head_dtype.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end test for LlmArgs.lm_head_output_dtype. + +With the default LM head, logits are rounded to the model dtype before the float32 +upcast, so every returned logit is exactly representable in the model dtype. With +lm_head_output_dtype="float32" the LM head GEMM writes its float32 accumulator, so the +returned logits are not rounded, and greedy decoding still follows them. +""" + +import json +import os +from types import SimpleNamespace + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM +from tensorrt_llm._torch.pyexecutor.config_utils import resolve_hf_torch_dtype +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.sampling_params import SamplingParams + +_MODEL = os.path.join("llama-models-v2", "TinyLlama-1.1B-Chat-v1.0") + + +def _model_dir() -> str: + """Locate the CI checkpoint or skip when its weights are unavailable.""" + root = llm_models_root() + path = os.path.join(root, _MODEL) if root else None + if not path or not os.path.isdir(path): + pytest.skip(f"{_MODEL} not available under llm_models_root()") + return path + + +def _rounded_to(logits: torch.Tensor, dtype: torch.dtype) -> bool: + """Check whether all logits remain exactly representable in the model dtype.""" + return torch.equal(logits, logits.to(dtype).float()) + + +@pytest.mark.parametrize("lm_head_output_dtype", ["auto", "float32"]) +def test_lm_head_output_dtype_returned_logits(lm_head_output_dtype: str) -> None: + """Check context and decode precision through the public LLM API.""" + model_dir = _model_dir() + with open(os.path.join(model_dir, "config.json")) as f: + model_config = SimpleNamespace(**json.load(f)) + model_dtype = resolve_hf_torch_dtype(model_config) or torch.bfloat16 + + sampling_params = SamplingParams( + max_tokens=8, + temperature=0.0, + return_context_logits=True, + return_generation_logits=True, + ) + with LLM( + model=model_dir, + lm_head_output_dtype=lm_head_output_dtype, + max_batch_size=8, + # Context logits are not recomputed for reused KV cache blocks. + kv_cache_config=KvCacheConfig(max_tokens=10000, enable_block_reuse=False), + ) as llm: + output = llm.generate(["The capital of France is"], sampling_params)[0] + + completion = output.outputs[0] + context_logits = output.context_logits + generation_logits = completion.generation_logits + assert context_logits.dtype == torch.float32 + assert generation_logits.dtype == torch.float32 + assert context_logits.shape[0] == len(output.prompt_token_ids) + assert generation_logits.shape[0] == len(completion.token_ids) + + rounded = lm_head_output_dtype == "auto" + assert _rounded_to(context_logits, model_dtype) == rounded + assert _rounded_to(generation_logits, model_dtype) == rounded + assert torch.argmax(generation_logits, dim=-1).tolist() == completion.token_ids diff --git a/tests/unittest/scripts/test_bolt.py b/tests/unittest/scripts/test_bolt.py index e53d8f7ab2d8..e16d16a24480 100644 --- a/tests/unittest/scripts/test_bolt.py +++ b/tests/unittest/scripts/test_bolt.py @@ -120,7 +120,7 @@ def _compile_so(path: Path, emit_relocs: bool) -> bool: cmd = ["gcc", "-shared", "-fPIC", "-o", str(path), str(src)] if emit_relocs: cmd.insert(3, "-Wl,--emit-relocs") - return subprocess.run(cmd, capture_output=True).returncode == 0 + return subprocess.run(cmd, capture_output=True, timeout=60).returncode == 0 needs_toolchain = pytest.mark.skipif(