ci: move Megatron-Bridge to d352aced

nemo_rl/models/generation/megatron/megatron_worker.py

@@ -186,6 +186,16 @@ def _setup_colocated_cuda_graph_managers(self) -> None:
                 "generation, but no CUDA-graph manager could be created for this model."
             )
 
+            # When the model-level manager owns block-scope graphs,
+            # construction deletes the decoder's fallback manager.
+            decoder = getattr(lang_module, "decoder", None)
+            if (
+                hasattr(lang_module, "cudagraph_manager")
+                and decoder is not None
+                and hasattr(decoder, "cudagraph_manager")
+            ):
+                del decoder.cudagraph_manager
+
         # Detach for training; this caches the managers built above.
         toggle_cuda_graphs(lang_module, set_to="none")
 

nemo_rl/models/megatron/draft/__init__.py

@@ -20,6 +20,7 @@
     get_eagle3_aux_hidden_state_layers,
 )
 from nemo_rl.models.megatron.draft.utils import (
+    draft_model_detached,
     export_eagle_weights_to_hf,
     load_hf_weights_to_eagle,
 )
@@ -29,6 +30,7 @@
     "HiddenStateCapture",
     "get_capture_context",
     "EagleModel",
+    "draft_model_detached",
     "load_hf_weights_to_eagle",
     "export_eagle_weights_to_hf",
     "get_eagle3_aux_hidden_state_layers",

nemo_rl/models/megatron/draft/utils.py

@@ -16,9 +16,10 @@
 
 import json
 import re
+from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Mapping
+from typing import Any, Callable, Iterator, Mapping
 
 import torch
 import torch.distributed as dist
@@ -989,6 +990,32 @@ def get_attached_draft_model(model: list[MegatronModule]) -> MegatronModule | No
     return None
 
 
+@contextmanager
+def draft_model_detached(model: list[MegatronModule]) -> Iterator[None]:
+    """Temporarily detach the nested draft model from its owner chunk.
+
+    Megatron-Bridge conversion only covers the base model's HF architecture;
+    draft weights are refit through `export_eagle_weights_to_hf` instead.
+    """
+    owner_chunk: MegatronModule | None = None
+    draft_model: MegatronModule | None = None
+    for model_chunk in reversed(model):
+        unwrapped_chunk = unwrap_model(model_chunk)
+        chunk_draft = getattr(unwrapped_chunk, "draft_model", None)
+        if chunk_draft is not None:
+            owner_chunk = unwrapped_chunk
+            draft_model = chunk_draft
+            break
+    if owner_chunk is None:
+        yield
+        return
+    delattr(owner_chunk, "draft_model")
+    try:
+        yield
+    finally:
+        owner_chunk.draft_model = draft_model
+
+
 def _export_layer_weights_to_hf(
     *,
     source_state: Mapping[str, Tensor],

nemo_rl/models/policy/workers/megatron_policy_worker.py

@@ -2299,15 +2299,19 @@ def _build_refit_conversion_tasks(self) -> list:
         emits a *pair* of tasks per FP8 weight (the FP8 data and a ``*_scale_inv``
         scale tensor).
         """
-        if self._is_fp8_export():
-            return self.megatron_bridge._model_bridge.build_export_fp8_tasks(
-                self.megatron_bridge.hf_pretrained, [self.model]
-            )
-        return [
-            task
-            for task in self.megatron_bridge.get_conversion_tasks([self.model])
-            if task is not None
-        ]
+        # Deferred import to avoid circular import issues.
+        from nemo_rl.models.megatron.draft import draft_model_detached
+
+        with draft_model_detached([self.model]):
+            if self._is_fp8_export():
+                return self.megatron_bridge._model_bridge.build_export_fp8_tasks(
+                    self.megatron_bridge.hf_pretrained, [self.model]
+                )
+            return [
+                task
+                for task in self.megatron_bridge.get_conversion_tasks([self.model])
+                if task is not None
+            ]
 
     def _calculate_refit_param_info(self) -> list[tuple[str, int]]:
         """Calculate parameter information for refit.

tests/unit/models/megatron/test_draft_grad_norm_group.py

@@ -34,3 +34,33 @@ def test_register_draft_grad_norm_group_is_idempotent_and_preserves_existing():
         assert mcore_opt.SEPARATE_GRAD_NORM_GROUPS == after_first  # no-op
     finally:
         mcore_opt.SEPARATE_GRAD_NORM_GROUPS = original
+
+
+@pytest.mark.mcore
+def test_draft_model_detached():
+    import torch.nn as nn
+
+    from nemo_rl.models.megatron.draft.utils import draft_model_detached
+
+    chunk = nn.Module()
+    chunk.decoder = nn.Linear(2, 2)
+    draft = nn.Module()
+    draft.eagle_module = nn.Linear(2, 2)
+    chunk.draft_model = draft
+
+    with draft_model_detached([chunk]):
+        assert not hasattr(chunk, "draft_model")
+        assert all(
+            not name.startswith("draft_model.") for name, _ in chunk.named_parameters()
+        )
+    assert chunk.draft_model is draft
+    assert any(name.startswith("draft_model.") for name, _ in chunk.named_parameters())
+
+    with pytest.raises(RuntimeError, match="boom"):
+        with draft_model_detached([chunk]):
+            raise RuntimeError("boom")
+    assert chunk.draft_model is draft
+
+    # No draft attached: a plain no-op.
+    with draft_model_detached([nn.Module()]):
+        pass