[baseten] fix(peft): extend recompute input-grad hook to HybridStack

src/megatron/bridge/peft/recompute.py

@@ -53,14 +53,20 @@ def maybe_enable_recompute_inputs_grad(model, peft_recompute_patched: Set[int] |
       This means CheckpointFunction.backward() is never called, and LoRA gradients
       inside the checkpoint are never computed.
 
-    Solution: Hook TransformerBlock.forward to ensure hidden_states.requires_grad=True
+    Solution: Hook the decoder block's forward to ensure hidden_states.requires_grad=True
     before it enters checkpointed computation. This doesn't unfreeze any parameters;
     it just ensures the autograd machinery calls checkpoint's backward.
 
+    ``HybridStack`` (Mamba hybrids such as NemotronH / Nemotron 3) honours
+    ``recompute_granularity='full'`` via the same reentrant checkpoint since
+    megatron-core 0.19 and needs the same fix — without it, adapter-only hybrid
+    training at PP=1 silently produces zero gradients.
+
     Borrowed (with modifications) from
     https://github.com/HollowMan6/verl/blob/4285f0601028aee7ddcb9ec5a15198ebfc69bba3/verl/utils/megatron_peft_utils.py
     """
 
+    from megatron.core.models.hybrid.hybrid_block import HybridStack
     from megatron.core.transformer.transformer_block import TransformerBlock
 
     patched_registry = peft_recompute_patched or PEFT_RECOMPUTE_PATCHED
@@ -83,8 +89,8 @@ def maybe_enable_recompute_inputs_grad(model, peft_recompute_patched: Set[int] |
             if not (trainable_adapter and not trainable_base):
                 continue  # Not adapter-only training, no fix needed
 
-            def _patch_transformer_block(module: torch.nn.Module) -> bool:
-                if isinstance(module, TransformerBlock):
+            def _patch_block(module: torch.nn.Module) -> bool:
+                if isinstance(module, (TransformerBlock, HybridStack)):
                     original_forward = module.forward
 
                     @wraps(original_forward)
@@ -104,18 +110,18 @@ def patched_forward(hidden_states, *args, _original_forward=original_forward, **
 
             patched = False
             for module in unwrapped_model.modules():
-                if _patch_transformer_block(module):
+                if _patch_block(module):
                     patched = True
             if patched:
                 patched_registry.add(id(unwrapped_model))
                 print_rank_0(
-                    "[PEFT+Recompute] Patched TransformerBlock.forward to enable grad on "
+                    "[PEFT+Recompute] Patched decoder block forward to enable grad on "
                     "hidden_states input. This ensures checkpoint backward is called when "
                     "only adapters are trainable (PP=1 with frozen base model).",
                 )
     except Exception as exc:  # pragma: no cover - best effort logging
         # Log but don't fail - user will see grad_norm=0 and can debug
-        print_rank_0(f"[PEFT+Recompute] Warning: Failed to patch TransformerBlock: {exc}")
+        print_rank_0(f"[PEFT+Recompute] Warning: Failed to patch decoder block: {exc}")
 
     return patched_registry
 

tests/unit_tests/peft/test_recompute.py

@@ -39,10 +39,10 @@ def forward(self, hidden_states, *args, **kwargs):
 
 
 class DummyModel(torch.nn.Module):
-    def __init__(self) -> None:
+    def __init__(self, block_cls=None) -> None:
         super().__init__()
         self.config = SimpleNamespace(recompute_method="uniform")
-        self.block = DummyTransformerBlock()
+        self.block = (block_cls or DummyTransformerBlock)()
 
         # Frozen base parameter (not trainable)
         self.base = torch.nn.Linear(1, 1, bias=False)
@@ -59,6 +59,10 @@ def modules(self):
             yield module
 
 
+class DummyHybridStack(DummyTransformerBlock):
+    """Distinct dummy type standing in for megatron's HybridStack."""
+
+
 def _patch_transformer_block(monkeypatch):
     import megatron.core.transformer.transformer_block as transformer_block
 
@@ -70,6 +74,17 @@ def _patch_transformer_block(monkeypatch):
     )
 
 
+def _patch_hybrid_stack(monkeypatch):
+    import megatron.core.models.hybrid.hybrid_block as hybrid_block
+
+    monkeypatch.setattr(
+        hybrid_block,
+        "HybridStack",
+        DummyHybridStack,
+        raising=False,
+    )
+
+
 def test_maybe_enable_recompute_inputs_grad_patches_block(monkeypatch):
     _patch_transformer_block(monkeypatch)
     recompute_mod.PEFT_RECOMPUTE_PATCHED.clear()
@@ -90,3 +105,22 @@ def test_maybe_enable_recompute_inputs_grad_patches_block(monkeypatch):
     # Second invocation should be a no-op (no duplicate patch)
     maybe_enable_recompute_inputs_grad(model, patched_registry)
     assert model.block.forward is patched_forward
+
+
+def test_maybe_enable_recompute_inputs_grad_patches_hybrid_stack(monkeypatch):
+    # HybridStack honours recompute_granularity='full' via the same reentrant
+    # checkpoint as TransformerBlock, so adapter-only training needs the same
+    # input-grad fix — without it LoRA gradients are silently zero at PP=1.
+    _patch_hybrid_stack(monkeypatch)
+    recompute_mod.PEFT_RECOMPUTE_PATCHED.clear()
+
+    model = DummyModel(block_cls=DummyHybridStack)
+    patched_registry = maybe_enable_recompute_inputs_grad(model, set())
+
+    assert id(model) in patched_registry
+
+    input_tensor = torch.zeros(2, 2)
+    assert input_tensor.requires_grad is False
+
+    model.block(input_tensor)
+    assert model.block.last_input_requires_grad is True