fix: Fix nanov3 mcore test (#2635)

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

Author: yfw

pyproject.toml

@@ -159,6 +159,8 @@ mcore = [
   "transformer-engine[pytorch,core_cu13] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1",
   "megatron-core",
   "megatron-bridge",
+  "mamba-ssm",
+  "causal-conv1d",
   "nvidia-modelopt[torch]; sys_platform != 'darwin'",
   "onnxscript",
   # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular)

tests/test_suites/llm/grpo-nanov3-30BA3B-2n8g-megatron-pack-cp.sh

@@ -36,8 +36,8 @@ uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_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 \
         'median(data["train/token_mult_prob_error"]) < 1.1' \
-        'mean(data["train/gen_kl_error"]) < 0.02' \
-        'max(data["train/reward"]) > 0.05'
+        "data['train/token_mult_prob_error']['$MAX_STEPS'] < 1.1" \
+        'mean(data["train/gen_kl_error"]) < 0.02'
 
     # Clean up checkpoint directory after successful run to save space.
     rm -rf "$CKPT_DIR"