Compile the Laya encoder for decision model training - #12778

Builds on #12585 (decision model fine-tuning for Laya and Clef), which is now merged. It now also carries @Etherll's #12824 (merged here with its commits kept), so the Laya training speedups are one PR.

Laya training is launch bound: at the default recipe (LoRA, 8 decisions of about 300 tokens per micro-batch) the GPU is busy for about 0.14 s of a 1.6 s step on a B200, and the time goes to many small ModernBERT kernels. Compiling each encoder layer fuses them.

What changes

• DecisionTrainer.train() compiles each Laya encoder layer in place (nn.Module.compile(dynamic = True)) for the run and puts the layers back to eager afterwards, so calibration, evaluation and saving run eagerly and parameter names, LoRA merging and saving are unchanged.
• Before training, one compiled forward and backward runs as a warm-up (random state and gradients left untouched). If Inductor cannot serve the platform, the layers go back to eager and the run trains as before instead of failing.
• On by default for GPU runs of at least 4,000 forward passes, where the 30 to 90 s compile pays for itself. UNSLOTH_DECISION_COMPILE=0 turns it off, UNSLOTH_DECISION_COMPILE=1 forces it on. Clef models are not affected (their head already compiles in Studio: fine-tune Laya and Cloudflare Clef decision models and serve them #12585).

Also from #12824

• Short decisions (at most 1,024 tokens, which covers Laya's data) compile with static 64-token shape buckets instead of dynamic shapes; padded positions are masked keys, so the loss is unchanged. Evaluation stays eager.
• The compiled layers use transformers' plain SDPA call (unsloth_zoo's attention wrapper made dynamo fall back on torch 2.11), and cuDNN SDPA is kept out of the decision forward, backward recompute and scoring (it rebuilds its plan for every new sequence length).
• Plain LoRA layers (one adapter, no dropout, no DoRA) skip PEFT's per-call checks and dtype round trip; identical to PEFT in fp32 and in the autocast dtype, and PEFT's own forward is used for everything else. Checked on PEFT 0.18.0 and 0.21.2.
• Scoring batches similar lengths together, so calibration and evaluation pad less.
• Compiled layers checkpoint with torch's own non-reentrant checkpoint: Unsloth's offloaded reentrant checkpoint gave compiled bf16 layers wrong gradients. torch's checkpoint is found even after unsloth_zoo patches torch.utils.checkpoint.checkpoint.

On one B200, LoRA r64 and full fine-tune at 2 decisions x 16 accumulation, interleaved runs: #12824's additions are about 1.12x (LoRA, 1.25 to 1.08 s/step) and 1.10x (full, 0.71 to 0.65 s/step) faster than the dynamic-shape compile alone, same peak memory, 0 graph breaks.

Results

Colab L4, one card per run, arms interleaved ABBA, fresh process and compile cache per arm, 24 steps, median s/step (identical base runs differed by about 1%):

• Over a full epoch: 0 graph breaks, 4 graphs, 0 recompiles.
• Full fine-tune gradients in float32 match eager (cosine 1.0); in bfloat16 both are equally close to the float32 result (cosine 0.9994 compiled, 0.9993 eager).
• Train, calibrate, merged save and reload through the vendored Laya: served answers agree with the trained model on 200/200 (LoRA) and 199/200 (full fine-tune), largest probability difference 0.013.

Merged with #12824: speed and accuracy per GPU

Pre-merge (7052f9d) vs merged (16c55e2), ABBA order, fresh process and compile cache per run, steady s/step from step 3:

Same-arm spread is at most 3% (9% on the RTX PRO 6000). Full fine-tune peak memory is unchanged; LoRA at 8x8 allocates 2.6 to 5% more from bucket padding.

Gradients vs an fp64 eager step, 30 to 40 batches per GPU: in bf16 the merged compile is level with eager (A100 0.99x, L4 1.10x, B200 tie; RTX PRO 6000 1.22x, CI [1.02, 1.51]); in fp32 compiled matches fp64 on all GPUs. A single batch can show a large bf16 gap in either direction (eager was the worst arm on L4 and A100 for one batch): it comes from bf16 rounding in layers 0 to 10, and each compiled layer alone is as accurate as eager. Holdout loss after 60 steps (3 seeds per arm, B200 and L4) moved by at most 0.015 CE, well inside the 0.03 to 0.05 seed spread.

Not included

• Compiling the whole training step (forward plus loss, fullgraph=True, no graph breaks) was another 1.13x over this, but takes 13 to 17 minutes to compile on an L4 and cannot trace Unsloth's offloaded gradient checkpointing, so it needs 3x the memory. Not worth it at Laya's typical run length.
• Turning gradient checkpointing off is 1.37x on top of this, at 1.53 GB to 4.6 GB peak. The default stays on.
• Fused AdamW and compiled autograd gave nothing measurable.

Tests

tests/test_decision_model.py covers the run-length threshold, both overrides, the return to eager, and a compile that raises (training still finishes, eagerly). The decision suite passes on CPU and on GPU with the compile on auto and forced on.