GRPO/Distillation validation collapses multi-dataset metrics into a single accuracy
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
## Summary
When `data.validation` is configured as a list of multiple datasets, GRPO and Distillation correctly load and dispatch them per task during rollout, but the validation aggregator collapses all rewards into one `validation/accuracy` and one `validation/avg_length`. There is no way to read per-dataset progress from the wandb panel or driver log.
DPO is unaffected: its `val_dataloader` is a `dict[str, StatefulDataLoader]` and its `validate()` loops per dataset with a `validation-/` wandb prefix (`nemo_rl/algorithms/dpo.py:332-377`).
## Reproduction
Any GRPO or Distillation recipe with `data.validation` as a list of two or more datasets:
```yaml
data:
validation:
- dataset_name: gsm8k
split: test
- dataset_name: ResponseDataset
data_path: data/math500.parquet
```
Driver log shows both load:
```
- Loaded validation dataset gsm8k with 1319 samples.
- Loaded validation dataset data-math500 with 500 samples.
✓ Validation dataset loaded with 1819 samples.
```
The wandb panel shows only:
```
validation/accuracy = 0.5957
validation/avg_length = 512.8
```
That single number is a sample-weighted mean across both datasets, which is meaningless if the per-dataset accuracies diverge.
## Root cause
`nemo_rl/data/utils.py:setup_response_data:200-211` concatenates all val datasets into one `AllTaskProcessedDataset`, and the GRPO/Distillation `setup()` builds a single `StatefulDataLoader` from it. By the time `validate()` runs, dataset identity is lost and the aggregator has nothing to split on.
`val_batch["task_name"]` does carry per-sample task names, but routing them to `accuracy_` keys leaves dashboard metric names with hyphenated or otherwise awkward suffixes and asymmetry with DPO; the cleaner fix is to keep dataset identity at the dataloader level.
## Proposed change
Adopt DPO's data flow for response-style validation:
1. `nemo_rl/data/utils.py:setup_response_data` returns `Optional[dict[str, AllTaskProcessedDataset]]` for the validation set instead of a single concatenated dataset.
2. `nemo_rl/algorithms/grpo.py` and `nemo_rl/algorithms/distillation.py` `setup()` build `val_dataloader: dict[str, StatefulDataLoader]`. `validate()` iterates the dict, factors the inner per-dataset logic into a helper, and logs each dataset under `prefix=f"validation-{name}"` and `prefix=f"timing/validation-{name}"`. The per-dataset wandb logging is moved inside `validate()` to avoid double-prefixing at the caller.
3. While in the same block, `max_val_samples` is widened to `NotRequired[int]`; when absent or `None` the full `val_dataloader` is iterated (per dataset). The Distillation validation truncation also switches from ceiling division to floor division so it matches GRPO.
4. Entry scripts (`run_grpo.py`, `run_distillation.py`, `run_grpo_sliding_puzzle.py`, `examples/nemo_gym/run_grpo_nemo_gym.py`) are updated to thread the new dict type. The puzzle runner wraps its single iterable val dataset under `{task_name: dataset}` before passing to `setup()`. The NeMo-Gym runner sums lengths across the dict for its existing `max_val_samples` / `val_batch_size` derivation and writes one `trajectory_collection_.jsonl` per dataset.
5. DPO is not touched.
## Backwards compatibility
The aggregated `validation/accuracy` and `validation/avg_length` wandb keys are no longer emitted; metrics surface only under the per-dataset `validation-/` prefix. Dashboards relying on the old keys need to be re-pointed at `validation-/accuracy` and `validation-/avg_length`.
The internal `val_metrics["accuracy"]` return value is preserved as the macro-mean across datasets so `save_state["val_reward"]` and `checkpointing.metric_name='val:accuracy'` keep working without configuration changes.
PR coming alongside this issue.
Contributor guide
Assessment
This issue has not been assessed yet.