perf(moe): pin the GC replay's NF4 expert stacks by default above a size floor - #2095

Summary

Under gradient checkpointing, a frozen NF4 expert stack is dequantized three times per layer per step: in the forward, in the checkpoint replay, and again in the backward. The replay pin (UNSLOTH_MOE_GC_REPLAY_PIN, opt-in since #1331) lets the backward reuse the replay's stack. This PR turns it on by default for stacks of at least _REPLAY_PIN_MIN_BYTES (512 MiB) that fit (_momentary_pin_fits, unchanged).

• UNSLOTH_MOE_GC_REPLAY_PIN=1 (or true / on / yes) keeps the previous opt-in: no size floor.
• UNSLOTH_MOE_GC_REPLAY_PIN=0 (or false / off / no) turns the pin off.
• Unset, empty or any other value means the new default: on from the floor up.
• UNSLOTH_MOE_RECOMPUTE still overrides both ways, and everything outside a replay is unchanged.

The floor is checked per decision, against the stack that decision is about:

• The bnb dispatcher decides on gate_up, as it does today, and then holds gate_up and down together.
• forward_native_grouped_mm (which the dispatcher falls back to when it declines, and which other callers such as the sonic fallback reach directly) checks each projection alone. At the default only stacks of at least 512 MiB pin there. On Qwen3-30B-A3B the gate_up stack (768 MiB) can pin and the down stack (384 MiB) recomputes. On gemma-4-26B the down stack (484 MiB) also stays under the floor.

Why revisit the opt-in, and what the floor does and doesn't cover

#1331 kept this off by default because it raised training peak memory (1.744 to 1.961 GiB on a tiny NF4 MoE, 23.7 to 24.0 GiB on gemma-4-26B) and the step-time gain did not reproduce on a quiet GPU.

• Tiny MoEs keep today's recompute under the floor. Faster MoE QLoRA: Triton NF4 dequant, fused combine, gemma-4 sliding layers at the window length #1331's tiny MoE held at least 0.21 GiB per layer (a pin raises peak by at most what it holds). I don't know its shapes, so I can't confirm it is under 512 MiB. The existing test_gc_recompute_pass_recomputes_for_4bit_by_default keeps passing unchanged, because its stacks are far below the floor.
• gemma-4-26B is unchanged by the floor. Its gate_up stack, (128, 1408, 2816) per Faster MoE QLoRA: Triton NF4 dequant, fused combine, gemma-4 sliding layers at the window length #1331, is 0.95 GiB, so it pins by default. Faster MoE QLoRA: Triton NF4 dequant, fused combine, gemma-4 sliding layers at the window length #1331's quiet-GPU result (it doesn't say which model) is the open question; I measured Qwen3-30B-A3B on an H100 and an L4, not gemma and not a B200.
• On Qwen3-30B-A3B the gain reproduced in three fresh runs per setting, with non-overlapping intervals (table below). It grows with the dequant's share of the step: about 16% at 256 rows per expert on the H100, where the step is 466 ms. It should be smaller at a B200's 8 TB/s.

Performance

unsloth/Qwen3-30B-A3B loaded by FastLanguageModel with load_in_4bit=True (3-D nested NF4 expert stacks, blocksize 64), LoRA r=16, use_gradient_checkpointing="unsloth", adamw_8bit, bf16, fixed 2048-token wikitext-2 sequences with the same data order and learning-rate schedule in every run. Three fresh processes per arm, alternating default / pin. Tokens/s over the steady steps (first 10 steps and the 2 profiled steps excluded); intervals are 95% t-intervals over the three runs, and the speedup interval is on the per-pair log ratio. torch 2.14.1+cu130, Triton 3.8.0, bitsandbytes 0.50.2, transformers 5.17.0, unsloth @ df4b4a5a83, unsloth_zoo @ ca97c8c. The pin arm sets UNSLOTH_MOE_GC_REPLAY_PIN=1, which takes the same dispatcher path this PR takes by default for these stacks (gate_up 768 MiB, above the floor). The end-to-end runs were not repeated on this branch.

• Dequant counts. Expert dequant launches were counted with the profiler over steps 35-36 (25-26 at batch 8) and averaged per step: 48 layers x 2 stacks x 3 passes = 288 by default, 192 with the pin. On the H100 the expert dequant time goes from 75.0 to 50.0 ms per step. On the L4 my 50 us duration cut-off also counts some attention-weight dequants, so only the difference is shown: 96 fewer launches and 887 -> 596 ms of dequant per step.
• L4 setup. On the L4, LoRA on attention + experts (642.5M trainable parameters) runs out of memory in 22 GiB on the default path at batch 1 and 2. The L4 row therefore uses LoRA on attention only; the frozen experts take the same dequant / pin path.
• Memory. The layer's gate_up and down stacks (0.75 + 0.375 GiB in bf16) are held from the replay through that layer's backward.

Numerics

Per layer, the pin is bit-identical to recompute. The check, on an RTX A5000 with this branch: a Qwen3-30B-A3B-shaped experts layer (random NF4 gate_up and down stacks of 128 experts, expert LoRA r=16), forward + backward inside the gradient-checkpoint replay marker, at 4,096 and 16,384 tokens, through both forward_moe_backend_bnb4bit (the path FastLanguageModel takes) and forward_native_grouped_mm, with the pin arm at the new default (variable unset) and at =1.

• The output, dX, routing-weight gradients and all four LoRA gradients are bit-for-bit equal between pin and recompute, and between two recompute runs.
• On the dispatcher path the pin run's check returned pin and its dequantize-and-swap path ran once per call, at the default as well as at =1; the recompute run did not take it.
• On the native path at the default only gate_up (768 MiB) is above the floor, so that arm pins gate_up and recomputes down.

End-to-end losses still vary run to run in both arms. Maximum per-step |difference| over all steps:

Validation

• Tests (RTX A5000, torch 2.14.1+cu130, bitsandbytes 0.50.2). tests/test_moe_bnb4bit_adaptive_recompute.py gives 1 failed and 19 passed on main (8b50cc4); the failure is the new default-pin test, ['provider'] != ['swap']. With this PR it gives 20 passed. The five other MoE test files (test_moe_recompute_gc_adaptive.py, test_moe_bnb4bit_selective_dequant.py, test_moe_sonic_dispatch_fallback.py, test_moe_selective_dequant_backward_view.py, test_moe_nf4_grouped.py): 103 passed, 1 skipped on both. The full suite was not run.
• Fit rule on a real GPU (unchanged by this PR; H100 and L4). A blocker tensor left free + cached memory about 256 MiB above and below 2.0 x a random NF4 stack of Qwen3-30B-A3B's gate_up shape (1.50 GiB threshold) and down shape (0.75 GiB). _momentary_pin_fits returned True above and False below in all 8 cases.
• Not covered end to end: A100, B200 and other cards (the tests and the bitwise check ran on an RTX A5000), multi-GPU (FSDP2 / DDP), larger MoEs (Qwen3-235B-A22B shapes would hold about 4.5 GiB per layer, which the fit check has to see room for), gemma-4-26B, models near the 512 MiB floor, other MoE families (gpt-oss experts do not load as Params4bit), and gradient checkpointing modes other than "unsloth".

Happy to keep it opt-in if you prefer; the data is here either way.

AI disclosure: I used Claude while writing this change and its tests; I reviewed every line, ran the tests on a rented GPU, and can answer questions about any of it.