diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 3a9ff57dbfe..7e68eed4f1a 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -2982,6 +2982,19 @@ def save_checkpoint( ckpt_cfg=self.mcore_state.cfg.checkpoint, blocking=True, ) + + # Onload model before saving it. + self.model = self.move_model( + self.model, "cuda", move_params=True, move_grads=False + ) + if ( + optimizer_path is not None + and self.optimizer is not None + and not self.optimizer_cpu_offload + ): + self.move_optimizer("cuda") + torch.cuda.synchronize() + self.mcore_state.cfg.checkpoint.save = weights_path optimizer_to_save = None diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 8cd4b715d1d..df31479a858 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -241,6 +241,55 @@ def cuda(self): assert events.index("finalize_async_save") < events.index("move_model") +def test_megatron_save_checkpoint_onloads_model_before_save(monkeypatch): + """Params offloaded by colocated generation must be onloaded before the save walks them.""" + import nemo_rl.models.policy.workers.megatron_policy_worker as worker_module + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + events = [] + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = _FakeTrainableModel() + worker.model.training = False + worker.optimizer = object() + worker.scheduler = None + worker.optimizer_cpu_offload = False + worker.should_disable_forward_pre_hook = False + worker.checkpointing_context = None + worker.mcore_state = SimpleNamespace( + cfg=SimpleNamespace( + checkpoint=SimpleNamespace(save="original_path", async_save=False) + ), + train_state=SimpleNamespace(floating_point_operations_so_far=0), + ) + worker.move_model = lambda model, device, move_params, move_grads: ( + events.append(f"move_model_{device}") or model + ) + worker.move_optimizer = lambda device: events.append(f"move_optimizer_{device}") + + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: events.append("synchronize")) + monkeypatch.setattr( + worker_module, + "maybe_finalize_async_save", + lambda *args, **kwargs: events.append("finalize_async_save"), + ) + monkeypatch.setattr( + worker_module, "save_checkpoint", lambda **kwargs: events.append("mcore_save") + ) + + MegatronPolicyWorkerImpl.save_checkpoint( + worker, weights_path="ckpt/weights", optimizer_path="ckpt/optim" + ) + + assert "mcore_save" in events + assert events.index("move_model_cuda") < events.index("mcore_save") + assert events.index("move_optimizer_cuda") < events.index("mcore_save") + assert events.index("synchronize") < events.index("mcore_save") + assert worker.mcore_state.cfg.checkpoint.save == "original_path" + + @pytest.mark.parametrize("cache_active", [True, False]) def test_megatron_finalize_async_save_releases_colocated_nvrx_cache( monkeypatch, cache_active