`compile_model=true` silently trains on NaN loss and crashes at `max-autotune` — miscompiled `make_att_2d_masks` in SmolVLA / PI0 / PI05

[zuoxingdong]

Ticket Type

🐛 Bug Report (Something isn't working)

Environment & System Info

- LeRobot version: main @ 95211b98 (0.6.0 dist metadata)
- Platform: Linux-5.15.92-x86_64-with-glibc2.35
- Python version: 3.12.13
- Huggingface Hub version: 1.23.0
- Transformers version: 5.5.4
- Datasets version: 4.8.5
- Numpy version: 2.2.6
- FFmpeg version: 7.1.1
- PyTorch version: 2.11.0+cu128
- Is PyTorch built with CUDA support?: True
- Cuda version: 12.8
- GPU model: NVIDIA H100 80GB HBM3
- Using GPU in script?: Yes

Description

Description

• What breaks: the shipped --policy.compile_model=true flag.
• Where: make_att_2d_masks (policies/common/vla_utils.py) — prefix-LM mask from cumsum + comparison ops.
• Why: inductor fuses these ops into a generated triton split-scan kernel (triton_spl_fused__to_copy_cat_cumsum_...) that is miscompiled at small dynamic shapes.
• Worst part: in compile_mode=default the run looks successful — loss is NaN from step 1, no error raised ⚠️.
• Fix (+7 lines, linked PR): @torch.compiler.disable on the helper — ~6 tiny mask ops run eagerly, everything else compiles (one graph break, ~0 cost).

Closest upstream: pytorch#131027, pytorch#124251.

symptom (policy + compile_model=true)	torch 2.11	torch 2.13
compile_mode=default → silent NaN loss ⚠️	❌	❌ still broken
fused-attention graphs + max-autotune → CUDA crash	❌	✅ fixed ≥2.12

Real training logs 🔎

compile_mode=default — silent NaN (run "succeeds"):

step:10 ... loss:nan grdn:nan ... losses_after_forward:nan
step:20 ... loss:nan grdn:nan ... losses_after_forward:nan

max-autotune + fused-attention graph — crash in the generated split-scan kernel:

File "/tmp/torchinductor_.../cuh3ijxww....py", line 5018, in partition_0
  triton_spl_fused__to_copy_cat_cumsum_ones_sub_unsqueeze_7.run(primals_9, buf9, buf31, workspace_0, 8, 194, stream=stream0)
RuntimeError: CUDA driver error: operation not supported on global/shared address space

(second flavor seen across runs: CUDA driver error: an illegal memory access was encountered)

Affected policies (audit of all src/lerobot/policies/* on 95211b98):

policy	make_att_2d_masks	compile_model flag	status
SmolVLA	✅	✅	🐛 reproduced
PI0	✅	✅	🐛 structurally affected¹
PI05	✅	✅	🐛 structurally affected¹
PI0-FAST	❌ (cumsum for position ids only)	✅	✅ not the affected pattern
Diffusion	❌	✅	✅ not affected
ACT, wall-x, evo1, eo1, groot, lingbot-va, vla-jepa, xvla, …	❌	❌ (no flag)	✅ not affected²

¹ identical code pattern and flags; not separately reproduced.
² lingbot-va compiles its own FlexAttention path; evo1/wall-x/act use cumsum for
indexing/position ids only — none route make_att_2d_masks through torch.compile.

Context & Reproduction

lerobot-train --policy.type=smolvla --policy.compile_model=true --policy.compile_mode=default ...
# → loss:nan from the first logged step, run "succeeds"

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem (torch 2.11 & 2.13, cu126/cu128/cu129).

Linked PR: #4179

Context & Reproduction

No response

Relevant logs or stack trace

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem.

Additional Info / Workarounds

No response

[raul-machine-learning]

Hi, could you paste the full, exact command line you are using to train the model using SDPA? Ideally using a public dataset that is known to work so we can test it.