From f3b1193b073501bf0a6843fc648d57b0042f20b1 Mon Sep 17 00:00:00 2001 From: Varad Pimpalkhute Date: Wed, 15 Jul 2026 21:14:44 +0000 Subject: [PATCH] Fix Megatron checkpoint save and resume state --- miles/backends/megatron_utils/model.py | 10 +- .../test_checkpoint_resume_correctness.py | 123 ++++++++++++++++++ 2 files changed, 130 insertions(+), 3 deletions(-) create mode 100644 tests/fast/backends/megatron_utils/test_checkpoint_resume_correctness.py diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 3c50770aa0d..85097426ded 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -19,7 +19,7 @@ from megatron.core.pipeline_parallel import get_forward_backward_func from megatron.core.utils import get_model_config from megatron.training.global_vars import get_args -from megatron.training.training import get_model +from megatron.training.training import get_model, preprocess_common_state_dict from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase from miles.utils.memory_utils import clear_memory @@ -728,7 +728,7 @@ def save( num_floating_point_operations_so_far=0, checkpointing_context=None, train_data_iterator=None, - preprocess_common_state_dict_fn=None, + preprocess_common_state_dict_fn=preprocess_common_state_dict, ) clear_memory() @@ -834,7 +834,11 @@ def initialize_model_and_optimizer( check_model_hashes(args, model, iteration) - opt_param_scheduler.step(increment=iteration * args.global_batch_size) + # A full resume restores the scheduler in load_checkpoint(). Advancing it + # again here would double-count the completed iterations. Only reconstruct + # its position when optimizer/scheduler state was intentionally not loaded. + if args.finetune or args.no_load_optim: + opt_param_scheduler.step(increment=iteration * args.global_batch_size) if is_megatron_main_rank(): init_train_step_counter(args.load, iteration) diff --git a/tests/fast/backends/megatron_utils/test_checkpoint_resume_correctness.py b/tests/fast/backends/megatron_utils/test_checkpoint_resume_correctness.py new file mode 100644 index 00000000000..e0bbf708bd5 --- /dev/null +++ b/tests/fast/backends/megatron_utils/test_checkpoint_resume_correctness.py @@ -0,0 +1,123 @@ +import importlib +import sys +import types +from argparse import Namespace +from enum import Enum +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + + +def _import_model_module(): + try: + return importlib.import_module("miles.backends.megatron_utils.model") + except (ImportError, ModuleNotFoundError) as error: + if "sglang" not in str(error): + raise + + # The checkpoint paths under test do not use the optional dumper. Keep the + # unit test runnable in lightweight environments that omit SGLang. + dumper_stub = types.ModuleType("miles.utils.dumper_utils") + + class DumperPhase(Enum): + FWD_ONLY = "fwd_only" + FWD_BWD = "fwd_bwd" + + dumper_stub.DumperMegatronUtil = object + dumper_stub.DumperPhase = DumperPhase + previous_dumper_module = sys.modules.get(dumper_stub.__name__) + sys.modules[dumper_stub.__name__] = dumper_stub + try: + return importlib.import_module("miles.backends.megatron_utils.model") + finally: + if previous_dumper_module is None: + sys.modules.pop(dumper_stub.__name__, None) + else: + sys.modules[dumper_stub.__name__] = previous_dumper_module + + +model_module = _import_model_module() + + +def test_save_uses_megatron_common_state_preprocessor(monkeypatch: pytest.MonkeyPatch) -> None: + save_checkpoint = Mock() + monkeypatch.setattr( + model_module, + "get_args", + lambda: Namespace(ci_test=False, ci_save_model_hash=False, save="unused"), + ) + monkeypatch.setattr(model_module, "is_lora_model", lambda model: False) + monkeypatch.setattr(model_module, "should_disable_forward_pre_hook", lambda args: False) + monkeypatch.setattr(model_module, "is_megatron_main_rank", lambda: False) + monkeypatch.setattr(model_module, "save_checkpoint", save_checkpoint) + monkeypatch.setattr(model_module, "clear_memory", Mock()) + + model = [object()] + optimizer = object() + scheduler = object() + model_module.save(12, model, optimizer, scheduler) + + save_checkpoint.assert_called_once_with( + 12, + model, + optimizer, + scheduler, + num_floating_point_operations_so_far=0, + checkpointing_context=None, + train_data_iterator=None, + preprocess_common_state_dict_fn=model_module.preprocess_common_state_dict, + ) + + +@pytest.mark.parametrize( + ("finetune", "no_load_optim", "expected_scheduler_steps"), + [ + (False, False, 0), + (True, False, 1), + (False, True, 1), + ], + ids=["normal-resume", "finetune", "optimizer-state-not-loaded"], +) +def test_scheduler_is_only_reconstructed_when_checkpoint_state_was_not_loaded( + monkeypatch: pytest.MonkeyPatch, + finetune: bool, + no_load_optim: bool, + expected_scheduler_steps: int, +) -> None: + args = Namespace( + finetune=finetune, + no_load_optim=no_load_optim, + global_batch_size=32, + ) + model = [SimpleNamespace()] + optimizer = object() + scheduler = Mock() + load_checkpoint = Mock(return_value=(7, 0)) + + monkeypatch.setattr( + model_module, + "setup_model_and_optimizer", + lambda args, role: (model, optimizer, scheduler), + ) + monkeypatch.setattr(model_module, "load_checkpoint", load_checkpoint) + monkeypatch.setattr(model_module, "check_peak_gpu_memory_after_load", Mock()) + monkeypatch.setattr(model_module, "check_model_hashes", Mock()) + monkeypatch.setattr(model_module, "is_megatron_main_rank", lambda: False) + monkeypatch.setattr(model_module, "clear_memory", Mock()) + + result = model_module.initialize_model_and_optimizer(args, role="actor") + + assert result == (model, optimizer, scheduler, 7) + assert model[0].role == "actor" + load_checkpoint.assert_called_once_with( + model, + optimizer, + scheduler, + checkpointing_context={}, + skip_load_to_model_and_opt=False, + ) + if expected_scheduler_steps: + scheduler.step.assert_called_once_with(increment=7 * args.global_batch_size) + else: + scheduler.step.assert_not_called()