From 395d6166e86f64f7c27f3c6e49cef89904add1aa Mon Sep 17 00:00:00 2001 From: Zaoyi Zheng <39311511+roy6n23@users.noreply.github.com> Date: Tue, 6 Oct 2026 02:34:33 +0000 Subject: [PATCH] fix(kl_div, tvd): make torch.compile work (NameError, then fp64 loop-carried scalar) (#1509) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 Co-authored-by: Arup De --- .../ops/backends/_ascend/ops/kl_div.py | 4 +-- .../ops/backends/_ascend/ops/tvd.py | 2 +- src/liger_kernel/ops/kl_div.py | 10 ++++++- src/liger_kernel/ops/tvd.py | 10 ++++++- test/transformers/test_kl_div.py | 29 +++++++++++++++++++ test/transformers/test_tvd.py | 22 ++++++++++++++ 6 files changed, 72 insertions(+), 5 deletions(-) diff --git a/src/liger_kernel/ops/backends/_ascend/ops/kl_div.py b/src/liger_kernel/ops/backends/_ascend/ops/kl_div.py index f7b2614f..7085589f 100644 --- a/src/liger_kernel/ops/backends/_ascend/ops/kl_div.py +++ b/src/liger_kernel/ops/backends/_ascend/ops/kl_div.py @@ -37,8 +37,8 @@ def _kldiv_kernel_forward( eps, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + reduction: tl.constexpr, # no default: always pass explicitly log_target: tl.constexpr = False, - reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN, ): pid = tl.program_id(0) num_progs = tl.num_programs(0) @@ -89,8 +89,8 @@ def _kldiv_kernel_backward( n_cols, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + reduction: tl.constexpr, # no default: always pass explicitly log_target: tl.constexpr = False, - reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN, ): pid = tl.program_id(0) num_progs = tl.num_programs(0) diff --git a/src/liger_kernel/ops/backends/_ascend/ops/tvd.py b/src/liger_kernel/ops/backends/_ascend/ops/tvd.py index 62a913a0..0dfe50f3 100644 --- a/src/liger_kernel/ops/backends/_ascend/ops/tvd.py +++ b/src/liger_kernel/ops/backends/_ascend/ops/tvd.py @@ -30,7 +30,7 @@ def _tv_distance_kernel( total_rows: tl.constexpr, # BT BLOCK_SIZE: tl.constexpr, HAS_LABEL: tl.constexpr, - reduction: tl.constexpr = "batchmean", + reduction: tl.constexpr, # no default: always pass explicitly ): thread_id = tl.program_id(0) num_threads = tl.num_programs(0) diff --git a/src/liger_kernel/ops/kl_div.py b/src/liger_kernel/ops/kl_div.py index d17da0a7..960aa540 100644 --- a/src/liger_kernel/ops/kl_div.py +++ b/src/liger_kernel/ops/kl_div.py @@ -54,8 +54,10 @@ def _kldiv_kernel_forward( n_cols, # int, number of columns in the input tensor eps, BLOCK_SIZE: tl.constexpr, + # no default: torch.compile re-emits this kernel's source in a module without the + # _REDUCTION_MODE_* globals, so a default naming one raises NameError. Always pass it. + reduction: tl.constexpr, log_target: tl.constexpr = False, - reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN, ): pid = tl.program_id(0).to(tl.int64) y_ptr += pid * y_stride @@ -64,6 +66,12 @@ def _kldiv_kernel_forward( base_offsets = tl.arange(0, BLOCK_SIZE) + # Pin the scalar to fp32. Eager Triton specializes a non-constexpr Python float to + # fp32, but Inductor specializes it to fp64; mixing an fp64 scalar into the loop body + # promotes the whole expression -- including tl.log -- to float64, and re-assigns the + # fp32 loop-carried `loss_sum` to fp64. This cast is a no-op on the eager path. + eps = eps.to(tl.float32) + loss_sum = 0.0 for i in range(0, n_cols, BLOCK_SIZE): offsets = i + base_offsets diff --git a/src/liger_kernel/ops/tvd.py b/src/liger_kernel/ops/tvd.py index aa3963b5..d5edacbd 100644 --- a/src/liger_kernel/ops/tvd.py +++ b/src/liger_kernel/ops/tvd.py @@ -52,7 +52,9 @@ def _tv_distance_kernel( scale, # pre-computed reduction scale for gradients (fused into kernel) BLOCK_SIZE: tl.constexpr, HAS_LABEL: tl.constexpr, - reduction: tl.constexpr = _REDUCTION_MODE_BATCHMEAN, + # no default: torch.compile re-emits this kernel's source in a module without the + # _REDUCTION_MODE_* globals, so a default naming one raises NameError. Always pass it. + reduction: tl.constexpr, ): pid = tl.program_id(0).to(tl.int64) p_ptr += pid * p_stride @@ -63,6 +65,12 @@ def _tv_distance_kernel( base_offsets = tl.arange(0, BLOCK_SIZE) + # Pin the scalar to fp32. Eager Triton specializes a non-constexpr Python float to + # fp32, but Inductor specializes it to fp64; mixing an fp64 scalar into the loop body + # promotes the gradient expression to float64 on every iteration and makes the + # compiled path disagree with eager. This cast is a no-op on the eager path. + scale = scale.to(tl.float32) + if HAS_LABEL: label = tl.load(label_ptr) if label == ignore_index: diff --git a/test/transformers/test_kl_div.py b/test/transformers/test_kl_div.py index 84386f4e..491923c7 100644 --- a/test/transformers/test_kl_div.py +++ b/test/transformers/test_kl_div.py @@ -102,3 +102,32 @@ def test_correctness_not_last(B, T, V, log_target, reduction, dtype, atol, rtol) log_target, is_last_layer=False, ) + + +@pytest.mark.parametrize("log_target", [True, False]) +@pytest.mark.parametrize("reduction", ["batchmean", "sum", "mean", "none"]) +def test_torch_compile_matches_eager(log_target, reduction): + # inductor re-emits the kernel source; a signature default that names a module global must + # still resolve there, and eps must stay fp32 (inductor passes Python floats as fp64). The + # unreduced loss matches eager bit for bit; a reduced loss may differ in the last bits + # because inductor fuses the wrapper's sum/division + torch._dynamo.reset() + torch.manual_seed(0) + liger_kldiv = LigerKLDIVLoss(reduction=reduction, log_target=log_target) + compiled = torch.compile(liger_kldiv) + input = torch.randn(64, 1271, device=device, dtype=torch.float32).log_softmax(dim=-1) + with torch.no_grad(): + target = torch.randn(64, 1271, device=device).softmax(dim=-1) + if log_target: + target = target.log() + x1 = input.detach().clone().requires_grad_(True) + x2 = input.detach().clone().requires_grad_(True) + out1 = liger_kldiv(x1, target) + out2 = compiled(x2, target) + if reduction == "none": + assert torch.equal(out1, out2) + return + torch.testing.assert_close(out2, out1) + out1.backward() + out2.backward() + torch.testing.assert_close(x2.grad, x1.grad) diff --git a/test/transformers/test_tvd.py b/test/transformers/test_tvd.py index 1c510557..2d3aa987 100644 --- a/test/transformers/test_tvd.py +++ b/test/transformers/test_tvd.py @@ -237,3 +237,25 @@ def test_correctness_with_ignore_index(B, T, V, reduction, dtype, atol, rtol, ig liger_tvd = LigerTVDLoss(reduction=reduction, ignore_index=ignore_index) torch_tvd = TorchTVDLoss(reduction=reduction, ignore_index=ignore_index) _test_correctness_with_ignore_index_once(liger_tvd, torch_tvd, ignore_index, B, T, V, dtype, atol, rtol, reduction) + + +@pytest.mark.parametrize("reduction", ["batchmean", "sum", "mean", "none"]) +def test_torch_compile_matches_eager(reduction): + # inductor re-emits the kernel source; a signature default that names a module global must + # still resolve there, and the compiled call must match eager bit for bit + torch._dynamo.reset() + torch.manual_seed(0) + device = infer_device() + liger_tvd = LigerTVDLoss(reduction=reduction) + compiled = torch.compile(liger_tvd) + p = torch.randn(64, 1271, device=device).softmax(dim=-1) + q = torch.randn(64, 1271, device=device).softmax(dim=-1) + x1 = p.detach().clone().requires_grad_(True) + x2 = p.detach().clone().requires_grad_(True) + out1 = liger_tvd(x1, q) + out2 = compiled(x2, q) + assert torch.equal(out1, out2) + if reduction != "none": + out1.backward() + out2.backward() + assert torch.equal(x1.grad, x2.grad)