diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 291abceb1c18..0cbb6442f798 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import json +import math import os from abc import ABC, abstractmethod from dataclasses import dataclass, field, fields @@ -27,6 +28,11 @@ MAX_TOP_LOGPROBS = 20 +# Smallest non-zero temperature forwarded to the sampling backend. Values below +# this overflow logits/temperature to inf/nan in float16/bfloat16 (mirrors the +# guard vLLM uses). temperature == 0 is greedy decoding and is left untouched. +MIN_SAMPLING_TEMPERATURE = 1e-2 + def validate_thinking_token_budget(value: Optional[Union[int, float, bool]]) -> Optional[int]: """Validate ``thinking_token_budget``; return ``None`` if unset.""" @@ -215,6 +221,7 @@ class SamplingParams: top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. None means using C++ runtime default 1.f. Defaults to None. seed (int, optional): Controls the random seed used by the random number generator in sampling. None means using C++ runtime default 0. Defaults to None. temperature (float, optional): Controls the modulation of logits when sampling new tokens. It can have values >= 0.f. Defaults to None. + Positive values smaller than 1e-2 are clamped up to 1e-2 to avoid inf/nan when dividing logits by the temperature in fp16/bf16; temperature = 0 stays greedy and non-finite values are rejected. The value None is treated as "not specified" in the following. If neither temperature, top_p, nor top_k are specified, sampling is greedy. If top_p < 1 and/or top_k > 1 are specified, sampling will proceed accordingly and temperature will default to temperature = 1. @@ -370,8 +377,22 @@ def _validate(self): raise ValueError(f"require 0 <= top_p <= 1, got top_p={self.top_p}") if self.top_k is not None and self.top_k < 0: raise ValueError(f"require top_k >= 0, got top_k={self.top_k}") - if self.temperature is not None and self.temperature < 0: - raise ValueError(f"require temperature >= 0, got temperature={self.temperature}") + if self.temperature is not None: + if not math.isfinite(self.temperature): + raise ValueError( + f"require temperature to be finite, got temperature={self.temperature}" + ) + if self.temperature < 0: + raise ValueError(f"require temperature >= 0, got temperature={self.temperature}") + # Clamp very small positive temperatures up to a numerically safe + # floor: logits / temperature otherwise overflows to inf/nan in + # fp16/bf16. temperature == 0 (greedy) is intentionally left as-is. + if 0 < self.temperature < MIN_SAMPLING_TEMPERATURE: + logger.debug( + f"Clamping temperature from {self.temperature} to " + f"{MIN_SAMPLING_TEMPERATURE} to avoid numerical instability." + ) + self.temperature = MIN_SAMPLING_TEMPERATURE if self.best_of is not None and self.best_of < self.n: raise ValueError(f"best_of ({self.best_of}) cannot be less than n ({self.n})") diff --git a/tests/unittest/llmapi/test_sampling_params_temperature.py b/tests/unittest/llmapi/test_sampling_params_temperature.py new file mode 100644 index 000000000000..66a1003522f6 --- /dev/null +++ b/tests/unittest/llmapi/test_sampling_params_temperature.py @@ -0,0 +1,33 @@ +import math + +import pytest + +from tensorrt_llm.sampling_params import MIN_SAMPLING_TEMPERATURE, SamplingParams + + +@pytest.mark.parametrize( + ("temperature", "expected"), + [ + (None, None), # unset -> unchanged + (0.0, 0.0), # greedy -> untouched + (1e-12, MIN_SAMPLING_TEMPERATURE), # tiny -> clamped (regression #15715) + (MIN_SAMPLING_TEMPERATURE - 1e-3, MIN_SAMPLING_TEMPERATURE), # just below floor -> clamped + (MIN_SAMPLING_TEMPERATURE, MIN_SAMPLING_TEMPERATURE), # boundary -> kept + (1.0, 1.0), # normal -> kept + ], +) +def test_temperature_clamped_below_floor(temperature, expected): + # #15715: 0 < temperature < MIN reaches the backend and overflows + # logits / temperature to inf/nan in fp16/bf16; clamp it to the floor. + assert SamplingParams(temperature=temperature, top_k=2).temperature == expected + + +def test_negative_temperature_rejected(): + with pytest.raises(ValueError): + SamplingParams(temperature=-0.5) + + +@pytest.mark.parametrize("temperature", [math.inf, -math.inf, math.nan]) +def test_non_finite_temperature_rejected(temperature): + with pytest.raises(ValueError): + SamplingParams(temperature=temperature)