Skip to content

[Fix] Checkpoint role-worker LR schedules - #2242

Open
matrix72c wants to merge 1 commit into
InternLM:stack/rl-ppofrom
matrix72c:fix/ppo-role-scheduler-checkpoint
Open

matrix72c wants to merge 1 commit into
InternLM:stack/rl-ppofrom
matrix72c:fix/ppo-role-scheduler-checkpoint

Conversation

@matrix72c

Copy link
Copy Markdown
Contributor

Stacked on #2186 (stack/rl-ppo).

Root cause

Role workers build a LambdaLR from the role's LRConfig, but RoleWorker.save() / resume() only go through DCP. The scheduler step is never written, so a resumed actor or critic restarts its LR curve from step 0, warmup included. When load_optimizer_args=True, the restored optimizer LR is also overwritten by the fresh schedule on the first step.

Fix

  • save_lr_scheduler() writes lr_scheduler.json into each role's DCP directory on rank 0, also when no_save_optimizer=True.
  • resume_lr_scheduler() loads it when load_checkpoint_cfg.load_scheduler is set (the default), matching the SFT trainer. A missing file raises with a hint to set load_scheduler=False, which restarts the schedule.
  • In both cases the optimizer LR is set from the scheduler, so it no longer depends on whether optimizer arguments were restored.

Test plan

tests/rl/test_scheduler_checkpoint.py checks constant/linear/cosine resumes against an uninterrupted run, actor and critic wiring through TrainingWorker.save/resume with and without optimizer arguments, saving with no_save_optimizer=True, and the missing-file / restart paths. Together with test_ppo_role_wiring.py, test_routed_experts_dtype.py and test_training_worker_rank.py: 65 passed on CPU.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant