From 417fba365dbcd778b78411f55e1ab13658866d35 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Tue, 1 Sep 2026 01:00:16 -0700 Subject: [PATCH] [https://nvbugs/6693991][fix] Restore verl rollout after sampler_type removal Removing TRTLLMSampler and the `sampler_type` argument was a breaking change for callers outside this repository, and one of them is pinned by our own CI: tests/integration/defs/verl/verl_config.yml pins verl at d324b01, whose rollout server passes `sampler_type` unconditionally. Every verl rollout test therefore died constructing the engine. In verl a single switch, TLLM_USE_TORCHSAMPLER (default "0"), selects both the sampler and the configuration that matches it: unset, it asks for `sampler_type="TRTLLMSampler"` and pairs it with `end_id=-1`, `detokenize=False` and `logprobs=1`. Merely letting the argument through would run TorchSampler against the C++ sampler's conventions -- `logprobs=1` returns the top-1 logprob where TorchSampler reports the sampled token's at `logprobs=0` -- turning a construction error into a wrong rollout. So set TLLM_USE_TORCHSAMPLER=1 in the CI config, which moves verl's whole profile onto TorchSampler, and accept `sampler_type` only for the values that still resolve to it (`auto`, `TorchSampler`), rejecting `TRTLLMSampler` with a message naming the settings to review. Removed arguments are declared once, in _TORCH_LLMARGS_REMOVED_ARG_CHECKS, which the two pre-pydantic kwarg gates and the model validator all consult. Unwaive all six cases in test_trtllm_rollout_utils.py: they are the whole file and share class-scoped fixtures that build a real replica, so they failed and recover together. Signed-off-by: ZhaoyangWang --- docs/source/features/sampling.md | 2 + tensorrt_llm/llmapi/llm.py | 17 ++++--- tensorrt_llm/llmapi/llm_args.py | 53 +++++++++++++++++++-- tests/integration/defs/verl/verl_config.yml | 4 ++ tests/integration/test_lists/waives.txt | 6 --- 5 files changed, 66 insertions(+), 16 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 03231bb9aac0..514f9a2a6e37 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -17,6 +17,8 @@ The PyTorch backend supports a wide variety of features, listed below: ## General usage +TorchSampler is the only sampler; the `sampler_type` argument and the TRTLLM Sampler it selected have been removed. `sampler_type="auto"` and `sampler_type="TorchSampler"` are still accepted and ignored, while `sampler_type="TRTLLMSampler"` is rejected. When migrating off TRTLLM Sampler, also review the settings tuned for it, notably `logprobs`: under TorchSampler, `logprobs=0` returns the sampled token's log probability. + Here is an example to run a model with basic usage of sampling parameters. This example prepares two identical prompts which will give different results due to the sampling parameters chosen: ```python diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index db866443236e..666a61364ce2 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -57,7 +57,8 @@ from ..logger import logger from ..sampling_params import LogitsProcessor, SamplingParams from ..scheduling_params import SchedulingParams -from .llm_args import (TORCH_LLMARGS_EXPLICIT_DOCSTRING, TorchLlmArgs, +from .llm_args import (TORCH_LLMARGS_EXPLICIT_DOCSTRING, + TORCH_LLMARGS_REMOVED_ARGS, TorchLlmArgs, validate_token_encoder_bucket_config) from .llm_utils import (CachedModelLoader, KvCacheRetentionConfig, LlmBuildStats, ModelLoader) @@ -403,6 +404,9 @@ def __init__(self, valid_keys = set( list(llm_args_cls.model_fields.keys()) + ['_mpi_session', 'backend']) + if issubclass(llm_args_cls, TorchLlmArgs): + # Values are vetted by TorchLlmArgs._drop_removed_args. + valid_keys |= TORCH_LLMARGS_REMOVED_ARGS for key in kwargs: if key not in valid_keys: raise ValueError( @@ -2008,14 +2012,13 @@ def _build_model(self): def _validate_args_for_torch_backend(self, kwargs: dict) -> None: """Validate that only arguments supported by the PyTorch backend are passed. """ - torchllm_fields = set(TorchLlmArgs.model_fields.keys()) + # Values of removed args are vetted by TorchLlmArgs._drop_removed_args. + accepted_keys = (set(TorchLlmArgs.model_fields.keys()) + | TORCH_LLMARGS_REMOVED_ARGS + | {'_mpi_session', 'backend'}) # Check if any arguments not supported by the PyTorch backend are passed. - unsupported_args = [ - key for key in kwargs - if key not in torchllm_fields and key not in ('_mpi_session', - 'backend') - ] + unsupported_args = [key for key in kwargs if key not in accepted_keys] if unsupported_args: raise ValueError( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 79e8daffb1ea..d2523da8b622 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -25,9 +25,9 @@ from dataclasses import dataclass from enum import Enum, EnumMeta from pathlib import Path -from typing import (TYPE_CHECKING, Annotated, Any, ClassVar, Dict, List, - Literal, Optional, Set, Tuple, Type, TypeAlias, TypeVar, - Union, get_args, get_origin) +from typing import (TYPE_CHECKING, Annotated, Any, Callable, ClassVar, Dict, + List, Literal, Optional, Set, Tuple, Type, TypeAlias, + TypeVar, Union, get_args, get_origin) import torch import yaml @@ -5330,6 +5330,39 @@ def validate_capture_num_tokens(cls, v): "The maximum number of CUDA streams to use for torch.compile.") +def _check_removed_sampler_type(value: Any) -> None: + """Validate a value passed to the removed ``sampler_type`` argument. + + ``auto`` and ``TorchSampler`` always resolved to TorchSampler, so dropping + the argument honors them. ``TRTLLMSampler`` cannot be honored: callers + asking for it tune the rest of their configuration to its conventions, so + substituting TorchSampler would silently run a mismatched configuration. + """ + if value in ("auto", "TorchSampler"): + logger.warning( + f"'sampler_type' was removed (got {value!r}) and is ignored; " + "TorchSampler is the only sampler. Remove the argument to silence " + "this warning.") + return + raise ValueError( + f"'sampler_type' was removed and {value!r} is no longer available; " + "TorchSampler is the only sampler. Drop the argument, and make sure " + "the surrounding configuration targets TorchSampler -- notably " + "`logprobs=0` for the sampled token's logprob and a real `end_id`, " + "which differ from the TRTLLMSampler conventions.") + + +# Arguments removed from `TorchLlmArgs` but still accepted on input, so that +# integrations pinned to an older TRT-LLM keep constructing. Declaring one here +# is enough: the pre-pydantic kwarg gates in `llm.py` consult +# `TORCH_LLMARGS_REMOVED_ARGS`, and `_drop_removed_args` runs the check below. +_TORCH_LLMARGS_REMOVED_ARG_CHECKS: Dict[str, Callable[[Any], None]] = { + "sampler_type": _check_removed_sampler_type, +} + +TORCH_LLMARGS_REMOVED_ARGS = frozenset(_TORCH_LLMARGS_REMOVED_ARG_CHECKS) + + class TorchLlmArgs(BaseLlmArgs): # PyTorch backend specific configurations generation_config: Literal["auto", "trtllm"] = Field( @@ -5454,6 +5487,20 @@ def init_multimodal_config(cls, v): "model's largest atomic item when necessary."), status="prototype") + @model_validator(mode="before") + @classmethod + def _drop_removed_args(cls, data: Any) -> Any: + """Drop removed arguments; their checks reject unserviceable values.""" + if not isinstance(data, dict): + return data + removed = TORCH_LLMARGS_REMOVED_ARGS.intersection(data) + if not removed: + return data + data = dict(data) + for key in sorted(removed): + _TORCH_LLMARGS_REMOVED_ARG_CHECKS[key](data.pop(key)) + return data + @field_validator("encoder_max_batch_size", "encoder_max_num_tokens") @classmethod def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: diff --git a/tests/integration/defs/verl/verl_config.yml b/tests/integration/defs/verl/verl_config.yml index 300f0d52a308..8a37855aca99 100644 --- a/tests/integration/defs/verl/verl_config.yml +++ b/tests/integration/defs/verl/verl_config.yml @@ -50,6 +50,10 @@ verl_config: - "PATH=\"${NVSHMEM_DIR}/bin:$PATH\"" - "TRTLLM_TEST_MODEL_PATH_ROOT=/tmp/verl-models" - "TORCH_CUDA_ARCH_LIST=" + # A verl-side switch that TRT-LLM never reads. Besides the removed + # sampler_type, it also selects verl's rollout config (end_id, detokenize, + # logprobs), so it must be set for that config to match TorchSampler. + - "TLLM_USE_TORCHSAMPLER=1" # Read-only CI model cache (flat layout: /scratch.../ModelName) ci_model_cache: "/scratch.trt_llm_data/llm-models" diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 8d89dce2ab92..d14c9669f4ca 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -385,10 +385,4 @@ verl/test_verl_cases.py::test_async_generate SKIP (https://nvbugs/6683838) verl/test_verl_cases.py::test_async_memory_management SKIP (https://nvbugs/6683838) verl/test_verl_cases.py::test_init_without_device_mesh SKIP (https://nvbugs/6683838) verl/test_verl_cases.py::test_inter_node_trtllm_rollout SKIP (https://nvbugs/6683838) -verl/test_verl_cases.py::test_multimodal_different_image_sizes SKIP (https://nvbugs/6693991) -verl/test_verl_cases.py::test_multimodal_generate_with_image SKIP (https://nvbugs/6693991) -verl/test_verl_cases.py::test_multimodal_text_only_fallback SKIP (https://nvbugs/6693991) verl/test_verl_cases.py::test_trtllm_abort SKIP (https://nvbugs/6272653) -verl/test_verl_cases.py::test_unimodal_batch_generate SKIP (https://nvbugs/6693991) -verl/test_verl_cases.py::test_unimodal_generate SKIP (https://nvbugs/6693991) -verl/test_verl_cases.py::test_wake_sleep_cycle SKIP (https://nvbugs/6693991)