Skip to content

Add iteration count training limit - #239

Merged
tesslerc merged 2 commits into
NVlabs:mainfrom
carlosdp:codex/training-max-iterations
Jul 29, 2026
Merged

Add iteration count training limit#239
tesslerc merged 2 commits into
NVlabs:mainfrom
carlosdp:codex/training-max-iterations

Conversation

@carlosdp

@carlosdp carlosdp commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Right now, the only way to specify when we want training to automatically end is by specifying the maximum environment steps. This is an not an intuitive number for us humans to reason about in the age of massive parallel environment RL (in practice, nobody would want just part of a rollout if max-steps isn't a multiple of env_count x rollout_steps).

Iteration count is much easier to reason about, and is often supported in other RL frameworks. This PR adds this as an option via CLI flag, as an alternative to max-steps, while leaving max-steps in place for those coding agents that like to show off they can reason about bigger numbers!

  • add --training-max-iterations as an alternative to
    --training-max-steps
  • define one iteration as a complete rollout across every environment and rank,
    followed by the configured optimization updates
  • reject commands that specify both duration flags
  • preserve the existing step-based default and whole-iteration floor behavior
  • forward iteration limits through the Slurm launcher and document resume
    semantics

Behavior

--training-max-iterations N runs to an absolute target of exactly N complete
trainer epochs; it never stops partway through a rollout. When omitted,
--training-max-steps retains its current behavior:

max_iterations = training_max_steps // total_parallel_envs // agent.num_steps

On resume, the selected limit remains the absolute total target rather than an
additional number of iterations.

@carlosdp
carlosdp force-pushed the codex/training-max-iterations branch from 5bbbe8b to ac0b219 Compare July 28, 2026 18:57
@tesslerc

Copy link
Copy Markdown
Collaborator

Thanks! Only ask is to avoid backward compatibility whenever possible.
Usually, we can simply recreate the config file, which allows us to avoid messy fallbacks for backward compat logic.
In each release, we simply re-create all the configs for all the released models. Ensuring it keeps working, but also properly aligned with the latest codebase.

The logic in training_limits.py can be moved directly into BaseAgentConfig. Since it's just 1 function, we can have BaseAgentConfig support a resolve_max_epochs(total_envs) function.

Then BaseAgent can simply call self.max_epochs = self.config.resolve_max_epochs(self._total_envs).

We can then also avoid backward compatibility clauses such as training_max_iterations = getattr(config, "training_max_iterations", None)

Also in train_agent.py we should revert from max_iterations = getattr(args, "training_max_iterations", None) to max_iterations = args.training_max_iterations.

Documentation:

Can you add a short note to configuration.rst or developer_tips.rst:

--training-max-iterations N sets the absolute number of rollout/optimization iterations for the experiment. Each iteration collects one num_steps rollout from every environment and performs the configured optimization updates. If an experiment is resumed at iteration 20 with a limit of 100, training continues to iteration 100; it does not run 100 additional iterations.

@carlosdp

Copy link
Copy Markdown
Contributor Author

Addressed in 88217f6:

  • Moved max-epoch resolution into BaseAgentConfig.resolve_max_epochs() and made BaseAgent call it directly.
  • Removed the standalone helper and the getattr compatibility fallbacks, updating test fixtures to use the current config/CLI contract.
  • Added the requested note describing absolute whole-iteration and resume semantics.

Validated with 64 focused tests and a warnings-as-errors docs build.

@tesslerc
tesslerc merged commit f1ce7c9 into NVlabs:main Jul 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants