diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py index b15e04ab74..29ad78dae2 100644 --- a/unsloth/models/__init__.py +++ b/unsloth/models/__init__.py @@ -20,4 +20,4 @@ from .mistral import FastMistralModel from .qwen2 import FastQwen2Model from .dpo import PatchDPOTrainer, PatchKTOTrainer from ._utils import is_bfloat16_supported -from .rl import PatchFastRL +from .rl import PatchFastRL, vLLMSamplingParams diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 3ab45cdf72..572caf594c 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -44,6 +44,7 @@ torch_compile_options = { def vLLMSamplingParams(**kwargs): + from vllm import SamplingParams sampling_params = SamplingParams(**kwargs) sampling_params._set_kwargs = kwargs return sampling_params