Block compile: keep the GatedDeltaNet self-check out of a non-reentrant checkpoint (#14195)

Addresses the GatedDeltaNet failure in #14063.

## Problem

On 2026.10.6, LoRA training of Qwen3.5 (a GatedDeltaNet hybrid) with
`use_gradient_checkpointing="unsloth"` under DDP or FSDP2 dies in the
backward of step 1. Reproduced here on Qwen3.5-0.8B with 2 GPUs; the
issue reports the same on the 27B model with 8 GPUs:

```
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: A different number of tensors was saved during the original forward and recomputation.
Number of tensors saved during forward: 55
Number of tensors saved during recomputation: 45.
```

Single-GPU runs are not affected. `UNSLOTH_BLOCK_COMPILE=0` avoids it.

## Root cause

Distributed runs use torch's non-reentrant checkpoint (`use_reentrant =
not is_distributed()` in `vision.py`), whose saved-tensor hooks are
active for the whole layer forward.

On the first GatedDeltaNet layer, the block-compile wrapper calls
`_prepare_gdn`, which runs `_gdn_op._fla_ops`. Its `_self_check` does a
forward and `torch.autograd.grad` against fla's own Function. Because
this happens inside the checkpointed layer forward:

1. The self-check's saved tensors are packed into the layer's checkpoint
frame.
2. Its `autograd.grad` unpacks them, which starts a recompute of the
whole layer during forward0.
3. That nested recompute enters the block-compile wrapper before the
layer is marked checked, so the trial loop catches torch's
`_StopRecomputationError` and demotes the layer kind `_FULL -> _SPLIT ->
_OFF`. The self-check then fails as well, so the fla opaque op is
disabled.
4. The original forward finishes compiled and saves the self-check's
tensors plus the compiled set (55). The real recompute in backward runs
the stock forward for that layer (45), and
`check_recomputed_tensors_match` raises.

The fla self-check runs for every user because unsloth_zoo vendors fla.

## Fix

Run the self-check in `_fla_ops` under its own `saved_tensors_hooks`, so
an enclosing checkpoint never sees its tensors. The pack hook detaches,
as torch's own recompute hook does, so a check that fails before
`autograd.grad` does not leave a reference cycle holding its tensors.

#14065 alone does not fix the reproduction. Re-raising
`_StopRecomputationError` removes the demotion, but the self-check's
tensors are still in the frame, so the run still fails with 55 vs 46.
With this change the nested recompute never happens.

## Verification

Qwen3.5-0.8B, LoRA r=16, `use_gradient_checkpointing="unsloth"`, 3
steps, 2 GPUs, torch 2.12.1+cu130:

| Launch | main | This PR |
| --- | --- | --- |
| `torchrun`, DDP wrapper | `CheckpointError` 55 vs 45, 18 layers OFF |
trains, 24 layers FULL |
| `accelerate` FSDP2, SFTTrainer | `CheckpointError` 55 vs 45, 18 layers
OFF | trains, 24 layers FULL |

Single GPU (reentrant) is unchanged: same step-0 loss, all layers FULL.

Tests:

- New `test_gdn_self_check_stays_out_of_a_non_reentrant_checkpoint` runs
`_fla_ops` inside `checkpoint(..., use_reentrant=False)` with a stand-in
self-check that calls `autograd.grad`. It fails on main with
`CheckpointError` (10 vs 5) and passes here. It runs on CPU without fla.
- `tests/test_block_compile.py`: 22 passed, 1 failed.
`test_qwen3_5_gdn_block_has_no_break` fails the same way on main.

## Not covered

The plain-attention case in the issue (`MuseGlimmerTextDecoderLayer`, 65
vs 64) and the reported slowdown.

tests/test_block_compile.py

@@ -362,6 +362,44 @@ def test_gdn_self_check_leaves_default_rng_alone():
     assert torch.equal(torch.cuda.get_rng_state(), state)
 
 
+def test_gdn_self_check_stays_out_of_a_non_reentrant_checkpoint(monkeypatch):
+    # #14063: the check's autograd.grad must not recompute the checkpointed layer mid-forward.
+    from torch.utils.checkpoint import checkpoint
+
+    G = _load("_gdn_op")
+    chunk = types.ModuleType("fla.ops.gated_delta_rule.chunk")
+    for name in (
+        "chunk_gated_delta_rule_fwd",
+        "chunk_gated_delta_rule_bwd",
+        "l2norm_fwd",
+        "l2norm_bwd",
+    ):
+        setattr(chunk, name, lambda *a, **k: None)
+    for name in ("fla", "fla.ops", "fla.ops.gated_delta_rule"):
+        monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
+    monkeypatch.setitem(sys.modules, "fla.ops.gated_delta_rule.chunk", chunk)
+
+    def self_check(q, k, v, g, beta):
+        ins = [t.requires_grad_() for t in (q, k, v)]
+        out = (ins[0] * ins[1]).sigmoid() * ins[2]
+        return all(gr is not None for gr in torch.autograd.grad(out.sum(), ins))
+
+    monkeypatch.setattr(G, "_self_check", self_check)
+    calls = []
+    w = torch.randn(8, 8, requires_grad = True)
+
+    def layer(x):
+        calls.append(1)
+        h = (x @ w).sigmoid()
+        G._fla_ops(lambda: (h, h, h, h, h))
+        return (h @ w).sigmoid()
+
+    checkpoint(layer, torch.randn(4, 8), use_reentrant = False).sum().backward()
+    assert G._FLA
+    assert len(calls) == 2  # forward, then one recompute in backward
+    assert w.grad is not None
+
+
 def test_sm_count_is_read_outside_the_trace(monkeypatch):
     # Stands in for Triton's hip_utils.get_device_properties, which Dynamo cannot trace.
     M = _load()

unsloth/models/_gdn_op.py

@@ -7,6 +7,7 @@ as one graph."""
 
 import torch
 import torch.nn.functional as F
+from torch.autograd.graph import saved_tensors_hooks
 
 
 def _causal_depthwise_conv(x, weight):
@@ -40,11 +41,12 @@ def _fla_ops(make_inputs):
                 C.l2norm_fwd,
                 C.l2norm_bwd,
             )
-            # The first call can come from a no_grad checkpoint forward.
+            # The first call can come from a no_grad checkpoint forward, or a non-reentrant one
+            # whose hooks pack the check's tensors; detached so a failed check leaves no cycle.
             with torch.inference_mode(False):
                 with torch.no_grad():
                     inputs = [t.detach().clone() for t in make_inputs()]
-                with torch.enable_grad():
+                with torch.enable_grad(), saved_tensors_hooks(lambda t: t.detach(), lambda t: t):
                     _FLA = fns
                     _FLA = fns if _self_check(*inputs) else False
         except Exception: