[Tunix] Add vLLM sampler for math eval.

PiperOrigin-RevId: 875875039

Author: lc5211

examples/deepscaler/math_eval_nb.py

@@ -231,7 +231,7 @@ def load_model(self):
           tokenizer=self.tokenizer,
           cache_config=cache_config,
       )
-    elif self.sampler_type == "sglang-jax":
+    elif self.sampler_type == "sglang_jax":
       from tunix.google.stubs import sglang_jax_sampler_stub as sglang_jax_sampler  # pylint: disable=g-import-not-at-top
 
       mapping_config = mappings.MappingConfig.build(
@@ -254,6 +254,27 @@ def load_model(self):
               mapping_config=mapping_config,
           ),
       )
+    elif self.sampler_type == "vllm":
+      from tunix.google.stubs import vllm_sampler_stub as vllm_sampler  # pylint: disable=g-import-not-at-top
+
+      mapping_config = mappings.MappingConfig.build(
+          mapping_obj=None,
+          model=self.model,
+          backend="vllm",
+      )
+      self.sampler_vllm = vllm_sampler.VllmSampler(
+          tokenizer=self.tokenizer,
+          config=vllm_sampler.VllmConfig(
+              mesh=self.mesh,
+              max_model_len=self.max_prompt_length
+              + self.max_generation_steps
+              + 100,
+              model_version=self.model_version,
+              hbm_utilization=0.4,
+              init_with_random_weights=False,
+              mapping_config=mapping_config,
+          ),
+      )
     else:
       raise ValueError(f"Unsupported sampler type: {self.sampler_type}")
 
@@ -358,7 +379,7 @@ def generate(
           eos_tokens=[stop_token_id],
           seed=jax.random.PRNGKey(seed) if seed is not None else None,
       )
-    elif self.sampler_type == "sglang-jax":
+    elif self.sampler_type == "sglang_jax":
       out_data = self.sampler_sglang(
           input_strings=prompts,
           max_generation_steps=safe_gen_length,
@@ -370,6 +391,18 @@ def generate(
           echo=False,
           pad_output=True,
       )
+    elif self.sampler_type == "vllm":
+      out_data = self.sampler_vllm(
+          input_strings=prompts,
+          max_generation_steps=safe_gen_length,
+          max_prompt_length=self.max_prompt_length,
+          temperature=temperature,
+          top_p=top_p,
+          top_k=top_k,
+          seed=seed,
+          echo=False,
+          pad_output=True,
+      )
     else:
       raise ValueError(f"Unsupported sampler type: {self.sampler_type}")
     return out_data.text
@@ -585,6 +618,7 @@ def evaluate(
 # %%
 # AIME-2024
 model_version = "agentica-org/DeepScaleR-1.5B-Preview"
+# model_version = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 dataset = AIME_2024_DATA_PATH
 model_config, model_path = MODEL_MAPPING[model_version]