adding support for vllm sampler kwargs.

Author: NicoGrande

tunix/generate/vllm_sampler.py

@@ -71,6 +71,9 @@ class VllmConfig:
       init=False, default_factory=dict
   )
 
+  # vLLM sampler args that can be directly passed in without additional processing, e.g. temperature, stop etc.
+  sampler_kwargs: Dict[str, Any] = dataclasses.field(default_factory=dict)
+
   def __post_init__(self, engine_kwargs: Optional[Dict[str, Any]]):
     engine_kwargs = engine_kwargs or {}
     self._processed_engine_kwargs = engine_kwargs
@@ -418,16 +421,21 @@ def __call__(
       if seed is not None:
         sampling_params.seed = seed
 
-      if kwargs:
+      kwargs.update(self.config.sampler_kwargs)
+      if kwargs:     
         try:
-          sampling_params.update(**kwargs)
           logging.log_first_n(
               logging.INFO,
               "Received additional kwargs that are not explicitly defined in"
               f" the method signature: {kwargs}. These will be forwarded to the"
               " underlying sampler, but please ensure that they are valid.",
               1,
-          )
+          )   
+          for key, value in kwargs.items():
+            logging.info(
+                "Sampler kwargs setting key '%s' with value '%s'.", key, value
+            )
+            setattr(sampling_params, key, value)
         except Exception as e:
           logging.log_first_n(
               logging.INFO,

tunix/rl/rollout/base_rollout.py

@@ -157,9 +157,12 @@ class RolloutConfig:
   # Maximum number of concurrent sequences allowed to be processed in vLLM.
   rollout_vllm_max_num_seqs: Optional[int] = None
 
-  # Additional keyword arguments forwarded directly to the vLLM sampler/engine.
+  # Additional keyword arguments forwarded directly to the vLLM engine constructor.
   rollout_vllm_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
 
+  # Additional keyword arguments forwarded directly to the vLLM sampler.
+  rollout_vllm_sampler_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
+
   # SG-Lang JAX specific rollout configs.
 
   # Model version for SG-Lang JAX rollout engine.

tunix/rl/rollout/vllm_rollout.py

@@ -68,6 +68,7 @@ def __init__(
                 "hf_config_path": rollout_config.rollout_vllm_hf_config_path,
                 **rollout_config.rollout_vllm_kwargs,
             },
+            sampler_kwargs=rollout_config.rollout_vllm_sampler_kwargs,
         ),
     )
     state = nnx.state(model)