Skip to content

Fix Megatron checkpoint save and resume state - #70

Draft
nightlessbaron wants to merge 1 commit into
prodfrom
agent/checkpoint-resume-correctness
Draft

Fix Megatron checkpoint save and resume state#70
nightlessbaron wants to merge 1 commit into
prodfrom
agent/checkpoint-resume-correctness

Conversation

@nightlessbaron

Copy link
Copy Markdown

What changed

  • Pass Megatron's preprocess_common_state_dict hook when saving non-LoRA checkpoints.
  • Avoid advancing the optimizer scheduler a second time after a normal checkpoint resume.
  • Continue reconstructing scheduler position for finetuning and --no-load-optim, where scheduler state is intentionally not restored.
  • Add four focused regression cases.

Why

A normal resume already restores scheduler state in load_checkpoint; stepping it again double-counts completed iterations. Checkpoint saves also need Megatron's common-state preprocessing to produce the expected distributed checkpoint state.

Validation

  • Focused checkpoint tests: 4 passed.
  • Black: passed.
  • Ruff: passed.
  • py_compile: passed.
  • git diff --check: passed.

Related Megatron optimizer-state fix: LLM360/Megatron-LM#6.

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