fix(data): add 'aime2026' to AIMEEvalDataConfig literal (#2469 follow-up)

[qiaochuz-nv]

Summary

One-line fix to the AIMEEvalDataConfig.dataset_name Literal so that
data.dataset_name=aime2026 is accepted by MasterConfig validation.

- dataset_name: Literal["aime2024", "aime2025"]
+ dataset_name: Literal["aime2024", "aime2025", "aime2026"]

Root cause

PR #2469 ("feat: add AIME-2026 benchmark") landed two halves of the
feature, but missed a third:

1. ✅ nemo_rl/data/datasets/eval_datasets/aime.py — the AIME-2026
   dataset loader.
2. ✅ nemo_rl/data/datasets/eval_datasets/__init__.py — the dispatcher
   was extended:

   elif dataset_name in ["aime2024", "aime2025", "aime2026"]:
       base_dataset = AIMEDataset(...)

3. ❌ nemo_rl/data/__init__.py — the AIMEEvalDataConfig TypedDict
   that constrains data.dataset_name was not updated:

   class AIMEEvalDataConfig(TypedDict):
       max_input_seq_length: int
       dataset_name: Literal["aime2024", "aime2025"]   # missing "aime2026"
       ...

examples/run_eval.py does MasterConfig(**config) (pydantic). With
the literal mismatch, pydantic rejects aime2026 against
AIMEEvalDataConfig, then walks every other arm of the EvalDataConfigType
union and rejects each in turn — surfacing as 9 validation errors before
any evaluation runs. The dispatcher branch is reachable code, but the
user-facing config can never get there.

Why the existing unit test in PR #2469 did not catch this

tests/unit/data/datasets/test_eval_dataset.py::test_aime_dataset was
added in PR #2469 and parametrizes over ["aime2024", "aime2025", "aime2026"],
but it does not exercise this code path:

1. The test is decorated with @pytest.mark.skip(reason="dataset download is flaky"),
   so it never runs in CI for any variant.
2. Even if the skip were removed, it calls load_eval_dataset(data_config)
   directly with a plain dict. That path hits the dispatcher
   (eval_datasets/__init__.py, which was updated by feat: add AIME-2026 benchmark. #2469) and never
   builds a MasterConfig/AIMEEvalDataConfig. So the literal mismatch
   in nemo_rl/data/__init__.py is invisible to this test.

The literal is only enforced at the MasterConfig(**config) boundary in
examples/run_eval.py, which the unit test bypasses.

Repro (before fix)

Container nemo-rl-nightly-20260521.sqsh (or any post-#2469 image),
single H100 on EOS interactive partition:

cd /opt/nemo-rl
uv run examples/run_eval.py \
  --config examples/configs/evals/math_eval.yaml \
  data.dataset_name=aime2026 \
  data.prompt_file=examples/prompts/cot.txt \
  generation.model_name=Qwen/Qwen3-0.6B \
  tokenizer.name=Qwen/Qwen3-0.6B \
  generation.num_prompts_per_step=2 \
  generation.max_new_tokens=64 \
  generation.vllm_cfg.max_model_len=1024 \
  generation.vllm_cfg.gpu_memory_utilization=0.70 \
  generation.vllm_cfg.enforce_eager=true \
  env.math.num_workers=2 \
  cluster.gpus_per_node=1 \
  eval.save_path=/tmp/aime2026_repro/results

Before fix — observed

Loaded configuration from: /opt/nemo-rl/examples/configs/evals/math_eval.yaml
...
Traceback (most recent call last):
  File "/opt/nemo-rl/examples/run_eval.py", line 135, in <module>
    main()
  File "/opt/nemo-rl/examples/run_eval.py", line 94, in main
    config = MasterConfig(**config)
pydantic_core._pydantic_core.ValidationError: 9 validation errors for MasterConfig
data.MMLUEvalDataConfig.dataset_name
  Input should be 'mmlu', 'mmlu_AR-XY', ...
    [type=literal_error, input_value='aime2026', input_type=str]
data.MMLUProEvalDataConfig.dataset_name
  Input should be 'mmlu_pro' [type=literal_error, input_value='aime2026', input_type=str]
data.AIMEEvalDataConfig.dataset_name
  Input should be 'aime2024' or 'aime2025'
    [type=literal_error, input_value='aime2026', input_type=str]
data.GPQAEvalDataConfig.dataset_name
  Input should be 'gpqa' or 'gpqa_diamond' ...
data.MathEvalDataConfig.dataset_name
  Input should be 'math' or 'math500' ...
data.MMAUEvalDataConfig.dataset_name
  Input should be 'mmau' or 'TwinkStart/MMAU' ...
data.LocalMathEvalDataConfig.problem_key   Field required
data.LocalMathEvalDataConfig.solution_key  Field required
data.LocalMathEvalDataConfig.file_format   Field required

Run never reaches dataset loading.

After fix — observed

Same command with the one-line literal fix applied (verified by patching
nemo_rl/data/__init__.py in-container to confirm the literal is the
only blocker):

[daily-pr] patching AIMEEvalDataConfig literal in /opt/nemo-rl/nemo_rl/data/__init__.py
129:    dataset_name: Literal["aime2024", "aime2025", "aime2026"]
...
Loaded configuration from: /opt/nemo-rl/examples/configs/evals/math_eval.yaml
Applied CLI overrides
Final config:
MasterConfig(
  eval={'metric': 'pass@k', ...},
  generation={'backend': 'vllm', ..., 'model_name': 'Qwen/Qwen3-0.6B', ...},
  tokenizer={'name': 'Qwen/Qwen3-0.6B', ...},
  data={'max_input_seq_length': 1024,
        'dataset_name': 'aime2026',
        'prompt_file': 'examples/prompts/cot.txt',
        'system_prompt_file': None},
  env={'math': {'num_workers': 2}},
  cluster={'gpus_per_node': 1, 'num_nodes': 1},
  ...
)
Using tokenizer's default chat template
Setting up data...

MasterConfig now validates and accepts aime2026 cleanly, the
dispatcher reaches AIMEDataset(variant="2026", ...), and execution
proceeds into setup_data — exactly the behavior PR #2469 intended.

(Full end-to-end completion of the eval was blocked downstream by an
unrelated EOS lustre inode-quota issue at the dataset-download step;
that is a cluster fs problem, not part of this regression.)

Detected by

NeMo daily-PR impact pipeline — auto-generated regression test
test_eval_aime2026_daily_pr covering RL PR #2469.

Test plan

• MasterConfig(data={..., "dataset_name": "aime2026"}) no longer
  raises pydantic.ValidationError.
• examples/run_eval.py --config .../math_eval.yaml data.dataset_name=aime2026 ...
  passes the validation step (verified above).
• No behavior change for aime2024 / aime2025.

Signed-off-by: Qiaochu Zhu qiaochuz@nvidia.com

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[qiaochuz-nv]

/ok to test d530bdf

[yuki-97]

thanks for the fix!

[qiaochuz-nv]

Retrospective QA-fix tracking issue: #4053. It records the fix authored by @qiaochuz-nv and delivered by this merged PR.