hendrycks sanity ci (#1900)

* do not oversample hendrycks sanity (trainer bottlenecks)

* put hendrycks sanity into nightly config

* remove duplicate utils

Author: mikasenghaas

configs/ci/nightly/hendrycks_math.toml

@@ -15,7 +15,6 @@ name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 [orchestrator]
 batch_size = 512
 rollouts_per_example = 8
-oversampling_factor = 2.0
 
 [[orchestrator.env]]
 id = "math-env" # included in lock file

examples/hendrycks_sanity/rl.toml

@@ -1,11 +1,10 @@
-from functools import partial
 from pathlib import Path
 from typing import Callable
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_no_error, check_number_goes_up_or_down, check_reward_in_range, strip_escape_codes
+from tests.utils import check_no_error, check_reward_goes_up, check_reward_in_range, strip_escape_codes
 
 pytestmark = [pytest.mark.gpu, pytest.mark.slow]
 
@@ -49,9 +48,6 @@ def rl_process(
     return run_process(cmd, timeout=TIMEOUT)
 
 
-check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})")
-
-
 @pytest.fixture(scope="module")
 def test_no_error(rl_process: ProcessResult, rl_output_dir: Path):
     """Tests that the RL process does not fail."""

tests/integration/test_alphabet_sort.py

@@ -1,11 +1,10 @@
-from functools import partial
 from pathlib import Path
 from typing import Callable
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_no_error, check_number_goes_up_or_down, check_reward_in_range, strip_escape_codes
+from tests.utils import check_no_error, check_reward_goes_up, check_reward_in_range, strip_escape_codes
 
 pytestmark = [pytest.mark.gpu, pytest.mark.slow]
 
@@ -70,9 +69,6 @@ def rl_resume_process(
     return run_process(cmd, timeout=TIMEOUT)
 
 
-check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})")
-
-
 @pytest.fixture(scope="module")
 def test_no_error(rl_process: ProcessResult, output_dir: Path):
     """Tests that the RL process does not fail."""

tests/integration/test_rl.py

@@ -1,11 +1,10 @@
-from functools import partial
 from pathlib import Path
 from typing import Callable
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_no_error, check_number_goes_up_or_down, check_reward_in_range, strip_escape_codes
+from tests.utils import check_no_error, check_reward_goes_up, check_reward_in_range, strip_escape_codes
 
 pytestmark = [pytest.mark.gpu, pytest.mark.slow]
 
@@ -70,9 +69,6 @@ def rl_resume_process(
     return run_process(cmd, env={"VLLM_ALLOW_RUNTIME_LORA_UPDATING": "True"}, timeout=TIMEOUT)
 
 
-check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})")
-
-
 @pytest.fixture(scope="module")
 def test_no_error(rl_process: ProcessResult, output_dir: Path):
     """Tests that the RL process does not fail."""

tests/integration/test_rl_lora.py

@@ -5,14 +5,13 @@
 import signal
 import subprocess
 import time
-from functools import partial
 from pathlib import Path
 from typing import Generator
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_number_goes_up_or_down, check_reward_in_range, strip_escape_codes
+from tests.utils import check_reward_goes_up, check_reward_in_range, strip_escape_codes
 
 pytestmark = [pytest.mark.gpu, pytest.mark.slow]
 
