fix(train): support eval-only mode (--num-rollout 0)

[EazyReal]

Port of THUDM/slime#2109 for miles.

Summary:

• compute estimated train iterations once
• clamp the scheduler-visible train_iters to at least 1 so Megatron LR decay steps stay valid
• keep the training loop semantics controlled by args.num_rollout; eval-only still performs zero training rollouts

Validation:

• uv run --with pytest --with torch pytest --confcutdir=tests/fast/backends/megatron_utils tests/fast/backends/megatron_utils/test_eval_only_optimizer_scheduler.py -q -> 2 passed

[gemini-code-assist]

Code Review

This pull request ensures that the optimizer parameter scheduler does not crash during eval-only runs (where num_rollout is 0) by clamping args.train_iters to a minimum of 1, and adds unit tests to verify this behavior. The reviewer suggests explicitly checking for num_rollout == 0 and raising a ValueError if estimated_train_iters is 0 during actual training, preventing configuration errors from being silently masked.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Clamping args.train_iters to 1 using max(1, estimated_train_iters) is appropriate for the num_rollout == 0 (eval-only) case. However, if num_rollout > 0 but estimated_train_iters is computed as 0 (due to integer division when total samples per rollout is less than global_batch_size), this silent clamping to 1 can mask a serious configuration error.

It is safer to explicitly handle the num_rollout == 0 case and raise a ValueError for invalid training configurations where estimated_train_iters <= 0 despite num_rollout > 0.

Suggested change

	estimated_train_iters = (
	args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
	)
	# ``num_rollout == 0`` is eval-only: no training runs, but the scheduler is
	# still constructed and Megatron asserts ``lr_decay_steps > 0``.
	args.train_iters = max(1, estimated_train_iters)
	estimated_train_iters = (
	args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
	)
	if args.num_rollout == 0:
	args.train_iters = 1
	else:
	if estimated_train_iters <= 0:
	raise ValueError(
	f"Invalid training configuration: total samples per rollout "
	f"({args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt}) "
	f"must be at least global_batch_size ({args.global_batch_size})."
	)
	args.train_iters = estimated_train_iters