From 798eecdb795b1ed4e2684f4c72c0e202f5eee440 Mon Sep 17 00:00:00 2001 From: NolenLiang Date: Thu, 23 Jul 2026 01:49:31 -0700 Subject: [PATCH] Fix Megatron async checkpoint CUDA memory retention Finalize NVRx saves before colocated offload and release persistent writer tensor caches and CUDA IPC mappings. This prevents stale checkpoint storage from blocking subsequent vLLM wakeups. Add lifecycle coverage for both offload paths and conditional cache cleanup. Signed-off-by: NolenLiang --- .../policy/workers/megatron_policy_worker.py | 63 +++++++- .../models/policy/test_megatron_worker.py | 148 ++++++++++++++++++ 2 files changed, 208 insertions(+), 3 deletions(-) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index eecd310a78..bcebfd11e5 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -37,6 +37,7 @@ ) from megatron.bridge.utils.common_utils import get_rank_safe from megatron.core import parallel_state +from megatron.core.dist_checkpointing.strategies.torch import get_async_strategy from megatron.core.distributed import DistributedDataParallel from megatron.core.distributed.fsdp.mcore_fsdp_adapter import ( FullyShardedDataParallel as custom_FSDP, @@ -161,6 +162,7 @@ class MegatronPolicyWorkerImpl( # ``self._train_step_state = None`` after finish/abort type-checks. _train_step_state: Optional[dict[str, Any]] = None _remote_sparse_refit: Any = None + _async_checkpoint_cuda_cache_active: bool = False def __repr__(self): """Customizes the actor's prefix in the Ray logs. @@ -2095,6 +2097,12 @@ def _clear_fp8_caches(self): @wrap_with_nvtx_name("megatron_policy_worker/offload_before_refit") def offload_before_refit(self): """Offload the optimizer and buffers to the CPU.""" + # An in-flight async checkpoint keeps references to the CUDA tensors in + # its sharded state dict until the write is finalized. Offloading swaps + # those tensors for CPU storage, so the checkpoint references would keep + # the old CUDA storage alive and defeat the offload. + self.finalize_async_save() + no_grad = torch.no_grad() no_grad.__enter__() allocated = torch.cuda.memory_allocated() / (1024**3) # Convert to GB @@ -2186,6 +2194,11 @@ def offload_before_refit(self): @wrap_with_nvtx_name("megatron_policy_worker/offload_after_refit") def offload_after_refit(self): """Offload as much as possible on the CPU.""" + # Finalize before replacing model-buffer storage. With cached NVRx async + # saves, the persistent writer otherwise retains CUDA IPC handles to the + # old storage after the model is moved to CPU. + self.finalize_async_save() + no_grad = torch.no_grad() no_grad.__enter__() self.model = self.move_model(self.model, "cpu") @@ -2296,6 +2309,10 @@ def save_checkpoint( original_save_path = self.mcore_state.cfg.checkpoint.save is_async = self.mcore_state.cfg.checkpoint.async_save + if is_async and self._requires_nvrx_cuda_cache_release(): + # Set this before saving so an exception path can still tear down a + # writer that may already have received CUDA IPC handles. + self._async_checkpoint_cuda_cache_active = True try: # Block until any previous async save is fully written to disk. @@ -2353,17 +2370,57 @@ def save_checkpoint( finally: self.mcore_state.cfg.checkpoint.save = original_save_path + def _requires_nvrx_cuda_cache_release(self) -> bool: + """Whether checkpoint finalization must also drop cached CUDA IPC handles.""" + ckpt_cfg = self.mcore_state.cfg.checkpoint + generation_cfg = self.cfg.get("generation") or {} + colocated_cfg = generation_cfg.get("colocated") or {} + return bool( + ckpt_cfg.async_save + and getattr(ckpt_cfg, "async_strategy", "nvrx") == "nvrx" + and getattr(ckpt_cfg, "use_persistent_ckpt_worker", False) + and getattr(ckpt_cfg, "ckpt_assume_constant_structure", False) + and not getattr(ckpt_cfg, "async_ckpt_use_cpu_shm", False) + and colocated_cfg.get("enabled", False) + ) + def finalize_async_save(self): - """Block until the in-flight async write completes and run finalize_fns. + """Finalize an async write and release unsafe colocated CUDA IPC caches. - Safe to call when async_save is disabled (no-op). - Does NOT terminate the persistent worker — it stays alive for the next save. + NVRx constant-structure saves cache CUDA tensor handles in the persistent + writer. That is safe while model/optimizer storage stays fixed, but a + colocated policy replaces that storage during CPU offload. In that case, + close the completed writer and invalidate its training-side cache; NVRx + starts a fresh persistent writer lazily for the next checkpoint. """ + release_cuda_cache = bool( + self._async_checkpoint_cuda_cache_active + and self._requires_nvrx_cuda_cache_release() + ) maybe_finalize_async_save( self.mcore_state, ckpt_cfg=self.mcore_state.cfg.checkpoint, blocking=True, + terminate=release_cuda_cache, ) + if release_cuda_cache: + _, async_modules = get_async_strategy( + self.mcore_state.cfg.checkpoint.async_strategy + ) + writer_cls = async_modules["FileSystemWriterAsync"] + cleanup_tensor_caches = getattr(writer_cls, "cleanup_tensor_caches", None) + if cleanup_tensor_caches is not None: + cleanup_tensor_caches() + else: + # Compatibility with older NVRx versions that predate the + # public cleanup helper. + cached_identifiers = getattr(writer_cls, "_cached_identifiers", None) + if cached_identifiers is not None: + cached_identifiers.clear() + gc.collect() + torch.cuda.ipc_collect() + torch.cuda.empty_cache() + self._async_checkpoint_cuda_cache_active = False def load_checkpoint(self, weights_path: str, optimizer_path: Optional[str] = None): """Load a training checkpoint. diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 9727b50d06..759ff20b1c 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -48,10 +48,14 @@ class _FakeTrainableModel: def __init__(self): self.train_called = False + self.eval_called = False def train(self): self.train_called = True + def eval(self): + self.eval_called = True + class _ModelWithNonSerializableExtraState(torch.nn.Module): def __init__(self): @@ -63,6 +67,150 @@ def get_extra_state(self): raise AssertionError("moving a module must not serialize its extra state") +def test_megatron_offload_before_refit_finalizes_async_save_first(monkeypatch): + """Async checkpoint tensor references must be released before GPU offload.""" + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + events = [] + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = object() + worker.optimizer = None + worker.optimizer_cpu_offload = False + worker.fp8_cfg = None + worker.cfg = {"megatron_cfg": {"clear_memory_caches_before_refit": False}} + worker.finalize_async_save = lambda: events.append("finalize_async_save") + worker.move_model = lambda model, device, move_params, move_grads: ( + events.append("move_model") or model + ) + + class _AllocatorWakeup: + def cuda(self): + events.append("wake_allocator") + + monkeypatch.setattr( + torch.cuda, + "memory_allocated", + lambda *args, **kwargs: events.append("memory_allocated") or 0, + ) + monkeypatch.setattr( + torch.cuda, + "memory_reserved", + lambda *args, **kwargs: events.append("memory_reserved") or 0, + ) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: events.append("empty_cache")) + monkeypatch.setattr(torch, "randn", lambda *args, **kwargs: _AllocatorWakeup()) + + MegatronPolicyWorkerImpl.offload_before_refit(worker) + + assert events[0] == "finalize_async_save" + assert events.index("finalize_async_save") < events.index("move_model") + + +def test_megatron_offload_after_refit_finalizes_before_model_move(monkeypatch): + """Checkpoint CUDA IPC handles must be dropped before model storage is replaced.""" + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + MegatronPolicyWorkerImpl, + ) + + events = [] + worker = object.__new__(MegatronPolicyWorkerImpl) + worker.model = _FakeTrainableModel() + worker.finalize_async_save = lambda: events.append("finalize_async_save") + worker.move_model = lambda model, device: events.append("move_model") or model + worker.offload_before_refit = lambda: events.append("offload_before_refit") + + class _AllocatorWakeup: + def cuda(self): + events.append("wake_allocator") + + monkeypatch.setattr( + torch.cuda, + "memory_allocated", + lambda *args, **kwargs: events.append("memory_allocated") or 0, + ) + monkeypatch.setattr( + torch.cuda, + "memory_reserved", + lambda *args, **kwargs: events.append("memory_reserved") or 0, + ) + monkeypatch.setattr(torch, "randn", lambda *args, **kwargs: _AllocatorWakeup()) + + MegatronPolicyWorkerImpl.offload_after_refit(worker) + + assert events[0] == "finalize_async_save" + assert events.index("finalize_async_save") < events.index("move_model") + + +@pytest.mark.parametrize("cache_active", [True, False]) +def test_megatron_finalize_async_save_releases_colocated_nvrx_cache( + monkeypatch, cache_active +): + """Only an active unsafe NVRx cache should terminate the persistent writer.""" + import nemo_rl.models.policy.workers.megatron_policy_worker as worker_module + + worker = object.__new__(worker_module.MegatronPolicyWorkerImpl) + worker.cfg = {"generation": {"colocated": {"enabled": True}}} + worker.mcore_state = SimpleNamespace( + cfg=SimpleNamespace( + checkpoint=SimpleNamespace( + async_save=True, + async_strategy="nvrx", + use_persistent_ckpt_worker=True, + ckpt_assume_constant_structure=True, + async_ckpt_use_cpu_shm=False, + ) + ) + ) + worker._async_checkpoint_cuda_cache_active = cache_active + events = [] + + monkeypatch.setattr( + worker_module, + "maybe_finalize_async_save", + lambda *args, **kwargs: events.append(("finalize", kwargs["terminate"])), + ) + + class _Writer: + @classmethod + def cleanup_tensor_caches(cls): + events.append(("cleanup_tensor_caches", None)) + + monkeypatch.setattr( + worker_module, + "get_async_strategy", + lambda strategy: (strategy, {"FileSystemWriterAsync": _Writer}), + ) + monkeypatch.setattr( + worker_module.gc, "collect", lambda: events.append(("gc_collect", None)) + ) + monkeypatch.setattr( + torch.cuda, + "ipc_collect", + lambda: events.append(("ipc_collect", None)), + ) + monkeypatch.setattr( + torch.cuda, + "empty_cache", + lambda: events.append(("empty_cache", None)), + ) + + worker_module.MegatronPolicyWorkerImpl.finalize_async_save(worker) + + assert events[0] == ("finalize", cache_active) + if cache_active: + assert events[1:] == [ + ("cleanup_tensor_caches", None), + ("gc_collect", None), + ("ipc_collect", None), + ("empty_cache", None), + ] + assert worker._async_checkpoint_cuda_cache_active is False + else: + assert events == [("finalize", False)] + + def test_megatron_move_model_does_not_serialize_extra_state(): from nemo_rl.models.policy.workers.megatron_policy_worker import ( MegatronPolicyWorkerImpl,