Enable accuracy checks in CI & change fsdp_distributed test to non-colocate async

tests/command_utils.py

@@ -116,3 +116,26 @@ def exec_command(cmd: str, capture_output: bool = False):
     result = subprocess.run(["bash", "-c", cmd], shell=False, check=True, capture_output=capture_output)
     if capture_output:
         return result.stdout
+
+
+_warned_bool_env_var_keys = set()
+
+
+# copied from SGLang
+def get_bool_env_var(name: str, default: str = "false") -> bool:
+    value = os.getenv(name, default)
+    value = value.lower()
+
+    truthy_values = ("true", "1")
+    falsy_values = ("false", "0")
+
+    if (value not in truthy_values) and (value not in falsy_values):
+        if value not in _warned_bool_env_var_keys:
+            print(f"get_bool_env_var({name}) see non-understandable value={value} and treat as false")
+        _warned_bool_env_var_keys.add(value)
+
+    return value in truthy_values
+
+
+def get_env_enable_infinite_run():
+    return get_bool_env_var("MILES_TEST_ENABLE_INFINITE_RUN", "false")

tests/test_quick_start_glm4_9B.py

@@ -1,9 +1,7 @@
-import os
-
 import command_utils as U
 
-ENABLE_EVAL = bool(int(os.environ.get("MILES_TEST_ENABLE_EVAL", "1")))
-TIGHT_DEVICE_MEMORY = bool(int(os.environ.get("MILES_TEST_TIGHT_DEVICE_MEMORY", "1")))
+ENABLE_EVAL = U.get_bool_env_var("MILES_TEST_ENABLE_EVAL", "1")
+TIGHT_DEVICE_MEMORY = U.get_bool_env_var("MILES_TEST_TIGHT_DEVICE_MEMORY", "1")
 
 MODEL_NAME = "GLM-Z1-9B-0414"
 MODEL_TYPE = "glm4-9B"
@@ -88,6 +86,8 @@ def execute():
 
     sglang_args = "--rollout-num-gpus-per-engine 2 " "--use-miles-router "
 
+    ci_args = "--ci-test "
+
     misc_args = (
         # default dropout in megatron is 0.1
         "--attention-dropout 0.0 "
@@ -97,7 +97,6 @@ def execute():
         "--attention-softmax-in-fp32 "
         # need to comment this when using model with MLA
         "--attention-backend flash "
-        "--ci-test "
         "--actor-num-nodes 1 "
         "--actor-num-gpus-per-node 4 "
         "--rollout-num-gpus 4 "
@@ -112,6 +111,7 @@ def execute():
         f"{perf_args} "
         f"{eval_args} "
         f"{sglang_args} "
+        f"{ci_args} "
         f"{misc_args} "
     )
 

tests/test_qwen2.5_0.5B_gsm8k.py

@@ -1,9 +1,8 @@
-import os
 import command_utils as U
 
 
-FEW_GPU = bool(int(os.environ.get("MILES_TEST_FEW_GPU", "1")))
-TIGHT_DEVICE_MEMORY = bool(int(os.environ.get("MILES_TEST_TIGHT_DEVICE_MEMORY", "1")))
+FEW_GPU = U.get_bool_env_var("MILES_TEST_FEW_GPU", "1")
+TIGHT_DEVICE_MEMORY = U.get_bool_env_var("MILES_TEST_TIGHT_DEVICE_MEMORY", "1")
 
 MODEL_NAME = "Qwen2.5-0.5B-Instruct"
 MODEL_TYPE = "qwen2.5-0.5B"
@@ -27,7 +26,7 @@ def execute():
         "--apply-chat-template "
         "--rollout-shuffle "
         "--rm-type math "
-        "--num-rollout 3000 "
+        f"--num-rollout {3000 if U.get_env_enable_infinite_run() else 250} "
         "--rollout-batch-size 32 "
         "--n-samples-per-prompt 8 "
         "--rollout-max-response-len 1024 "
@@ -80,6 +79,13 @@ def execute():
         "--rollout-num-gpus-per-engine 1 " f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} "
     )
 
+    ci_args = (
+        "--ci-test "
+        "--ci-disable-kl-checker "
+        "--ci-metric-checker-key eval/gsm8k "
+        "--ci-metric-checker-threshold 0.55 "  # loose threshold at 250 step
+    )
+
     misc_args = (
         # default dropout in megatron is 0.1
         "--attention-dropout 0.0 "
@@ -103,6 +109,7 @@ def execute():
         f"{perf_args} "
         f"{eval_args} "
         f"{sglang_args} "
+        f"{ci_args} "
         f"{misc_args} "
     )
 

tests/test_qwen2.5_0.5B_gsm8k_async.py

@@ -1,8 +1,7 @@
-import os
 import command_utils as U
 
-FEW_GPU = bool(int(os.environ.get("MILES_TEST_FEW_GPU", "1")))
-TIGHT_DEVICE_MEMORY = bool(int(os.environ.get("MILES_TEST_TIGHT_DEVICE_MEMORY", "1")))
+FEW_GPU = U.get_bool_env_var("MILES_TEST_FEW_GPU", "1")
+TIGHT_DEVICE_MEMORY = U.get_bool_env_var("MILES_TEST_TIGHT_DEVICE_MEMORY", "1")
 
 MODEL_NAME = "Qwen2.5-0.5B-Instruct"
 MODEL_TYPE = "qwen2.5-0.5B"
@@ -26,7 +25,7 @@ def execute():
         "--apply-chat-template "
         "--rollout-shuffle "
         "--rm-type math "
-        "--num-rollout 3000 "
+        f"--num-rollout {3000 if U.get_env_enable_infinite_run() else 250} "
         "--rollout-batch-size 32 "
         "--n-samples-per-prompt 8 "
         "--rollout-max-response-len 1024 "
@@ -79,6 +78,13 @@ def execute():
         "--rollout-num-gpus-per-engine 1 " f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} "
     )
 
+    ci_args = (
+        "--ci-test "
+        "--ci-disable-kl-checker "
+        "--ci-metric-checker-key eval/gsm8k "
+        "--ci-metric-checker-threshold 0.55 "  # loose threshold at 250 step
+    )
+
     misc_args = (
         # default dropout in megatron is 0.1
         "--attention-dropout 0.0 "
@@ -102,6 +108,7 @@ def execute():
         f"{perf_args} "
         f"{eval_args} "
         f"{sglang_args} "
+        f"{ci_args} "
         f"{misc_args} "
     )
 

tests/test_qwen3_0.6B_fsdp_colocated_2xGPU.py

@@ -19,7 +19,7 @@ def execute():
         "--apply-chat-template "
         "--rollout-shuffle "
         "--rm-type math "
-        "--num-rollout 3000 "
+        f"--num-rollout {3000 if U.get_env_enable_infinite_run() else 60} "
         "--rollout-batch-size 32 "
         "--n-samples-per-prompt 8 "
         "--rollout-max-response-len 1024 "
@@ -66,6 +66,13 @@ def execute():
         "--update-weight-buffer-size 536870912 "  # 512MB
     )
 
+    ci_args = (
+        "--ci-test "
+        "--ci-disable-kl-checker "
+        "--ci-metric-checker-key eval/gsm8k "
+        "--ci-metric-checker-threshold 0.71 "  # loose threshold at 60 step
+    )
+
     misc_args = "--actor-num-nodes 1 " "--actor-num-gpus-per-node 2 " "--colocate " "--train-backend fsdp "
 
     train_args = (
@@ -77,6 +84,7 @@ def execute():
         f"{U.get_default_wandb_args(__file__)} "
         f"{eval_args} "
         f"{fsdp_args} "
+        f"{ci_args} "
         f"{misc_args} "
     )
 

tests/test_qwen3_0.6B_fsdp_distributed.py

@@ -1,10 +1,9 @@
-import os
 import command_utils as U
 
 MODEL_NAME = "Qwen3-0.6B"
 
 
-FEW_GPU = bool(int(os.environ.get("MILES_TEST_FEW_GPU", "1")))
+FEW_GPU = U.get_bool_env_var("MILES_TEST_FEW_GPU", "1")
 
 
 def prepare():
@@ -23,7 +22,7 @@ def execute():
         "--apply-chat-template "
         "--rollout-shuffle "
         "--rm-type math "
-        "--num-rollout 3000 "
+        f"--num-rollout {3000 if U.get_env_enable_infinite_run() else 60} "
         "--rollout-batch-size 32 "
         "--n-samples-per-prompt 8 "
         "--rollout-max-response-len 1024 "
@@ -65,11 +64,18 @@ def execute():
 
     misc_args = (
         "--actor-num-nodes 1 "
-        f"--actor-num-gpus-per-node {2 if FEW_GPU else 4} "
-        "--colocate "
+        f"--actor-num-gpus-per-node {1 if FEW_GPU else 2} "
+        f"--rollout-num-gpus {1 if FEW_GPU else 2} "
         "--train-backend fsdp "
     )
 
+    ci_args = (
+        "--ci-test "
+        "--ci-disable-kl-checker "
+        "--ci-metric-checker-key eval/gsm8k "
+        "--ci-metric-checker-threshold 0.71 "  # loose threshold at 60 step
+    )
+
     train_args = (
         f"{ckpt_args} "
         f"{rollout_args} "
@@ -78,13 +84,15 @@ def execute():
         f"{U.get_default_wandb_args(__file__)} "
         f"{eval_args} "
         f"{sglang_args} "
+        f"{ci_args} "
         f"{misc_args} "
     )
 
     U.execute_train(
         train_args=train_args,
         num_gpus=2 if FEW_GPU else 4,
         model_type=None,
+        train_script="train_async.py",
     )
 
 

tests/test_qwen3_30B_A3B.py

@@ -95,6 +95,8 @@ def execute():
         "--sglang-disable-radix-cache "
     )
 
+    ci_args = "--ci-test "
+
     misc_args = (
         # default dropout in megatron is 0.1
         "--attention-dropout 0.0 "
@@ -106,7 +108,6 @@ def execute():
         "--attention-backend flash "
         "--moe-token-dispatcher-type flex "
         "--moe-enable-deepep "
-        "--ci-test "
         "--actor-num-nodes 1 "
         "--actor-num-gpus-per-node 8 "
         "--colocate "
@@ -121,6 +122,7 @@ def execute():
         f"{perf_args} "
         f"{eval_args} "
         f"{sglang_args} "
+        f"{ci_args} "
         f"{misc_args} "
     )