diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 291abceb1c18..177996a75817 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 @@ -26,6 +27,7 @@ from tensorrt_llm.logger import logger MAX_TOP_LOGPROBS = 20 +MIN_SAMPLING_TEMPERATURE = 1e-2 def validate_thinking_token_budget(value: Optional[Union[int, float, bool]]) -> Optional[int]: @@ -214,7 +216,8 @@ class SamplingParams: top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. Indicates where to reset the decay. None means using C++ runtime default 1. Defaults to None. 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. + temperature (float, optional): Controls the modulation of logits when sampling new tokens. It must be finite and can have values >= 0.f. Defaults to None. + Positive values smaller than 1e-2 are normalized to 1e-2 for numerical stability, while 0 preserves greedy decoding. 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 +373,19 @@ 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}") + if 0 < self.temperature < MIN_SAMPLING_TEMPERATURE: + logger.debug( + f"Clamping temperature from {self.temperature} to {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.py b/tests/unittest/llmapi/test_sampling_params.py index 8988cc137b93..fee70f2bda11 100644 --- a/tests/unittest/llmapi/test_sampling_params.py +++ b/tests/unittest/llmapi/test_sampling_params.py @@ -14,6 +14,7 @@ # limitations under the License. import asyncio import json +import math import pytest import torch @@ -22,7 +23,12 @@ ThinkingBudgetLogitsProcessor, add_thinking_budget_logits_processor, ) -from tensorrt_llm.sampling_params import MAX_TOP_LOGPROBS, SamplingParams, check_logprobs_limit +from tensorrt_llm.sampling_params import ( + MAX_TOP_LOGPROBS, + MIN_SAMPLING_TEMPERATURE, + SamplingParams, + check_logprobs_limit, +) from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, CompletionRequest, @@ -53,6 +59,36 @@ def test_logprobs_request_limit(field): SamplingParams(**{field: MAX_TOP_LOGPROBS + 1}) +@pytest.mark.parametrize( + ("temperature", "expected"), + [ + (0.0, 0.0), + (1e-6, MIN_SAMPLING_TEMPERATURE), + (MIN_SAMPLING_TEMPERATURE - 1e-3, MIN_SAMPLING_TEMPERATURE), + (MIN_SAMPLING_TEMPERATURE, MIN_SAMPLING_TEMPERATURE), + (1.0, 1.0), + ], +) +def test_temperature_clamps_small_nonzero_values(temperature, expected): + """Small non-zero temperatures are clamped while zero remains greedy.""" + sampling_params = SamplingParams(temperature=temperature, top_k=2) + + assert sampling_params.temperature == expected + assert sampling_params._get_sampling_config().temperature == pytest.approx(expected) + assert sampling_params._greedy_decoding is (temperature == 0) + + +@pytest.mark.parametrize("temperature", [math.inf, -math.inf, math.nan]) +def test_temperature_rejects_non_finite_values(temperature): + with pytest.raises(ValueError, match="require temperature to be finite"): + SamplingParams(temperature=temperature) + + +def test_temperature_rejects_negative_values(): + with pytest.raises(ValueError, match="require temperature >= 0"): + SamplingParams(temperature=-0.5) + + def test_chat_top_logprobs_request_limit(): with pytest.raises(ValueError, match=f"less than or equal to {MAX_TOP_LOGPROBS}"): ChatCompletionRequest(