fix: don't overwrite an explicitly set --start-rollout-id

[keepkeen]

What

--start-rollout-id is silently discarded whenever the run does not resume from a Megatron checkpoint.

Root cause

The flag documents itself as a fallback:

The starting rollout step, if not set, will try to load the step from --load when doing continue training, otherwise will be set to 0, meaning training from start.

But slime_validate_args assigns it unconditionally in both branches:

if args.megatron_to_hf_mode == "bridge":
    ...
    else:
        if args.load is None:
            args.load = args.ref_load or args.hf_checkpoint
        args.start_rollout_id = 0          # <-- clobbers an explicit value
else:
    if args.load is None or not os.path.exists(...):
        ...
        args.start_rollout_id = 0          # <-- same here

That branch is taken exactly when there is no resumable Megatron checkpoint — i.e. the situation where an explicit --start-rollout-id is the only way to say where to start. create_actor_model afterwards only fills in the value when it is None:

if args.start_rollout_id is None:
    args.start_rollout_id = start_rollout_ids[0]

so the 0 survives, and train.py runs range(0, num_rollout).

Reproduction

Driving the real slime_validate_args through the harness in tests/test_megatron_argument_validation.py:

megatron_to_hf_mode=raw      user asked for 100 -> got 0
megatron_to_hf_mode=bridge   user asked for 100 -> got 0

So --ref-load /ckpt/base --start-rollout-id 100 --num-rollout 200 silently retrains rollouts 0-199 instead of 100-199.

Fix

Guard both assignments with if args.start_rollout_id is None:, so they only supply the documented default.

No new failure mode downstream: RolloutDataSource.load already returns quietly when rollout/global_dataset_state_dict_{id}.pt is absent, so starting at an id with no saved dataset state just begins from a fresh dataset.

Test

Four cases added to the existing (CI-wired) tests/test_megatron_argument_validation.py, parametrized over raw / bridge: an explicit start_rollout_id=100 survives validation, and None still becomes 0. The two "preserve" cases fail on main with assert 0 == 100.