gsm8k-checkpoint-eval

Files (1)

• SKILL.md

SKILL.md (first lines)

# Math Checkpoint Evaluation

Use this workflow for native Nemotron Diffusion NeMo-RL Megatron checkpoints. The evaluation has two major steps:

1. Convert the raw Megatron checkpoint step directory to Hugging Face format.
2. Evaluate the converted Hugging Face checkpoint on GSM8K, AIME, or related math benchmarks with the standalone evaluator or the `yonggon-rl` NemoSkills pipeline.

Do not run conversion directly on the login node. Run conversion inside the container specified below so CUDA, Transformer Engine, cuDNN, and C++ runtime dependencies are consistent.

## Step 1: Convert Checkpoint

Required inputs:

- `STEP_DIR`: raw checkpoint step directory containing `config.yaml` and `policy/weights/iter_0000000`.
- `OUT`: destination Hugging Face checkpoint directory.

Use this container for conversion:

```bash
CONTAINER=/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/sfawzy/nemo-rl-nightly.sqsh
```

Use the repo wrapper for conversion. Do not inline or recreate the converter command from `tools/nemotron_diffusion/convert_nemotron_diffusion_checkpoint_to_hf.sh`; `convert_checkpoint.sh` owns the native Nemotron Diffusion conversion settings.

For diffuGRPO/native NeMo-RL checkpoints, set the same uv tag/env/cache family used by training. `convert_checkpoint.sh` does not derive these paths from `ENV_TAG` by itself, so pass both `ENV_TAG` and the `NEMOTRON_UV_*` overrides explicitly:

```bash
cd /home/snorouzi/diffusion_RL/RL

ENV_TAG=mb_3rdparty_sglagn_local_fork \
UV_PROJECT_ENVIRONMENT=/lustre/fsw/portfolios/coreai/users/snorouzi/nemorl_uv_driver_envs/diffusion_RL_RL_diffu_grpo \
UV_CACHE_DIR=/lustre/fsw/portfolios/coreai/users/snorouzi/uv_cache_diffu_grpo \
NEMOTRON_UV_PROJECT_ENVIRONMENT=/lustre/fsw/portfolios/coreai/users/snorouzi/nemorl_uv_driver_envs/diffusion_RL_RL_mb_3rdparty_sglagn_local_fork \
NEMOTRON_UV_CACHE_DIR=/lustre/fsw/portfolios/coreai/users/snorouzi/uv_cache_diffu_grpo \
NEMOTRON_HF_HOME=/lustre/fsw/portfolios/coreai/users/snorouzi/hf_home \
STEP_DIR=/path/to/run/checkpoints/step_N \
OUT=/lustre/fsw/portfolios/coreai/users/snorouzi/checkpoints/<name>_step_N_hf \
./convert_checkpoint.sh
```

When running under Slurm, load the container above and run the same wrapper command inside `/home/snorouzi/diffusion_RL/RL`. Example:

```bash
srun --container-image=/lustre/fsw/portfolios/coreai/projects/coreai_dlalgo_llm/users/sfawzy/nemo-rl-nightly.sqsh \
  --container-mounts=/home/snorouzi:/home/snorouzi,/lustre:/lustre \
  --container-workdir=/home/snorouzi/diffusion_RL/RL \
  bash -lc 'export ENV_TAG=diffu_grpo
export UV_PROJECT_ENVIRONMENT=/lustre/fsw/portfolios/coreai/users/snorouzi/nemorl_uv_driver_envs/diffusion_RL_RL_diffu_grpo
export UV_CACHE_DIR=/lustre/fsw/portfolios/coreai/users/snorouzi/uv_cache_diffu_grpo
export NEMOTRON_UV_PROJECT_ENVIRONMENT=$UV_PROJECT_ENVIRONMENT
export NEMOTRON_UV_CACHE_DIR=$UV_CACHE_DIR
export NEMOTRON_HF_HOME=/lustre/fsw/portfolios/coreai/users/snorouzi/hf_home
STEP_DIR=/path/to/run/checkpoints/step_N OUT=/lustre/fsw/portfolios/coreai/users/snorouzi/checkpoints/<name>_step_N_hf ./convert_checkpoint.sh'
```

Agents should not generate a full custom conversion script unless the user explicitly asks for one.

After conversion succeeds, verify the Hugging Face directory contains at least:

```bash
ls -l "${OUT}"/config.json \
      "${OUT}"/model.safetensors \
      "${OUT}"/tokenizer.json \
      "${OUT}"/configuration_ministral_dlm.py \
      "${OUT}"/modeling_ministral_dlm.py
```

## Step 2: Evaluate Converted Checkpoint

Default evaluation must be standlone No-Ray NemoRL validation.

### 1. Standalone No-Ray NeMoRL Validation Replication

