Fix FLCE torch.compile failure from aten.addmm.dtype_out overload

[yashb98]

Summary

Fixes #1327.

FLCE's grad_weight fast path calls torch.addmm(..., out_dtype=torch.float32, out=...), which dispatches to aten.addmm.dtype_out. TorchInductor cannot lower that overload yet (the Tier-2 fix, pytorch/pytorch#190936, is still open), so any torch.compile'd FLCE backward with bf16/fp16 inputs and accum_dtype=torch.float32 dies at lowering time with:

LoweringException: TypeError: tuned_addmm() takes 3 positional arguments but 4 were given

This PR gates the fast path on not torch.compiler.is_compiling(). Under torch.compile the backward takes the existing generic matmul path, which Inductor lowers fine; eager mode is unchanged and keeps the fused addmm accumulation. When the PyTorch-side fix lands this guard can be revisited.

Details

• One-line guard in fused_linear_cross_entropy_forward; no behavioural change outside torch.compile.
• Adds test_torch_compile_fp32_accum_backward: compiles FLCE with bf16 operands + accum_dtype=torch.float32, runs backward, and checks loss / grad_x / grad_w against eager. It fails on current main with the exact tuned_addmm() signature from the issue and passes with this change.
• Verified the compiled fallback is numerically correct, not just non-crashing: grad_x is bitwise identical to eager, grad_w differs only at bf16 rounding (max abs diff 0.00195 on the repro shape).

Testing Done

Reproduced the issue on current main (9ac26bb) with the script from #1327 before changing anything:

PyTorch: 2.13.0+cu130, Triton 3.7.1, GB10 (compute capability 12.1)
Result:  REPRODUCED - LoweringException: TypeError: tuned_addmm() takes 3 positional arguments but 4 were given

After the change: Result: NOT REPRODUCED (compiled backward completes).

• test/transformers/test_fused_linear_cross_entropy.py: 142 passed (was 141 + the new regression test)

• New test, unpatched main: fails with the tuned_addmm() error above

• New test, with this change: passes

• Full unit suite (pytest --ignore=test/convergence test/), patched: 3916 passed, 2 failed — the 2 failures are test_grpo_loss_vs_trl[...-luspo-sequence-1.5], and the identical 2 fail on unpatched main in the same full-suite run (3915 passed, 2 failed). They pass in isolation in both arms and live in grpo_loss, which this change does not touch: a pre-existing, order-dependent flake.

• ruff check and ruff format --check on both touched files: clean

• Hardware Type: NVIDIA GB10 (DGX Spark, sm_121, aarch64)

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

One honest caveat: I did not run make test-convergence (model-scale convergence runs); the change only reroutes an existing eager path under compilation, so convergence behaviour is unaffected in eager and the compiled path is checked against eager in the new test.

Benchmarks

GB10 (DGX Spark, sm_121), torch 2.13.0+cu130, bf16 operands with fp32 accumulation (the affected configuration), forward + backward, median of 8 runs after 3 warmup. Branch is rebased on current main (91ae44a).

shape	eager, main	eager, this PR	torch.compile, main	torch.compile, this PR
2048 tok, H=1024, V=151936 (Qwen3-0.6B class)	1232.9 ms	1238.5 ms (+0.5%)	LoweringException (the bug)	780.3 ms
4096 tok, H=4096, V=128256 (Llama3-8B class)	1047.1 ms	1053.3 ms (+0.6%)	LoweringException (the bug)	709.8 ms

Two takeaways: eager is unchanged (parity within noise, as expected since the guard only fires under compilation), and the compiled path this PR unblocks is not merely working but about 1.5x faster than eager on both shapes. The unpatched compiled run still fails with the exact tuned_addmm() LoweringException from the issue, re-confirmed on main at 91ae44a today.

[yashb98]

CI has never run on this PR, all four workflows have been sitting at action_required since 2026-07-31. Could a maintainer approve the run?

In the meantime, here is a before/after against current main @ 08e8ccdc. GB10 (DGX Spark, sm_121, aarch64), torch 2.13.0+cu130, Triton 3.7.1, same box as the PR description.

Unpatched main with only the new test applied:

E       torch._inductor.exc.InductorError: LoweringException: TypeError: tuned_addmm() takes 3 positional arguments but 4 were given
E         target: aten.addmm.dtype
E          File ".../liger_kernel/ops/fused_linear_cross_entropy.py", line 249, in torch_dynamo_resume_in_fused_linear_cross_entropy_forward_at_93

1 failed, 141 deselected, 15 warnings in 2.86s

Same tree plus the one line from this PR:

1 passed, 141 deselected, 15 warnings in 2.23s

Whole file:

142 passed, 15 warnings in 5.18s

One thing that may be worth recording in the issue: on torch 2.11.0+cu130 the same test fails earlier and differently, in Dynamo rather than in Inductor lowering, with TypeError: unsupported operand type(s) for *: 'torch.dtype' and 'FakeTensor' at the same line 249. That is the Tier 1 signature #1327 describes for torch versions without the core decomposition fix. Both signatures come from the same aten.addmm.dtype_out dispatch, and this one-line guard fixes both; I ran the patched test on 2.11.0 and 2.13.0 and it passes on each.

The branch shows as behind by 3 commits, but neither file it touches is among the files those commits changed, and the runs above are on 08e8ccdc itself, so the rebase is already covered. Happy to rebase whenever you want it.