NVIDIA-NeMo / NVIDIA-NeMo/RL

GRPO/Distillation validation collapses multi-dataset metrics into a single accuracy

Open
#2,501 0 comments 1 reaction 1 assignee Claimed by @terrykong View on GitHub
bug community-request Feature waiting-on-maintainers
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.