test metrics

Signed-off-by: Ahmad Kiswani <[email protected]>

Author: ahmadki

tests/test_suites/llm/grpo-qwen3-next-80ba3b-8n8g-megatron.sh

@@ -32,10 +32,9 @@ uv run examples/run_grpo_math.py \
 uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
 
 # Only run metrics if the target step is reached
-# TODO(ahmadki): set metrics
-# if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
-#     uv run tests/check_metrics.py $JSON_METRICS \
-#         'data["train/token_mult_prob_error"]["30"] < 1.1' \
-#         'data["train/reward"]["30"] > 0.43' \
-#         'mean(data["timing/train/total_step_time"], -6, -1) < 220'
-# fi
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'data["train/token_mult_prob_error"]["30"] < 1.1' \
+        'data["train/reward"]["30"] > 0.43' \
+        'mean(data["timing/train/total_step_time"], -6, -1) < 305'
+fi

tests/test_suites/llm/sft-qwen3-next-80ba3b-8n8g-megatron.sh

@@ -4,10 +4,10 @@ source $SCRIPT_DIR/common.env
 
 # ===== BEGIN CONFIG =====
 NUM_NODES=8
-STEPS_PER_RUN=80
-MAX_STEPS=80
+STEPS_PER_RUN=200
+MAX_STEPS=200
 NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
-NUM_MINUTES=30
+NUM_MINUTES=120
 # ===== END CONFIG =====
 
 exit_if_max_steps_reached
@@ -32,9 +32,9 @@ uv run examples/run_sft.py \
 uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
 
 # Only run metrics if the target step is reached
-# # TODO(ahmadki): set metrics
-# if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
-#     uv run tests/check_metrics.py $JSON_METRICS \
-#         'data["train/loss"]["80"] < 0.301' \
-#         'data["validation/val_loss"]["80"] < 0.304'
-# fi
+if [[ $(jq 'to_entries | .[] | select(.key == "train/loss") | .value | keys | map(tonumber) | max' $JSON_METRICS) -ge $MAX_STEPS ]]; then
+    uv run tests/check_metrics.py $JSON_METRICS \
+        'data["train/loss"]["200"] < 0.3' \
+        'data["validation/val_loss"]["200"] < 0.3' \
+        'mean(data["timing/train/total_step_time"], -6, -1) < 20'
+fi