[feat][checkpoint] Allow selective state restore on resume

[bvolpato]

Summary

Allow checkpoint weights to seed a new training phase while independently controlling global-step, dataloader, optimizer, and learning-rate scheduler restoration.

Exact resume remains unchanged because every option defaults to true.

Changes

• Add trainer.resume_load_global_step, trainer.resume_load_dataloader_state, trainer.resume_load_optimizer_states, and trainer.resume_load_lr_scheduler_states.
• Reset training to step zero when global-step restoration is disabled.
• Skip stale fully-async UID and epoch state for a new training phase.
• Preserve logical epoch progress and epoch-end saves in sync and fully-async trainers when global-step restoration skips the dataloader cursor.
• Bound fully-async producers to the shortened first epoch's remaining groups when the cursor is skipped. Filtered groups return one permit so replacement sampling still works without leaving surplus buffered generation at epoch teardown.
• Pass optimizer and scheduler restore options to policy and critic checkpoint loading.
• Document selective-state resume and separate checkpoint output paths.
• Add focused sync and fully-async regression coverage.

Testing

• pytest tests/train/test_trainer.py tests/train/test_config.py tests/train/test_fully_async_trainer.py -q (179 passed)
• Same focused suites on a clean merge with current upstream/main at a6871df (183 passed)
• pre-commit run --all-files
• uv lock --check
• documentation type check and production build (66 pages)
• git diff --check

Megatron dependency

Merge #2149 before using trainer.resume_load_optimizer_states=false with Megatron. It refreshes optimizer master parameters from the loaded model; without it, the first optimizer step can overwrite checkpoint weights with stale master values. Default full-state resume is unaffected.

[gemini-code-assist]

Code Review

This pull request introduces four new configuration parameters (resume_load_global_step, resume_load_dataloader_state, resume_load_optimizer_states, and resume_load_lr_scheduler_states) to allow users to selectively disable restoring specific states when resuming from a checkpoint. The trainer logic and documentation have been updated accordingly, and new unit tests have been added. Feedback suggests explicitly checking if self.train_dataloader is initialized before attempting to restore its state to avoid potential AttributeError warnings when no dataloader exists.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit b2cde16. Configure here.}