test: add deepscaler tests + pipe-clean configs + fix eval for deepscaler

docs/guides/grpo-deepscaler.md

@@ -5,12 +5,12 @@ This guide explains how to use NeMo RL to train long Chain of Thought (CoT) reas
 
 ## Train the Model
 We follow the DeepScaleR recipe and train the model in three stages. In the first stage, we train with an 8K context window. In the second stage, we train with a 16K context window. In the third stage, we train with a 24K context window.
-To train the model using NeMo RL, use the `examples/configs/grpo-deepscaler-1.5b-8K.yaml` config file. This file closely matches the experiment settings in the original DeepScaleR recipe. We then train with `examples/configs/grpo-deepscaler-1.5b-16K.yaml` and `examples/configs/grpo-deepscaler-1.5b-24K.yaml` for the second and third stages, respectively.
+To train the model using NeMo RL, use the `examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K.yaml` config file. This file closely matches the experiment settings in the original DeepScaleR recipe. We then train with `examples/configs/recipes/llm/grpo-deepscaler-1.5b-16K.yaml` and `examples/configs/recipes/llm/grpo-deepscaler-1.5b-24K.yaml` for the second and third stages, respectively.
 
 ```sh
-uv run examples/run_grpo_math.py --config=examples/configs/grpo-deepscaler-1.5b-8K.yaml
-uv run examples/run_grpo_math.py --config=examples/configs/grpo-deepscaler-1.5b-16K.yaml policy.model_name=/path/to/8K/checkpoint/hf
-uv run examples/run_grpo_math.py --config=examples/configs/grpo-deepscaler-1.5b-24K.yaml policy.model_name=/path/to/16K/checkpoint/hf
+uv run examples/run_grpo_math.py --config=examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K.yaml
+uv run examples/run_grpo_math.py --config=examples/configs/recipes/llm/grpo-deepscaler-1.5b-16K.yaml policy.model_name=/path/to/8K/checkpoint/hf
+uv run examples/run_grpo_math.py --config=examples/configs/recipes/llm/grpo-deepscaler-1.5b-24K.yaml policy.model_name=/path/to/16K/checkpoint/hf
 ```
 
 At the end of each stage, you need to specify the Hugging Face checkpoint to continue training with. To get this checkpoint, we convert a model checkpoint to a Hugging Face checkpoint with the following command:
