mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
395d6166e8
## Summary On current PyTorch, `torch.compile` of `LigerKLDIVLoss` and `LigerTVDLoss` fails for every reduction mode. There are two stacked causes; the second only shows once the first is fixed. **1. `NameError` from the `reduction` default.** Both kernels declare `reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN`, a module global. Inductor re-emits a user-defined Triton kernel's source in its own module and brings along the globals the body uses, but not names that only appear in a default, so evaluating the default fails: ``` InductorError: NameError: name '_REDUCTION_MODE_BATCHMEAN' is not defined ``` The callers always pass `reduction`, so it is now a required argument with no default (a forgotten argument raises `TypeError` instead of silently picking a reduction). **2. fp64 loop-carried scalar (kl_div).** Continues the sweep offered in #1350, same root cause: a non-`constexpr` Python float is fp32 under eager Triton but fp64 under Inductor. In `_kldiv_kernel_forward`, `eps` feeds `tl.log(tl.maximum(y_true, eps))` inside the reduction loop, so with fix 1 alone `log_target=False` fails like #1362 did: ``` AssertionError('Loop-carried variable loss_sum has initial type fp32 but is re-assigned to fp64 in loop!') ``` `eps` (kl_div) and `scale` (tvd) are now cast with `.to(tl.float32)` before the loop, as in `rms_norm.py` and #1350/#1362. That is a no-op on the eager path. The `_ascend` mirrors get the same signature change (review feedback); the fp32 casts are not applied there, matching the scope of #1350/#1362. I have no Ascend hardware, so those kernels are not run. ## Testing Done H100 80GB HBM3, torch 2.13.0+cu130, Triton 3.7.1. GPU tests ran on `main` at `492bfeb`; the branch is rebased onto `b297821`, which does not touch these files (`make checkstyle` re-run after the rebase). | | `torch.compile` of kl_div / tvd | |---|---| | `main` (`492bfeb`) | every case fails: `NameError` above | | fix 1 only | tvd compiles; kl_div with `log_target=False` fails with the fp64 loop-carried error above | | fix 1 + fix 2 (this PR) | compiles in every reduction / `log_target` case | - New `test_torch_compile_matches_eager` in `test_kl_div.py` (8 cases) and `test_tvd.py` (4 cases). The unreduced output matches eager bit for bit (`torch.equal`); reduced losses use `assert_close`. - Why not bit-identical for reduced kl_div: I compared the kernel's per-row output (before the wrapper's `.sum()`) between eager and compiled, and it is bit-identical in all 6 `reduction × log_target` cases. Compiling only the `.sum()` / division on the eager rows reproduces the compiled loss exactly, so the remaining difference (1.2e-7 for `batchmean`, 7.6e-6 for `sum`) comes from Inductor's own reduction order, not from the kernel. - `test_kl_div.py` + `test_tvd.py`: **340 passed**. Re-run on the follow-up commit `f34c737` (required `reduction`, same GPU model and versions): 340 passed. - `make checkstyle`: passes. - `make test`: not clean, so not checked. I ran it on the same box with the PyTorch / transformers versions above, which are newer than CI. `test_grpo_loss.py::test_grpo_loss_with_vllm_is_ratio_reduced[2-128-1000-0.04-luspo-sequence]` and `test_fused_linear_jsd.py::test_low_probability_vocab_block_is_finite` failed. After the `test_grpo_loss` failure, every later test in the same process errored at setup (including `test_kl_div.py`, which passes 340/340 when run on its own), and the run hit my 40-minute cap at 84%. Both failures are in files this PR does not touch. I don't have a `main` baseline from that box to show they are pre-existing. - Hardware Type: H100 80GB HBM3 - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence I used an AI assistant (Claude) to help investigate and write this change; I reviewed and ran everything above. @vaibhavjindal --------- Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Arup De <[email protected]>