@@ -364,9 +363,6 @@ def multi_run_result(
                 p.kill()
 
 
-check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})")
-
-
 def test_remaining_orchestrators_complete(
     multi_run_result: dict[str, ProcessResult],
     output_dir: Path,

tests/integration/test_rl_multi_run_lora.py

@@ -1,11 +1,10 @@
-from functools import partial
 from pathlib import Path
 from typing import Callable
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_number_goes_up_or_down, strip_escape_codes
+from tests.utils import check_loss_goes_down, strip_escape_codes
 
 pytestmark = [pytest.mark.slow, pytest.mark.gpu]
 
@@ -80,9 +79,6 @@ def sft_resume_process(
     return run_process(cmd, timeout=TIMEOUT)
 
 
-check_loss_goes_down = partial(check_number_goes_up_or_down, go_up=False, pattern=r"Loss:\s*(\d+\.\d{4})")
-
-
 def test_no_error(sft_process: ProcessResult):
     """Tests that the SFT process does not fail."""
     assert sft_process.returncode == 0, f"Process has non-zero return code ({sft_process})"

tests/integration/test_sft.py

@@ -1,11 +1,10 @@
-from functools import partial
 from pathlib import Path
 from typing import Callable
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_mismatch_kl_in_range, check_no_error, check_number_goes_up_or_down, strip_escape_codes
+from tests.utils import check_mismatch_kl_in_range, check_no_error, check_reward_goes_up, strip_escape_codes
 
 pytestmark = [pytest.mark.gpu, pytest.mark.slow]
 
@@ -39,7 +38,6 @@ def rl_process(
     return run_process(cmd)
 
 
-check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})")
 MISMATCH_KL_MIN = 0.0
 MISMATCH_KL_MAX = 0.0015
 

tests/nightly/test_acereason_math.py

@@ -1,11 +1,10 @@
-from functools import partial
 from pathlib import Path
 from typing import Callable
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_no_error, check_number_goes_up_or_down, strip_escape_codes
+from tests.utils import check_no_error, check_reward_goes_up, strip_escape_codes
 
 pytestmark = [pytest.mark.gpu, pytest.mark.slow]
 
@@ -39,9 +38,6 @@ def rl_process(
     return run_process(cmd)
 
 
-check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})")
-
-
 @pytest.fixture(scope="module")
 def test_no_error(rl_process: ProcessResult, output_dir: Path):
     """Tests that the RL process does not fail."""

tests/nightly/test_alphabet_sort.py

@@ -1,19 +1,24 @@
-from functools import partial
 from pathlib import Path
 from typing import Callable
 
 import pytest
 
 from tests.conftest import ProcessResult
-from tests.utils import check_mismatch_kl_in_range, check_no_error, check_number_goes_up_or_down, strip_escape_codes
+from tests.utils import (
+    check_mismatch_kl_in_range,
+    check_no_error,
+    check_reward_goes_up,
+    check_reward_in_range,
+    strip_escape_codes,
+)
 
 pytestmark = [pytest.mark.gpu, pytest.mark.slow]
 
 
 @pytest.fixture(scope="module")
 def wandb_name(branch_name: str) -> str:
     """Fixture for W&B name for RL CI integration tests."""
-    return f"hendrycks-math-{branch_name}"
+    return f"hendrycks-sanity-{branch_name}"
 
 
 @pytest.fixture(scope="module")
@@ -28,20 +33,22 @@ def rl_process(
         "run",
         "rl",
         "@",
-        "configs/ci/nightly/hendrycks_math.toml",
+        "examples/hendrycks_sanity/rl.toml",
         "--wandb.project",
         wandb_project,
         "--wandb.name",
         wandb_name,
         "--output-dir",
         output_dir.as_posix(),
+        "--max-steps",
+        "1000",  # do less steps to finish in time
     ]
     return run_process(cmd)
 
 
-check_reward_goes_up = partial(check_number_goes_up_or_down, go_up=True, pattern=r"Reward:\s*(\d+\.\d{4})")
+MIN_TRAIN_REWARD = 0.75
 MISMATCH_KL_MIN = 0.0
-MISMATCH_KL_MAX = 0.0004
+MISMATCH_KL_MAX = 0.0005
 
 
 @pytest.fixture(scope="module")
@@ -51,12 +58,19 @@ def test_no_error(rl_process: ProcessResult, output_dir: Path):
 
 
 def test_reward_goes_up(rl_process: ProcessResult, test_no_error, output_dir: Path):
-    """Tests that the reward goes up in the RL process"""
+    """Tests that the train reward goes up during training"""
     with open(output_dir / "logs" / "orchestrator.stdout", "r") as f:
         orchestrator_stdout = strip_escape_codes(f.read()).splitlines()
     check_reward_goes_up(orchestrator_stdout)
 
 
+def test_reward_reaches_threshold(rl_process: ProcessResult, test_no_error, output_dir: Path):
+    """Tests that the train reward reaches a minimum threshold"""
+    with open(output_dir / "logs" / "orchestrator.stdout", "r") as f:
+        orchestrator_stdout = strip_escape_codes(f.read()).splitlines()
+    check_reward_in_range(orchestrator_stdout, min_threshold=MIN_TRAIN_REWARD)
+
+
 def test_mismatch_kl_in_band(rl_process: ProcessResult, test_no_error, output_dir: Path):
     """Tests that mismatch KL stays within the expected band."""
     with open(output_dir / "logs" / "trainer.stdout", "r") as f: