diff --git a/tests/command_utils.py b/tests/command_utils.py index e4c10eea3f3..bf5c33ef310 100644 --- a/tests/command_utils.py +++ b/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") diff --git a/tests/test_quick_start_glm4_9B.py b/tests/test_quick_start_glm4_9B.py index 21b964f0eab..b269cedaef7 100644 --- a/tests/test_quick_start_glm4_9B.py +++ b/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} " ) diff --git a/tests/test_qwen2.5_0.5B_gsm8k.py b/tests/test_qwen2.5_0.5B_gsm8k.py index febae5144fd..9da332b06be 100644 --- a/tests/test_qwen2.5_0.5B_gsm8k.py +++ b/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} " ) diff --git a/tests/test_qwen2.5_0.5B_gsm8k_async.py b/tests/test_qwen2.5_0.5B_gsm8k_async.py index dcfcf56b064..97dfb69ed06 100644 --- a/tests/test_qwen2.5_0.5B_gsm8k_async.py +++ b/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} " ) diff --git a/tests/test_qwen3_0.6B_fsdp_colocated_2xGPU.py b/tests/test_qwen3_0.6B_fsdp_colocated_2xGPU.py index 82349109c30..f104e3ad7c6 100644 --- a/tests/test_qwen3_0.6B_fsdp_colocated_2xGPU.py +++ b/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} " ) diff --git a/tests/test_qwen3_0.6B_fsdp_distributed.py b/tests/test_qwen3_0.6B_fsdp_distributed.py index 2727d3e6cb1..54230d7c50a 100644 --- a/tests/test_qwen3_0.6B_fsdp_distributed.py +++ b/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,6 +84,7 @@ def execute(): f"{U.get_default_wandb_args(__file__)} " f"{eval_args} " f"{sglang_args} " + f"{ci_args} " f"{misc_args} " ) @@ -85,6 +92,7 @@ def execute(): train_args=train_args, num_gpus=2 if FEW_GPU else 4, model_type=None, + train_script="train_async.py", ) diff --git a/tests/test_qwen3_30B_A3B.py b/tests/test_qwen3_30B_A3B.py index e1293b8354b..b476d7efa94 100644 --- a/tests/test_qwen3_30B_A3B.py +++ b/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} " )