@@ -19,7 +19,7 @@ At the end of each stage, you need to specify the Hugging Face checkpoint to con
 uv run examples/converters/convert_dcp_to_hf.py --config=results/grpo-deepscaler-1.5b-8K/step_240/config.yaml --dcp-ckpt-path=results/grpo-deepscaler-1.5b-8K/step_240/policy/weights --hf-ckpt-path=results/grpo-deepscaler-1.5b-8K/step_240/hf
 ```
 
-When running the next command, we use the Hugging Face checkpoint as the initial checkpoint. We train with an 8K context window for 240 steps, a 16K context window for 290 steps, and a 24K context window for 50 steps. The 8K and 16K steps can be run on a single 8XH100 80GB node, while the 24K step requires four nodes. If you're running on 8XA100 80GB, you will need at least 1 node for 8K training and four nodes for 16-24k training.
+When running the next command, we use the Hugging Face checkpoint as the initial checkpoint. We train with an 8K context window for 240 steps, a 16K context window for 290 steps, and a 24K context window for 50 steps. We run all experiments on a single 8XH100 80GB node. If you're running on 8XA100 80GB, you will need at least 1 node for 8K training and 2 nodes for 16-24k training.
 
 ## Training Curve
 When using the above commands, we get the following training curve:

examples/configs/grpo-deepscaler-1.5b-16K.yaml → examples/configs/recipes/llm/grpo-deepscaler-1.5b-16K.yaml

@@ -8,6 +8,7 @@ loss_fn:
 
 policy:
   max_total_sequence_length: 16384
+  logprob_batch_size: 2
 
   dtensor_cfg:
     enabled: true

examples/configs/grpo-deepscaler-1.5b-24K.yaml → examples/configs/recipes/llm/grpo-deepscaler-1.5b-24K.yaml

@@ -8,6 +8,7 @@ loss_fn:
 
 policy:
   max_total_sequence_length: 24576
+  logprob_batch_size: 2
 
   dtensor_cfg:
     enabled: true
@@ -44,9 +45,6 @@ policy:
       gpu_memory_utilization: 0.8
       enforce_eager: True
       max_model_len: ${policy.max_total_sequence_length}
-      # For most cases, use "dummy" to load the initial weights, since they will be overwritten during refit
-      # For Gemma models, we need to use "auto" due to a vllm bug
-      load_format: dummy
 
 cluster:
   gpus_per_node: 8

examples/configs/grpo-deepscaler-1.5b-8K.yaml → examples/configs/recipes/llm/grpo-deepscaler-1.5b-8K.yaml

@@ -103,9 +103,6 @@ policy:
       gpu_memory_utilization: 0.6
       max_model_len: ${policy.max_total_sequence_length}
       enforce_eager: True
-      # For most cases, use "dummy" to load the initial weights, since they will be overwritten during refit
-      # For Gemma models, we need to use "auto" due to a vllm bug
-      load_format: dummy
     colocated:
       # true: generation shares training GPUs
       # false: uses dedicated generation resources

nemo_rl/evals/eval.py

@@ -427,7 +427,7 @@ def _save_evaluation_data_to_json(evaluation_data, master_config, save_path):
         "model_name": master_config["generation"]["model_name"],
         "dataset_name": master_config["data"]["dataset_name"],
         "metric": master_config["eval"]["metric"],
-        "pass_k_value": master_config["eval"]["pass_k_value"],
+        "k_value": master_config["eval"]["k_value"],
         "num_tests_per_prompt": master_config["eval"]["num_tests_per_prompt"],
         "temperature": master_config["generation"]["temperature"],
         "top_p": master_config["generation"]["top_p"],

tests/test_suites/llm/grpo-deepscaler-1.5b-16K.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=1
+STEPS_PER_RUN=30
+MAX_STEPS=30
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=240
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Use checkpoint created from the 8K checkpoint in grpo-deepscaler-1.5b-8K.sh
+if [[ -z "$CACHED_MODEL_PATH" ]]; then
+    echo "Need to set CACHED_MODEL_PATH to the path to the trained 8K checkpoint"
+    exit 1
+fi
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_grpo_math.py \
+    --config $CONFIG_PATH \
+    policy.model_name=$CACHED_MODEL_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+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 \
+        'mean(data["train/token_mult_prob_error"]) < 1.1' \
+        "data['train/token_mult_prob_error']['$MAX_STEPS'] < 1.1"
+fi
+
+# TODO: enable in subsequent PR to do a quick accuracy check
+## Convert 16k checkpoint
+#uv run examples/converters/convert_dcp_to_hf.py \
+#  --config=$CKPT_DIR/step_${MAX_STEPS}/config.yaml \
+#  --dcp-ckpt-path=$CKPT_DIR/step_${MAX_STEPS}/policy/weights \
+#  --hf-ckpt-path=$CKPT_DIR/grpo-deepscaler-16k-${MAX_STEPS}-hf
+#
+## Run eval
+#uv run examples/run_eval.py \
+#    generation.model_name=$CKPT_DIR/grpo-deepscaler-16k-${MAX_STEPS}-hf \
+#    data.prompt_file=examples/prompts/cot.txt \
+#    generation.vllm_cfg.max_model_len=32768 2>&1 | tee ${RUN_LOG}.aime-16k
+#
+#cat ${RUN_LOG}.aime-16k       | grep "score=" | sed 's/.*score=\([^ ]*\).*/{"score": \1}/' > ${RUN_LOG}-16k-metric.json
+# 
+#uv run tests/check_metrics.py ${RUN_LOG}-16k-metric.json \
+#  'data["score"] >= 0.25' \
+

tests/test_suites/llm/grpo-deepscaler-1.5b-24K.sh

@@ -0,0 +1,65 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=4
+STEPS_PER_RUN=30
+MAX_STEPS=30
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=240
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Use checkpoint created from the 16K checkpoint in grpo-deepscaler-1.5b-16K.sh
+if [[ -z "$CACHED_MODEL_PATH" ]]; then
+    echo "Need to set CACHED_MODEL_PATH to the path to the trained 16K checkpoint"
+    exit 1
+fi
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_grpo_math.py \
+    --config $CONFIG_PATH \
+    policy.model_name=$CACHED_MODEL_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+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 \
+        'mean(data["train/token_mult_prob_error"]) < 1.1' \
+        "data['train/token_mult_prob_error']['$MAX_STEPS'] < 1.1"
+fi
+
+# TODO: enable in subsequent PR to do a quick accuracy check
+## Convert 24k checkpoint
+#uv run examples/converters/convert_dcp_to_hf.py \
+#  --config=$CKPT_DIR/step_${MAX_STEPS}/config.yaml \
+#  --dcp-ckpt-path=$CKPT_DIR/step_${MAX_STEPS}/policy/weights \
+#  --hf-ckpt-path=$CKPT_DIR/grpo-deepscaler-24k-${MAX_STEPS}-hf
+#
+## Run eval
+#uv run examples/run_eval.py \
+#    generation.model_name=$CKPT_DIR/grpo-deepscaler-24k-${MAX_STEPS}-hf \
+#    data.prompt_file=examples/prompts/cot.txt \
+#    generation.vllm_cfg.max_model_len=32768 2>&1 | tee ${RUN_LOG}.aime-24k
+#
+#cat ${RUN_LOG}.aime-24k       | grep "score=" | sed 's/.*score=\([^ ]*\).*/{"score": \1}/' > ${RUN_LOG}-24k-metric.json
+# 
+#uv run tests/check_metrics.py ${RUN_LOG}-24k-metric.json \
+#  'data["score"] >= 0.25' \
+

tests/test_suites/llm/grpo-deepscaler-1.5b-8K.sh

@@ -0,0 +1,67 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=1
+STEPS_PER_RUN=40
+MAX_STEPS=40
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))  # Round up
+NUM_MINUTES=240
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+# Run the experiment
+cd $PROJECT_ROOT
+uv run examples/run_grpo_math.py \
+    --config $CONFIG_PATH \
+    grpo.max_num_steps=$MAX_STEPS \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=True \
+    logger.wandb.project=nemo-rl \
+    logger.wandb.name=$EXP_NAME \
+    logger.monitor_gpus=True \
+    logger.tensorboard_enabled=True \
+    checkpointing.enabled=True \
+    checkpointing.checkpoint_dir=$CKPT_DIR \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+# Convert tensorboard logs to json
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+# Only run metrics if the target step is reached
+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 \
+        'mean(data["train/token_mult_prob_error"]) < 1.1' \
+        "data['train/token_mult_prob_error']['$MAX_STEPS'] < 1.1"
+fi
+
+# TODO: enable in subsequent PR to do a quick accuracy check
+## Convert 8k checkpoint
+#uv run examples/converters/convert_dcp_to_hf.py \
+#  --config=$CKPT_DIR/step_${MAX_STEPS}/config.yaml \
+#  --dcp-ckpt-path=$CKPT_DIR/step_${MAX_STEPS}/policy/weights \
+#  --hf-ckpt-path=$CKPT_DIR/grpo-deepscaler-8k-${MAX_STEPS}-hf
+#
+## Run eval
+#uv run examples/run_eval.py \
+#    generation.model_name=$CKPT_DIR/grpo-deepscaler-8k-${MAX_STEPS}-hf \
+#    data.prompt_file=examples/prompts/cot.txt \
+#    generation.vllm_cfg.max_model_len=32768 2>&1 | tee ${RUN_LOG}.aime-8k
+#
+#cat ${RUN_LOG}.aime-8k       | grep "score=" | sed 's/.*score=\([^ ]*\).*/{"score": \1}/' > ${RUN_LOG}-8k-metric.json
+# 
+#uv run tests/check_metrics.py ${RUN_LOG}-8k-metric.json \
+#  'data["score"] >= 0.25' \
+#
+##uv run examples/run_eval.py \
+##    generation.model_name=deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
+##    data.prompt_file=examples/prompts/cot.txt \
+##    generation.vllm_cfg.max_model_len=32768 2>&1 | tee ${RUN_LOG}.aime-baseline
+#
+##cat ${RUN_LOG}.aime-baseline | grep "score=" | sed 's/.*score=\([^ ]*\).*/{"score": \1}/' > ${RUN_LOG}-baseline-metric.json
+#
+##uv run tests/check_metrics.py ${RUN_LOG}-baseline-metric.json \
+##  'data["score"] == 0.2' \

tests/test_suites/nightly.txt

@@ -13,6 +13,11 @@ tests/test_suites/llm/grpo-qwen2.5-7b-instruct-4n8g-fsdp2tp4sp.v3.sh
 # Functional 32b run
 tests/test_suites/llm/grpo-qwen2.5-32b-32n8g-fsdp2tp8sp-actckpt.v3.sh
 
+# Deepscaler (short tests)
+tests/test_suites/llm/grpo-deepscaler-1.5b-16K.sh
+tests/test_suites/llm/grpo-deepscaler-1.5b-24K.sh
+tests/test_suites/llm/grpo-deepscaler-1.5b-8K.sh
+
 #######
 # SFT #
 #######

tests/unit/test_recipes_and_test_suites.py

@@ -283,6 +283,8 @@ def test_all_recipes_can_merge_configs_with_base_config(
 ):
     from omegaconf import OmegaConf
 
+    from nemo_rl.utils.config import load_config
+
     base_yaml = os.path.join(project_root, algo_base_yaml)
     base_config = OmegaConf.load(base_yaml)
     # Would result in an error if we couldn't merge our config with the recipe's config
@@ -293,7 +295,7 @@ def test_all_recipes_can_merge_configs_with_base_config(
             #  test_all_recipes_start_with_algo_hyphen()
             continue
         recipe_yaml_path = os.path.join(recipes_dir, recipe_yaml)
-        recipe_config = OmegaConf.load(recipe_yaml_path)
+        recipe_config = load_config(recipe_yaml_path)
         OmegaConf.set_struct(recipe_config, True)
         # This will raise a error if the config can't be merged
         print(f"Merging {recipe_yaml} with {base_yaml}")