Use this path when the user asks to reproduce the NeMoRL validation schedule, compare against training-time validation, or avoid Ray while keeping the NeMoRL prompt/generation/grading logic. Prefer the checked-in submit wrapper; do not reconstruct the full Slurm command by hand.

Set benchmark-specific budgets explicitly. The wrapper defaults may be tuned for recent AIME runs, so do not rely on defaults for GSM8K.

Recommended standalone budgets:

- `BENCHMARK=gsm8k`: `MAX_NEW_TOKENS=750`, `MAX_STEPS=BS` (e.g. `32` for `BS=32`), `CONTEXT_LENGTH=1024`.
- `BENCHMARK=aime2024`/`aime24` or `BENCHMARK=aime2025`/`aime25`: `MAX_NEW_TOKENS=8192`, `MAX_STEPS=BS` (e.g. `16` for `BS=16`), `CONTEXT_LENGTH=20480`.

CRITICAL — `MAX_STEPS` is the per-BLOCK denoising-step cap, NOT the total token budget. Always set `MAX_STEPS=BS` (equal to the block size). Do NOT set `MAX_STEPS=MAX_NEW_TOKENS`: a block of `BS` tokens can never need more than `BS` denoising steps (worst case = 1 token revealed per step), so a value larger than `BS` is meaningless and writes a misleading `max_steps` into `dllm_config.yaml`. The latest validated AIME24/AIME25 runs used `BS=16` with `MAX_STEPS=16`.

```bash
cd /home/snorouzi/diffusion_RL/RL

CKPT=/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_genai/users/snorouzi/checkpoints/<name>_step_N_hf \
TOKENIZER=/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_genai/users/snorouzi/checkpoints/Nemotron-Labs-Diffusion-3B \
BENCHMARK=gsm8k \
ALG=FastDiffuser \
BS=32 \
TEMP=0.0 \
MAX_NEW_TOKENS=750 \
MAX_STEPS=32 \
CONTEXT_LENGTH=1024 \
TAG=<name>_step_N_nemorl_rep_fastdiffuser_b32_t1 \
./submit_standalone_gsm8k_eval.sh
```

Key parameters:

- `BENCHMARK=gsm8k` by default. Use `BENCHMARK=aime2024`/`aime24` or `BENCHMARK=aime2025`/`aime25` to run the same standalone NeMoRL-style evaluator on AIME.
- `ALG=FastDiffuser` for diffusion/FastDiffuser eval.
- `ALG=LinearSpec` for linear speculation eval.
- `BS=16` or `BS=32` for block size.
- `TEMP=0.0` or `TEMP=1.0` confirm with the user.
- `MAX_NEW_TOKENS`, `MAX_STEPS`, and `CONTEXT_LENGTH` must be set per benchmark. Set `MAX_STEPS=BS` (the per-block denoising cap), NOT `MAX_NEW_TOKENS`. Use `MAX_NEW_TOKENS=750`, `MAX_STEPS=BS`, `CONTEXT_LENGTH=1024` for GSM8K and `MAX_NEW_TOKENS=8192`, `MAX_STEPS=BS`, `CONTEXT_LENGTH=20480` for AIME.
- AIME standalone eval loads the Nemo Skills AIME JSONLs from `/lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_genai/users/snorouzi/eval_data/nemo_skills_aime/{aime24,aime25}/test.jsonl` by default. Override with `NEMO_SKILLS_AIME_DATA_DIR` if needed; do not commit benchmark JSONLs into the repo.
- CoT prompting MUST be used (see the mandatory note at the top of Step 2): always set `PROMPT_FILE=/home/snorouzi/diffusion_RL/RL/examples/prompts/cot.txt`. Do NOT use `generic_math.txt` or `aime_no_cot.txt` — they are NOT the training prompt and lower accuracy. For chat-completions parity with Yonggon/NemoSkills you may add `GENERATION_API=chat_completions` and `TOP_P=0.95`, but to faithfully match NeMo-RL training-time prompt construction use `GENERATION_API=generate` (it applies the same chat template as training).
- `NUM_SAMPLES=-1` by default for the full benchmark; set `NUM_SAMPLES=16` for a smoke test.

Examples:

```bash
# GSM8K, FastDiffuser, NeMoRL validation-style budget
BENCHMARK=gsm8k ALG=FastDiffuser BS=32 TEMP=0 MAX_NEW_TOKENS=750 MAX_STEPS=32 CONTEXT_LENGTH=1024 TAG=<name>_gsm8k_fd_b32_t1 ./submit_standalone_gsm8k_eval.sh

# AIME 2024, FastDiffuser, Yonggon/NemoSkills-style high budget

From sajadn/diffusion_RL, Apache-2.0.

More from sajadn/diffusion_RL