NVIDIA-NeMo / NVIDIA-NeMo/RL

refactor(algorithms): align max_num_steps / max_num_epochs semantics with NeMo Automodel

Open
#2,511 0 comments 0 reactions 1 assignee Claimed by @terrykong View on GitHub
community-request Feature waiting-on-maintainers
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## Summary

Adopt the **"step overrides epoch"** convention (with `-1` sentinel) used by NeMo Automodel — and by the broader HuggingFace ecosystem that most NeMo-RL users come from — so that one knob is unambiguously primary instead of relying on `min(max_num_epochs * len(dataloader), max_num_steps)` with magic-number sentinels (`1000000`).

## Motivation

### 1. Consistency with NeMo Automodel (NVIDIA's sibling training stack)

NeMo Automodel's [`StepScheduler`](https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/training/step_scheduler.py) already resolves training duration as:

- `max_steps` provided ⇒ step-driven (overrides `num_epochs`)
- `max_steps` absent / `None` ⇒ derived from `num_epochs × len(dataloader)`

Users frequently move between Automodel (SFT/pretraining) and NeMo-RL (RLHF), and having two NVIDIA training stacks disagree on basic loop semantics is a recurring source of confusion.

### 2. The HuggingFace ecosystem convention

| Framework | Sentinel | Interaction |
|---|---|---|
| HF Transformers `TrainingArguments` | `max_steps = -1` | `max_steps > 0` overrides `num_train_epochs` |
| TRL (`SFTConfig`, `GRPOConfig`) | inherits from above | same |
| NeMo Automodel | `max_steps = None` | same |
| **NeMo-RL (current)** | magic number `1000000` | `min(epochs × len, steps)` with `AND` in loop |
| PyTorch Lightning | `-1` both | earliest-wins (`min`) |

NeMo-RL is the outlier among the trainers that RLHF practitioners actually use.

### 3. Internal inconsistency

Only `nemo_rl/algorithms/rm.py` currently honors `max_num_steps == -1` as "no cap" (rm.py:501-504). Every other algorithm (`sft`, `grpo`, `dpo`, `distillation`) requires a magic value like `max_num_steps: 1000000` to express "no step cap". Examples:

- `examples/configs/grpo_math_1B.yaml` → `max_num_steps: 1000000`
- `examples/configs/grpo_math_1B_megatron.yaml` → `max_num_steps: 1000000`
- `examples/configs/sft_openmathinstruct2.yaml` → `max_num_steps: 1000000`

This violates `CONTRIBUTING.md`'s "be as explicit as possible" and `.claude/skills/config-conventions` guidance against arbitrary defaults / magic values.

## Proposed Behavior

```python
# nemo_rl/algorithms/utils.py (new helper, shared across algos)
def resolve_train_horizon(
max_num_epochs: int,
max_num_steps: int,
iters_per_epoch: int,
) -> int:
"""HF/Automodel-style: `max_num_steps > 0` overrides `max_num_epochs`.

- `max_num_steps > 0`: step-driven; `max_num_epochs` is ignored (allowed to
loop through the dataloader multiple times).
- `max_num_steps == -1`: epoch-driven; total = `max_num_epochs * iters_per_epoch`.
Requires `max_num_epochs > 0`.
"""
```

The training loop predicate simplifies from:

```python
while current_epoch < max_num_epochs and total_steps < max_num_steps:
```

to a single-axis condition driven by the resolved horizon.

## Affected Files

- `nemo_rl/algorithms/sft.py` (L189-190, L418-420, L469-473, L642)
- `nemo_rl/algorithms/dpo.py` (L244-245, L573-574, L600-602, L786)
- `nemo_rl/algorithms/grpo.py` (L567-568, L1427, L1881-1884, L2254; async path L2684/2976/3263)
- `nemo_rl/algorithms/distillation.py` (L392-393, L444-445)
- `nemo_rl/algorithms/rm.py` (extend existing `-1` handling to `max_num_epochs` for symmetry)
- New helper module (or `nemo_rl/algorithms/utils.py`) + unit test

## Exemplar Migration

| Config | Before | After |
|---|---|---|
| `examples/configs/grpo_math_1B.yaml` | `max_num_steps: 1000000` | `max_num_steps: -1` |
| `examples/configs/grpo_math_1B_megatron.yaml` | `max_num_steps: 1000000` | `max_num_steps: -1` |
| `examples/configs/sft_openmathinstruct2.yaml` | `max_num_steps: 1000000` | `max_num_steps: -1` |
| `examples/configs/sft.yaml` | `epochs: 1, steps: 60` | unchanged (still step-driven) |
| `examples/configs/dpo.yaml` | `epochs: 1, steps: 150` | unchanged |
| `examples/configs/rm.yaml` | `epochs: 1, steps: -1` | unchanged |

Recipe yamls under `examples/configs/recipes/**/*.yaml` inherit via `defaults:` and require no changes unless they explicitly override these keys.

## Backward Compatibility

- **Positive values on both axes**: behavior changes subtly — currently `min(...)`, after this PR `max_num_steps` wins. Audit all positive-positive yamls and confirm they were already intended as step-driven (a manual review of the table above suggests all current exemplars are unaffected).
- **Magic-number sentinel (`1000000`)**: will start training for many more steps than before. Migrate to `-1` in the same PR — this is the only true behavior-breaking case.
- A deprecation warning could be emitted when a user passes `max_num_steps >= 1_000_000` with `max_num_epochs > 0`, suggesting `-1` instead.

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.