fix: reset dataloader index when dataset changes on resume

[jinglinglingling]

What does this PR do ?

When resuming a single-dataloader run with a different data.train.dataset_name, StatefulDataLoader.load_state_dict() silently skips the first N samples of the new dataset (where N = samples_yielded from the previous run), or hard-crashes with StopIteration during iter() if len(new_dataset) <= N.

This PR tags each train_dataloader.pt save with the active dataset_name and skips the state restore when it differs from the current config, printing a warning so the swap is visible. Legacy saves (raw state dict) are still loaded as before for backward compatibility.

Applied to GRPO (sync + async), SFT, DPO, RM, Distillation. The multi-dataloader save path is intentionally left untouched in this PR.

Issues

N/A

Usage

No new API. Behavior is automatic:

# step 1: train on dataset A and checkpoint
uv run examples/run_grpo_math.py --config ... data.train.dataset_name=OpenMathInstruct-2

# step 2: resume from the checkpoint with a *different* dataset
uv run examples/run_grpo_math.py --config ... data.train.dataset_name=gsm8k checkpointing.checkpoint_must_exist=true

Previously, step 2 would silently skip the first N samples of gsm8k (or crash if N >= len(gsm8k)). Now the new dataset starts from index 0 and a warning is printed.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Wrote standalone empirical + identity-logic tests locally (not upstreamed)
• Ran existing unit tests locally:
  • tests/unit/algorithms/test_sft.py: 5 passed
  • tests/unit/algorithms/test_dpo.py: 4 passed
  • tests/unit/algorithms/test_rm.py: 3 passed
  • tests/unit/algorithms/test_distillation.py: 10 passed
  • tests/unit/algorithms/test_grpo.py: running at PR time, relying on CI
• No new documentation required

Additional Information

• Backward compatible: old checkpoints with the legacy raw-state format still load via an isinstance fallback.
• The save payload is now {"state": <state>, "dataset_name": <name>}.
• Multi-dataloader save path is unchanged in this PR; per-dataset state semantics can be added in a follow-up.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jinglinglingling]

/ok to test 410b0f0

[jinglinglingling]

/ok to test 9627592

[jinglinglingling]

/ok to test a6fe6ba

[jinglinglingling]

/ok to test e68e7bc

[jinglinglingling]

/ok to test 6401a47

[yuki-97]

thanks for helping fix! some comments.

[jinglinglingling]

/ok to test eaa5a40

[yuki-97]

LGTM, thanks @jinglinglingling . change to run L1 to make sure resume works well.

[yuki-97]

/ok to test fc0138c