Add Nemotron 3 Nano tests (#18119)

Signed-off-by: Shahar Mor <smor@nvidia.com>

python/pyproject.toml

@@ -128,6 +128,7 @@ test = [
   "bitsandbytes",
   "expecttest",
   "jsonlines",
+  "lm-eval[api]>=0.4.9.2",
   "matplotlib",
   "pandas",
   "parameterized",

python/sglang/test/kits/lm_eval_kit.py

@@ -0,0 +1,109 @@
+"""
+This module provides a mixin class for running lm-eval harness evaluations
+against SGLang servers
+"""
+
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import requests
+import yaml
+
+
+@contextmanager
+def scoped_env_vars(new_env: dict[str, str] | None):
+    """Context manager to temporarily set environment variables."""
+    if not new_env:
+        yield
+        return
+
+    old_values = {}
+    new_keys = []
+
+    try:
+        for key, value in new_env.items():
+            if key in os.environ:
+                old_values[key] = os.environ[key]
+            else:
+                new_keys.append(key)
+            os.environ[key] = str(value)
+        yield
+    finally:
+        for key, value in old_values.items():
+            os.environ[key] = value
+        for key in new_keys:
+            os.environ.pop(key, None)
+
+
+class LMEvalMixin:
+    """
+    Mixin class for running lm-eval harness evaluations.
+    """
+
+    other_args: list[str] = []
+    model_config_name: str = ""
+    default_rtol: float = 0.08
+
+    def test_lm_eval(self):
+        """Run lm-eval evaluation and validate results."""
+        # Flush cache before evaluation
+        requests.get(self.base_url + "/flush_cache")
+
+        eval_config = yaml.safe_load(
+            Path(self.model_config_name).read_text(encoding="utf-8")
+        )
+        results = self.launch_lm_eval(eval_config)
+
+        rtol = eval_config.get("rtol", self.default_rtol)
+
+        success = True
+        for task in eval_config["tasks"]:
+            for metric in task["metrics"]:
+                ground_truth = metric["value"]
+                measured_value = results["results"][task["name"]][metric["name"]]
+                print(
+                    f"{task['name']} | {metric['name']}: "
+                    f"ground_truth={ground_truth:.3f} | "
+                    f"measured={measured_value:.3f} | rtol={rtol}"
+                )
+                success = success and np.isclose(
+                    ground_truth, measured_value, rtol=rtol
+                )
+
+        self.assertTrue(success, f"lm-eval validation failed")
+
+    def launch_lm_eval(self, eval_config: dict[str, Any]) -> dict:
+        """
+        Args:
+            eval_config: Configuration dictionary with model and task settings
+        """
+        import lm_eval
+
+        batch_size = eval_config.get("batch_size", "auto")
+        backend = eval_config.get("backend", "local-completions")
+        num_concurrent = eval_config.get("num_concurrent", 1)
+
+        model_args = {
+            "model": eval_config["model_name"],
+            "base_url": self.base_url + "/v1/completions",
+            "num_concurrent": num_concurrent,
+        }
+
+        env_vars = eval_config.get("env_vars", None)
+        with scoped_env_vars(env_vars):
+            results = lm_eval.simple_evaluate(
+                model=backend,
+                model_args=model_args,
+                tasks=[task["name"] for task in eval_config["tasks"]],
+                num_fewshot=eval_config.get("num_fewshot", 0),
+                limit=eval_config.get("limit", None),
+                apply_chat_template=eval_config.get("apply_chat_template", False),
+                fewshot_as_multiturn=eval_config.get("fewshot_as_multiturn", False),
+                gen_kwargs=eval_config.get("gen_kwargs"),
+                batch_size=batch_size,
+            )
+
+        return results

test/lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml

@@ -0,0 +1,13 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.847
+  - name: "exact_match,flexible-extract"
+    value: 0.556
+limit: 1319
+num_concurrent: 128
+num_fewshot: 5
+apply_chat_template: false
+fewshot_as_multiturn: true

test/lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml

@@ -0,0 +1,13 @@
+model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
+tasks:
+- name: "gsm8k"
+  metrics:
+  - name: "exact_match,strict-match"
+    value: 0.847
+  - name: "exact_match,flexible-extract"
+    value: 0.556
+limit: 1319
+num_concurrent: 128
+num_fewshot: 5
+apply_chat_template: false
+fewshot_as_multiturn: true

test/registered/models/test_nvidia_nemotron_3_nano.py

@@ -0,0 +1,41 @@
+import unittest
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.lm_eval_kit import LMEvalMixin
+from sglang.test.server_fixtures.default_fixture import DefaultServerBase
+
+register_cuda_ci(est_time=180, suite="stage-b-test-large-2-gpu")
+
+NEMOTRON_3_NANO_THINKING_ARGS = [
+    "--trust-remote-code",
+    "--tool-call-parser",
+    "qwen3_coder",
+    "--reasoning-parser",
+    "deepseek-r1",
+]
+
+
+class TestNvidiaNemotron3Nano30BBF16(LMEvalMixin, DefaultServerBase):
+    """Test Nemotron-3-Nano-30B BF16 model with lm-eval GSM8K evaluation."""
+
+    model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
+    model_config_name = "lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.yaml"
+    other_args = [
+        "--tp-size",
+        "2",
+    ] + NEMOTRON_3_NANO_THINKING_ARGS
+
+
+class TestNvidiaNemotron3Nano30BFP8(LMEvalMixin, DefaultServerBase):
+    """Test Nemotron-3-Nano-30B FP8 model with lm-eval GSM8K evaluation."""
+
+    model = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8"
+    model_config_name = "lm_eval_configs/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8.yaml"
+    other_args = [
+        "--tp-size",
+        "2",
+    ] + NEMOTRON_3_NANO_THINKING_ARGS
+
+
+if __name__ == "__main__":
+    unittest.main()