[6/N] Opt-in partial final training step

[Zhichenzzz]

Summary

Stacked on #1936. New opt-in flag --allow-partial-train-step: instead of dropping the trailing rollouts that don't fill a whole global_batch_size step, train them as one smaller final step. This is the endgame the series was built for — genuinely variable per-step global batch sizes.

• Rollout-side schedule + dynamic batch size only (the static path can't satisfy its fixed-size mbs alignment on an arbitrary remainder); a partial step smaller than dp_size samples is skipped with a log line.
• No other machinery changes: loss normalization and the LR scheduler increment already use the true per-step rollout count (from [1/N]), and the metric reduction is already count-weighted.
• Saves up to global_batch_size - 1 rollouts per rollout batch under dynamic sampling / filtering — samples are the expensive part of RL.

Validation

• CPU: 101 tests green, incl. partial-step opt-in coverage (uneven [2, 2, 1] schedule covering every sample) and the below-dp_size skip guard; default behavior unchanged when the flag is off.
• GPU e2e (Qwen2.5-0.5B gsm8k, dp8, dynamic batch, mooncake, --ci-test, --global-batch-size 24 --disable-rollout-trim-samples):

Rollout-side DP schedule: num_samples=32, num_rollouts=32, global_batch_sizes=[24, 8], num_microbatches=[1, 1]

All 3 rollout batches trained with uneven [24, 8] steps — the first genuinely variable global batch size training in miles — job succeeded, CI logprob asserts pass, rewards healthy (0.47-0.59). Without the flag those 8 rollouts per batch (25%) would have been discarded.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.