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

[pstefa1707]

Problem

Adapter-only (LoRA) training of hybrid Mamba models (Nemotron 3 Super) returns a real loss but grad_norm=0.0 — every LoRA tensor gets zero gradients.

Root cause

• Since megatron-core 0.19, HybridStack honours recompute_granularity='full' through the same reentrant tensor_parallel.checkpoint as TransformerBlock.
• A reentrant checkpoint only records a backward node when one of its tensor inputs requires grad; parameters inside the chunk (the LoRA adapters) are invisible to autograd at apply time.
• With a frozen base model at PP=1, the embedding output has requires_grad=False, so every checkpointed chunk cascades to grad_fn=None and the whole decoder drops out of the autograd graph.
• maybe_enable_recompute_inputs_grad already fixes exactly this — but it only patched TransformerBlock, so pure transformers (Qwen, Kimi) work while every hybrid model silently trains nothing.

Fix

Patch HybridStack the same way as TransformerBlock (guarded import for older megatron-core without the hybrid block). Unit test added mirroring the existing TransformerBlock test.

Verification

On Nemotron-3-Super-120B-A12B LoRA (TP=8, EP=8, PP=1, B200, recompute_granularity='full'):

• Before: grad_norm=0.0, decoder output grad_fn=None, 1 CheckpointFunctionBackward node in the loss graph (the MTP branch only).
• After: 89 CheckpointFunctionBackward nodes, grad_norm 4.65 → 72.9 → 54.6 and loss 4.43 → 2.54 → 0.76 over repeated steps on a fixed sample; both LoRA A and B matrices receive gradients.
• tests/unit_tests/peft/test_recompute.py passes (2 passed) in the trainer image venv; ruff check/format clean.

🤖 Generated with Claude Code