mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
main
537
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Vaibhav JindalandCopilot
|
74777f44fb |
fix(test): drop unsupported device kwarg from config-based RoPE con… (#1514)
## Summary Fixes the cuTile / cuTeDSL / correctness RoPE CI jobs that started failing across both H100 and B200 (e.g. run 37538343198) with: ``` TypeError: LlamaRotaryEmbedding.__init__() got an unexpected keyword argument 'device' ``` ## Root cause Newer `transformers` (observed on 5.19.0 in CI) removed the `device` keyword from the **config-based** `*RotaryEmbedding.__init__` (`__init__(self, config)`). Our `transformers_version_dispatch(...)` helper still forwarded `device` in `after_kwargs`, so every RoPE test that builds a rotary embedding raised the `TypeError` **at setup**, before any kernel ran — which is why all cuTile + cuTeDSL jobs failed on both architectures at once. The on-pod/mirror `transformers` (<= 5.15.0) still accepted `device`, so it only surfaced in CI where public PyPI resolves a newer version. ## Fix The config-based constructor derives device from the input tensor in `forward` (`inv_freq` is moved to `x.device`), so passing `device` at construction is redundant. This PR drops `"device": device` from every `after_kwargs`, while keeping it in `before_kwargs` for the pre-4.48.0 positional-`dim` path. Changed call sites: - `test/transformers/test_rope.py` (3) - `test/cutile/test_rope.py` (3) - `test/cutedsl/test_rope.py` (3) - `benchmark/scripts/benchmark_rope.py` (1) - `benchmark/scripts/benchmark_llama4_rope.py` (1) - `src/liger_kernel/utils.py` — docstring example updated ## Verification - `test/transformers/test_rope.py`: **81 passed** on transformers 5.15.0 (backward-compatible). - Simulated the 5.19.0 strict signature (reject `device`): the **old** pattern raises the exact CI `TypeError`; the **new** pattern constructs successfully and `forward` returns `cos`/`sin` on the input's CUDA device. - `ruff check` + `ruff format --check`: clean. - cuTile/cuTeDSL backends aren't installed on the dev pod, but all three test files use the identical dispatch pattern, so the fix applies uniformly. Co-authored-by: Copilot <[email protected]> |
||
|
|
b5cdbf7e34 |
perf(norm): skip dW/dB calculation when parameters are frozen (#1067) (#1440)
## Description Addresses #1067 (previously explored in #1068) This PR brings the LoRA normalization optimization across the finish line: - Rebased on top of modern `main` (cleanly resolving the #1116 benchmark conflicts that blocked #1068). - Incorporates maintainer feedback from @Tcc0403 on #1068 (passes `None` directly for unused gradient pointers). - Comprehensive unit tests: 24/24 passing on Float32 and BFloat16. - Hardware-validated on NVIDIA A10G and NVIDIA A100 GPUs (up to 1.38x backward speedup). Full co-author credit is preserved for @yukiu00 in the commit. --- ### Motivation In parameter-efficient fine-tuning (LoRA / PEFT), normalization parameters (weights and biases) are frozen (`requires_grad=False`). Previously, Liger-Kernel unconditionally: 1. Allocated intermediate reduction buffers (`_dW`, `_DB` of size `sm_count * n_cols`) in HBM. 2. Allocated registers and ran Triton accumulation loops across tokens. 3. Wrote intermediate gradients to global HBM. 4. Launched PyTorch `.sum(dim=0)` global reductions to produce tensors that PyTorch autograd immediately discarded. --- ### Key Changes - **`rms_norm.py`**: Added `COMPUTE_DW: tl.constexpr` to `_rms_norm_backward_kernel` and `_block_rms_norm_backward_kernel`. Prunes accumulation and memory writes when `compute_dw=False`; skips `_dW` buffer allocation and reduction; queries `ctx.needs_input_grad[1]`. - **`layer_norm.py`**: Added `COMPUTE_DW` and `COMPUTE_DB` flags to prune weight/bias gradients independently based on `ctx.needs_input_grad[1]` and `[2]`. - **`fused_add_rms_norm.py`**: Added `COMPUTE_DW` to skip `_dW` allocation and reduction based on `ctx.needs_input_grad[2]`. - **`test/ops/test_norm_lora_parity.py`**: Added 24 unit tests verifying numerical parity (FP32 & BF16), checking that `w.grad is None` / `b.grad is None` when frozen, and confirming exact `x.grad` equality. - **`benchmark/scripts/benchmark_norm_lora.py`**: Added dedicated standalone LoRA backward microbenchmark. --- ### Benchmark Results #### NVIDIA A100-SXM4-40GB (`triton.testing.do_bench`, bfloat16) | Operation | Shape (B, D) | Full FT Bwd (ms) | LoRA Bwd (ms) | Speedup | | :--- | :--- | :--- | :--- | :--- | | **LayerNorm Bwd** | (2048, 4096) | 0.0806 | 0.0585 | **1.38x (+38%)** | | **RMSNorm Bwd** | (2048, 4096) | 0.0871 | 0.0707 | **1.23x (+23%)** | | **LayerNorm Bwd** | (4096, 4096) | 0.1217 | 0.0991 | **1.23x (+23%)** | | **RMSNorm Bwd** | (4096, 4096) | 0.1305 | 0.1177 | **1.11x (+11%)** | | **LayerNorm Bwd** | (4096, 8192) | 0.1966 | 0.1729 | **1.14x (+14%)** | | **RMSNorm Bwd** | (4096, 8192) | 0.2275 | 0.2135 | **1.07x (+7%)** | | **LayerNorm Bwd** | (8192, 4096) | 0.2070 | 0.1843 | **1.12x (+12%)** | | **RMSNorm Bwd** | (8192, 4096) | 0.2384 | 0.2243 | **1.06x (+6%)** | | **LayerNorm Bwd** | (8192, 8192) | 0.3514 | 0.3255 | **1.08x (+8%)** | | **RMSNorm Bwd** | (8192, 8192) | 0.4235 | 0.4077 | **1.04x (+4%)** | #### NVIDIA A10G (`triton.testing.do_bench`, bfloat16) - **LayerNorm Bwd**: **1.05x – 1.18x speedup** across shapes. - **RMSNorm Bwd**: **1.02x – 1.07x speedup** across shapes. --- ### Test Validation - `test/ops/test_norm_lora_parity.py`: **24 / 24 PASSED (100%)** - Existing ops correctness: **108 / 108 PASSED (100%)** - `make checkstyle` (`ruff check` & `ruff format`): **PASSED** Co-authored-by: Yuki Uehara <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Arup De <[email protected]> |
||
|
|
590a608574 |
fix(monkey_patch): let the norm flags cover the whole vision tower (#1443)
## Summary Five patchers gate the *per-layer* norms of a vision tower on the corresponding flag, but patch the tower's *own* norms unconditionally. Passing the flag as `False` together with a `model=` instance therefore still swaps Liger norms in: | patcher | flag | patched regardless | | --- | --- | --- | | `llama4`, `mllama` | `layer_norm` | `vision_model.layernorm_pre`, `layernorm_post` | | `gemma3`, `paligemma` | `layer_norm` | `vision_model.post_layernorm` | | `glm4v_moe` | `rms_norm` | `visual.post_conv_layernorm`, `post_layernorm` | The class-level branch of the very same functions gets it right (`if layer_norm and model is None:` in mllama/gemma3/paligemma, `if rms_norm:` in glm4v_moe), so one function answers the same flag differently depending on whether you hand it a model. And `smolvlm`, `internvl`, `muse_glimmer` and `glm4v` already gate the tower norm exactly like their encoder layers, which is why these five read as oversights rather than deliberate exceptions. Both flags default to `True` on all five, so nothing changes unless a caller explicitly opts out. ## Testing Done Reproduced on transformers 5.9.0 / torch 2.11.0 before the change, patching an instance with the flag off: ``` paligemma (layer_norm=False) vision_model.post_layernorm PATCHED <-- ignored the flag encoder.layers[0].layer_norm1 untouched mllama (layer_norm=False) vision_model.layernorm_pre PATCHED <-- ignored the flag vision_model.layernorm_post PATCHED <-- ignored the flag transformer.layers[0].input_layernorm untouched ``` gemma3, llama4 and glm4v_moe behave the same way. After the change every one of those reads `untouched`. - `pytest test/transformers/test_monkey_patch.py`: 63 passed / 1 skipped on `main`, 68 passed / 1 skipped here. The five added tests each fail on `main` and pass with the fix; the skip is `muse_glimmer` missing from the installed transformers. - `ruff check .` and `ruff format --check .` with the v0.14.11 that `.pre-commit-config.yaml` pins: clean. - Convergence is untouched by construction. Both `test_mini_models_multimodal.py` files pin `"rms_norm": True` and `kwargs["layer_norm"] = True` for these models, glm4v_moe is not in the convergence set at all, and this change only alters the `False` branch. Also ran on GPU. This branch against `main`, both on the same A100 40GB in one job: | | this branch | `main` | | --- | --- | --- | | `test_mini_models_multimodal.py` bf16 (`-k "mllama or paligemma or gemma3 or llama4 or glm4v"`) | 5 passed | 5 passed | | the same, fp32 | 4 passed, 1 xfailed | 4 passed, 1 xfailed | | `test_monkey_patch.py` | 68 passed, 1 skipped | 63 passed, 1 skipped | Full `make test` scope, separately on an A100 80GB: 4496 passed, 18 failed, 931 skipped, 15 xfailed. Re-running just `test/chunked_loss/test_grpo_loss.py` on `main` gives the same 18 failing test ids (`luspo`/`cispo` cases that miss tolerance on sm80), and that file does not import `monkey_patch`. Hardware Type: A100 40GB and A100 80GB, plus CPU for the unit suite. The patchers are pure Python and `test_monkey_patch.py` has no CUDA gating, so that suite is the one that actually exercises them. Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Arup De <[email protected]> |
||
|
|
fc330e197e |
perf(distillation): skip the hard (CE) loss when it is neither weighted nor returned (#1489)
## Summary `LigerFusedLinearDistillationBase` computes the student's float32 log-softmax over the vocabulary, and the NLL (the hard loss), on every chunk even when `weight_hard_loss == 0`. In that case the hard loss is multiplied by zero and discarded. This happens in pure distillation, for example `LigerFusedLinearJSDLoss(weight_hard_loss=0.0, weight_soft_loss=1.0)`. This PR computes the hard loss, and the log-softmax it needs, only when `weight_hard_loss != 0` or when `return_soft_hard_loss=True` asks for it. Loss and gradients are unchanged: the existing correctness tests already cover `weight_hard_loss=0` against the torch reference. The hard loss is still reported whenever it is requested. ## Details Forward + backward in bf16, `LigerFusedLinearJSDLoss(weight_hard_loss=0.0, weight_soft_loss=1.0, beta=0.5)`, trainable student head, B·T = 4096, H = 4096. Each figure is the median of 10 timed iterations after 3 warm-ups, then the median over separate processes, with main and this branch run alternately. | V | compiled | main | this PR | |---|---|---|---| | 128,256 | yes (default) | 42.0 ms, 6.51 GiB | 41.9 ms, 6.51 GiB | | 151,936 | yes (default) | 71.1 ms, 7.98 GiB | **66.1 ms, 7.69 GiB (−7%)** | | 128,256 | no | 85.8 ms, 8.75 GiB | **75.5 ms, 8.26 GiB (−12%)** | | 151,936 | no | 101.1 ms, 10.34 GiB | **89.0 ms, 9.75 GiB (−12%)** | This touches the same function as #1487 (frozen-parameter gradients in the same base). The two are independent, but whichever lands second needs a small rebase. Under `torch.compile`, the compiler already absorbs some of the redundant work, so the saving depends on the shape. With `weight_hard_loss=0.5` the timings are unchanged (42.0 vs 42.1 ms at V = 128,256; 71.0 vs 71.1 ms at V = 151,936). ## Testing Done - Hardware Type: NVIDIA H100 80GB HBM3 (torch 2.14.0+cu130, triton 3.8.0) - [ ] run `make test` to ensure correctness: I ran the test files for the changed base rather than the full suite. `test/chunked_loss/test_jsd_loss.py` and `test/chunked_loss/test_cosine_loss.py`: 110 passed on this branch (104 on main, plus 6 new). They also pass on CPU. - [x] run `make checkstyle` to ensure code style (ruff 0.15.22 `check` and `format --check` on the changed files) - [ ] run `make test-convergence` to ensure convergence: not run; the convergence tests do not exercise the chunked losses. New tests, for both losses: - The CE is computed exactly when it is weighted or returned: skipped at weight 0 without a return, and computed at weight 0 with a return and at weight 0.5. - With `weight_hard_loss=0` and `return_soft_hard_loss=True`, the returned hard loss equals the cross-entropy. Co-authored-by: Jari Kolehmainen <[email protected]> Co-authored-by: Arup De <[email protected]> |
||
|
cherry77-cloudandArup De
|
29eb8ca223 |
Fix Falcon-H1 output scaling in LCE forward (#1477)
## Summary Falcon-H1's patched `lce_forward` omits `self.model.lm_head_multiplier` in both the fused CE and materialized-logits paths. For models with a non-unit multiplier, including the official Falcon-H1 checkpoints, this changes the class probabilities and training objective. Setting `skip_logits=False` also produces incorrectly scaled logits. Apply the multiplier after `lm_head` when materializing logits. For fused CE, fold it into the hidden states before the bias-free output projection, preserving `CE(cHW^T, y)` without materializing the full logits. The reordered multiplication can introduce low-precision rounding differences, which are covered by the local numerical checks below. ## Testing Done - Hardware: NVIDIA RTX A1000 Laptop GPU, 4 GB; PyTorch 2.13.0+cu130, Transformers 5.14.1, Triton 3.7.1. - 72 local regression cases passed, comparing the upstream and patched forwards with multipliers `1.0`, `0.013020833333333334`, and `0.01953125` in FP32/BF16. Coverage includes logits, loss, hidden-state/head gradients, both loss paths, batch loss normalization, and inference slicing. Before the fix, all 48 non-unit cases failed. These local test additions are not included in this production-code-only patch. - Existing Falcon-H1 instance-patching test passed. - All four existing Falcon-H1 convergence cases passed: FP32/BF16, fused CE/with logits, 32 training steps each. The local runner used batch size 1 and the 7B multiplier to fit the 4 GB GPU; architecture, assertions, and tolerances were unchanged. The original batch size 16 exhausted GPU memory in the upstream baseline. - [ ] `make test` (full suite not run; focused validation above) - [x] `make checkstyle` - [ ] `make test-convergence` (full suite not run; four Falcon-H1 cases run as described above) Co-authored-by: Arup De <[email protected]> |
||
|
LOGO127andArup De
|
c508f2b4aa |
fix(grpo): respect independent token and mask strides (#1461)
## Summary
Fix the default Triton selective-log-softmax helper using the token-ID
row stride to address the mask and assuming unit column strides for both
inputs.
The helper independently slices token IDs and the mask to the trailing
completion length. Those views can have different row/column strides,
including zero strides from expansion. The old kernel can silently
select the wrong mask entry or token ID. Pass each view's two strides to
the kernel; keep logit-contiguity requirements, output layout,
temperature scaling and loss math unchanged. No new input
materialization is introduced.
## Testing Done
- Hardware Type: NVIDIA A100-80G-PCIe.
- Base:
|
||
|
luca huangandArup De
|
e2039eb8c7 |
fix(cutedsl): fuse Gemma RMSNorm weight offset in FP32 (#1462)
## Summary The CuTe DSL Gemma paths in RMSNorm and FusedAddRMSNorm add the offset to BF16/FP16 weights before promoting them to FP32. This introduces extra rounding in the forward output and input gradient. For example, BF16 `0.00390625 + 1` rounds to `1.0`, while the FP32 sum is `1.00390625`. Pass the weight and offset separately, then add the offset in FP32 inside the forward and backward kernels. This also removes the separate weight-add kernel and temporary vector. The fused offset handling follows [Quack](https://github.com/Dao-AILab/quack/blob/10ce505a7ecf68e46753009bd08dcf9bd6c059ee/quack/rmsnorm.py). ## Testing Done - Hardware Type: NVIDIA H100 and NVIDIA B200. - Existing selected suites: **233 passed, 0 failed, 0 skipped on each GPU, before and after**; 125 deselected. Covers RMSNorm, LayerNorm, FusedAddRMSNorm and CUDA Graph replay. - No new pytest cases or changes to existing test tolerances. The existing BF16 Gemma RMSNorm test at M=128, N=4096 passes on both revisions. Its maximum absolute error against Triton improves identically on both GPUs: | Output | Before | After | |---|---:|---:| | Y | 0.0625 | 0 | | dX | 0.0625 | 0.0078125 | | dW | 0 | 0 | Cache-cleared benchmarks show lower forward latency across the seven measured BF16 RMSNorm shapes: **3.8–31.4% on H100 and 13.9–42.1% on B200**. Removing the BF16 weight vector saves **8/16/36 KiB** of forward/full peak allocation at N=4096/8192/18432. Full-call timings have more variation; details are below. - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Full-repository and model-convergence suites were not run. No end-to-end training speed, memory or convergence claim is made. <details> <summary>Environment, test command and original pytest output</summary> Python 3.12.10, PyTorch 2.9.1, CUDA 12.8, Triton 3.5.1, CuTe DSL 4.6.0, TVM FFI 0.1.13.post3. CUDA-visible memory: H100 79.2 GiB; B200 178.4 GiB. Run on the base commit `95b01e9` and this branch: ```bash LIGER_KERNEL_IMPL=cutedsl python -m pytest \ test/cutedsl/test_rms_norm.py test/cutedsl/test_rms_norm_stream.py \ test/ops/test_rms_norm.py test/ops/test_layer_norm.py \ test/ops/test_fused_add_rms_norm.py \ -k "nvidia-cutedsl or test_rms_norm_parity or test_rms_norm_mixed_weight_dtype or graph" -q ``` Original pytest summary lines, labelled by GPU and revision: ```text H100 — before ========= 233 passed, 125 deselected, 2 warnings in 116.55s (0:01:56) ========== H100 — after ========== 233 passed, 125 deselected, 2 warnings in 90.10s (0:01:30) ========== B200 — before ========= 233 passed, 125 deselected, 2 warnings in 153.33s (0:02:33) ========== B200 — after ========= 233 passed, 125 deselected, 2 warnings in 102.62s (0:01:42) ========== ``` The error table uses `test_rms_norm_parity[False-True-1.0-gemma-dtype1-2-64-4096]`. Some irregular parity shapes exercise the existing inline fallback. Supplemental diagnostics on each GPU passed 122 numerical cases against a CuTe FP32-pre-add control and 26 controls for non-Gemma/zero-offset behavior, offset cache keys and CUDA Graph replay with changed X, W and dY. Coverage includes BF16/FP16/FP32, mixed FP32 weights, strided inputs, in-place/out-of-place backward, residual-output gradients and hidden sizes 512–18432, including boundaries around 4096, 8192 and 16384. These diagnostics are separate from the committed pytest suite. </details> <details> <summary>PyTorch accuracy comparison and minimal reproduction</summary> Relative L2 error against PyTorch/autograd, using BF16 X/W, M=128, N=4096, offset=1, seed=42 and weight scale=0.01: | GPU | Op | Y before → after | dX before → after | |---|---|---:|---:| | H100 | RMSNorm | 0.00292 → 2.3e-05 | 0.00288 → 1.22e-05 | | H100 | FusedAddRMSNorm | 0.00289 → 2.53e-05 | 0.00191 → 0.000253 | | B200 | RMSNorm | 0.00292 → 2.3e-05 | 0.00288 → 1.21e-05 | | B200 | FusedAddRMSNorm | 0.00289 → 2.53e-05 | 0.00191 → 0.000256 | FusedAddRMSNorm includes a residual-output gradient. Its backward uses a rounded saved residual, which contributes to the remaining difference from PyTorch. dW is unchanged for fixed inputs and upstream gradients. Run this unchanged on each revision using the existing Gemma PyTorch reference: ```python import torch from liger_kernel.ops.backends._cutedsl.rms_norm import _LigerRMSNormCuTeDSLFunction from liger_kernel.testing._op_helpers import pytorch_reference_rms_norm torch.manual_seed(42) x = torch.randn(128, 4096, device="cuda", dtype=torch.bfloat16, requires_grad=True) w = (torch.randn(4096, device="cuda") * 0.01).bfloat16().requires_grad_() dy = torch.randn_like(x) xr = x.detach().clone().requires_grad_() wr = w.detach().clone().requires_grad_() y = _LigerRMSNormCuTeDSLFunction.apply(x, w, 1e-6, 1.0, "gemma", False, None) yr = pytorch_reference_rms_norm(xr, wr, 1e-6, 1.0, "gemma") actual = (y, *torch.autograd.grad(y, (x, w), dy)) expected = (yr, *torch.autograd.grad(yr, (xr, wr), dy.clone())) for name, a, b in zip(("Y", "dX", "dW"), actual, expected): delta = a.detach().double() - b.detach().double() rel_l2 = delta.norm() / b.detach().double().norm().clamp_min(1e-30) print(name, "max_abs=", delta.abs().max().item(), "relative_l2=", rel_l2.item()) ``` </details> <details> <summary>Benchmark results and measurement details</summary> Cache-cleared `triton.testing.do_bench`, warmup=5 ms, rep=30 ms, median of nine shuffled-provider rounds. Compilation and input/dY generation are excluded. `in_place=False`; full includes `autograd.grad`. An identical baseline is interleaved as an A/A control. GPU clocks were not locked; full timings include host enqueue gaps. Measured 16 configurations × forward/backward/full = 48 combinations per GPU, including seven BF16 RMSNorm shapes, three FusedAddRMSNorm shapes, other dtypes and zero-offset controls. Profiling confirms removal of one weight-add kernel from each affected forward. BF16 results, before → after (µs): | GPU | Op | M × N | Forward | Forward + backward | |---|---|---|---:|---:| | H100 | RMSNorm | 1 × 4096 | 8.864 → 6.080 | 102.640 → 95.184 | | H100 | RMSNorm | 128 × 4096 | 9.728 → 6.880 | 99.856 → 87.680 | | H100 | RMSNorm | 2048 × 4096 | 20.864 → 18.944 | 103.072 → 98.160 | | H100 | RMSNorm | 8192 × 4096 | 54.048 → 51.968 | 141.472 → 139.488 | | H100 | RMSNorm | 128 × 8192 | 10.592 → 8.000 | 103.744 → 92.160 | | H100 | RMSNorm | 2048 × 8192 | 31.808 → 29.504 | 107.216 → 94.672 | | H100 | RMSNorm | 128 × 18432 | 12.448 → 10.016 | 102.352 → 92.320 | | H100 | FusedAddRMSNorm | 128 × 4096 | 10.528 → 7.872 | 130.656 → 120.576 | | H100 | FusedAddRMSNorm | 2048 × 4096 | 30.800 → 28.224 | 130.672 → 119.888 | | H100 | FusedAddRMSNorm | 2048 × 8192 | 53.440 → 50.720 | 128.032 → 116.304 | | B200 | RMSNorm | 1 × 4096 | 11.104 → 6.432 | 61.120 → 55.680 | | B200 | RMSNorm | 128 × 4096 | 11.296 → 8.192 | 63.648 → 58.624 | | B200 | RMSNorm | 2048 × 4096 | 17.280 → 14.336 | 71.104 → 64.448 | | B200 | RMSNorm | 8192 × 4096 | 33.664 → 28.992 | 97.440 → 95.232 | | B200 | RMSNorm | 128 × 8192 | 12.288 → 8.192 | 71.008 → 62.528 | | B200 | RMSNorm | 2048 × 8192 | 23.360 → 18.496 | 66.256 → 64.528 | | B200 | RMSNorm | 128 × 18432 | 13.216 → 10.048 | 59.968 → 61.008 | | B200 | FusedAddRMSNorm | 128 × 4096 | 13.152 → 8.192 | 71.760 → 64.432 | | B200 | FusedAddRMSNorm | 2048 × 4096 | 22.496 → 18.464 | 85.216 → 77.312 | | B200 | FusedAddRMSNorm | 2048 × 8192 | 33.568 → 28.704 | 92.096 → 79.008 | H100 RMSNorm backward-only changes range from −0.8% to +4.7%; B200 backward-only medians are unchanged. Backward incremental allocation and full-call live allocation are unchanged. The B200 128×18432 full result has substantial A/A variation, and its apparent slowdown changed sign on repeat. A separate 21-round zero-offset repeat measured 53.760 → 56.096 µs (+4.3%), while the identical-baseline A/A control differed by +5.2%; forward was unchanged and backward differed by −0.15%. These measurements do not resolve a stable kernel regression or establish zero overhead. </details> Co-authored-by: Arup De <[email protected]> |
||
|
luca huangandArup De
|
80fb156fce |
fix(monkey_patch): patch Q/K norms on existing Qwen models (#1470)
## Summary Patch `self_attn.q_norm` and `k_norm` when applying Liger to existing Qwen3, Qwen3 MoE, Qwen3 Next, Qwen3.5 and Qwen3.5 MoE instances. Reuse each model's RMSNorm helper and respect `rms_norm`. Extend the existing instance tests following the Qwen3 VL/VL MoE and Gemma3 pattern: construct native HF instances first, then check Q/K norm forward bindings before and after patching. Also verify preserved weights and epsilon, casting settings, and unchanged Linear Attention norms. ## Testing Done - Hardware Type: H100 80GB - Instance tests: 21 passed. All 13 new Q/K regression assertions fail before the fix and pass after it. - RMSNorm numerical tests: 64 passed. - Convergence: 16 passed, 1 known baseline failure, 5 existing skips. The BF16 Qwen3 MoE with-logits failure also occurs on the base commit. These results are from prior validation; tests were not rerun after switching to explicit model classes. PyTorch 2.9.1 / CUDA 12.8, Triton 3.5.1, Transformers 5.15.1. Hybrid convergence uses FLA 0.5.2 and TileLang 0.1.14. Before uses the existing Liger instance patch with native Q/K norms; After adds this fix. Both use matched random weights and inputs, BF16, SDPA, identical other kernel settings, and no `torch.compile`. Each runs in a separate process with warmup and five measurement rounds. Speedup is Before / After. For these configurations, the patch covers all attention layers in Qwen3-8B (36/36) and Qwen3-30B-A3B (48/48). In Qwen3.5-9B and Qwen3.5-35B-A3B, it covers only the Full Attention layers (8/32 and 10/40); Linear Attention norms are unchanged. Q/K norms and a single Full Attention module on one H100 80GB, forward + backward median time in ms: | Model | Batch × seq | Q+K Before → After | Speedup | Attention Before → After | Speedup | | --- | ---: | ---: | ---: | ---: | ---: | | Qwen3-8B | 4 × 2048 | 2.439 → 0.562 | 4.34× | 8.883 → 6.705 | 1.32× | | Qwen3-8B | 1 × 8192 | 2.435 → 0.542 | 4.50× | 13.552 → 11.269 | 1.20× | | Qwen3.5-9B | 4 × 2048 | 2.560 → 0.739 | 3.46× | 10.763 → 8.802 | 1.22× | | Qwen3.5-9B | 1 × 8192 | 2.564 → 0.729 | 3.52× | 16.063 → 14.108 | 1.14× | | Qwen3-30B-A3B | 4 × 2048 | 2.219 → 0.643 | 3.45× | 7.071 → 5.280 | 1.34× | | Qwen3-30B-A3B | 1 × 8192 | 2.212 → 0.659 | 3.36× | 11.597 → 9.858 | 1.18× | | Qwen3.5-35B-A3B | 4 × 2048 | 2.292 → 0.740 | 3.10× | 8.248 → 6.324 | 1.30× | | Qwen3.5-35B-A3B | 1 × 8192 | 2.288 → 0.691 | 3.31× | 13.475 → 11.539 | 1.17× | - [ ] run `make test` (targeted tests above; full suite not run) - [x] run `make checkstyle` (prior validation) - [ ] run `make test-convergence` (targeted tests above; full suite not run) Co-authored-by: Arup De <[email protected]> |
||
|
|
ae47958e39 |
fix: respect gradient column strides in fused linear CE (#1500)
## Summary Fix incorrect gradient scaling in fused linear cross entropy when its cached input or weight gradients have a non-unit column stride. `zeros_like` can preserve a transposed tensor's layout, but `element_mul_kernel` previously addressed every column as contiguous. For a `(3, 4)` weight gradient with stride `(1, 3)`, its three programs only cover storage offsets 0 through 5, with overlapping accesses, leaving half of the gradient unscaled. Consequently, even `(0.0 * loss).backward()` can return nonzero gradients. Pass the actual column strides from FLCE backward and use them in the kernel's address calculation, with 64-bit offsets. The new constexpr argument defaults to 1 for existing callers. This preserves the gradient layout and avoids allocating a contiguous copy. The change is two files, +7/-2; no test files are included. Minimal CUDA example: ```python import torch from liger_kernel.ops.fused_linear_cross_entropy import LigerFusedLinearCrossEntropyFunction x = torch.ones(2, 4, device="cuda", requires_grad=True) w = torch.zeros(4, 3, device="cuda").t().requires_grad_() target = torch.tensor([0, 0], device="cuda") loss = LigerFusedLinearCrossEntropyFunction.apply(x, w, target)[0] (0.0 * loss).backward() print(w.grad) ``` Before: six of twelve weight-gradient elements remain nonzero, with maximum magnitude 0.66666663. After: all twelve are zero. This affects the cached-gradient `mean`/`sum` backward paths when scaling is required; the unit upstream gradient bypasses the kernel. Open-PR searches found no matching fix. Related #1292 changes device/host synchronization, and #1338 addresses compiled loss-buffer handling; neither handles column strides in gradient scaling. ## Testing Done - Hardware: NVIDIA RTX A1000 Laptop GPU (4 GiB), PyTorch 2.13.0+cu130, Triton 3.7.1. - Existing FLCE suite: **141 passed**. - External end-to-end CUDA autograd harness: **360/360 comparisons passed** against PyTorch. Covered contiguous/transposed inputs and weights, FP32/FP16/BF16, default/FP32 accumulation, `mean`/`sum`/`none`, and upstream coefficients 0, 1, 0.5, 2, and -1. The baseline failed 144 of these comparisons. The harness is outside the repository. - Changed-file Ruff lint/format and `git diff --check`: passed. ```text python -m pytest test/transformers/test_fused_linear_cross_entropy.py -q --disable-warnings --maxfail=1 -o log_cli=false 141 passed, 15 warnings in 3.60s pre-commit run --files src/liger_kernel/ops/fused_linear_cross_entropy.py src/liger_kernel/ops/utils.py ruff check...............................................................Passed ruff format..............................................................Passed ``` Small performance check on the final implementation, using the repository's benchmark helpers: BF16, input `(256, 512)`, weight `(2048, 512)`, three measurements per case. Values below are forward+backward medians; these are single-GPU microbenchmarks, not model-throughput results. | Layout | Before (ms) | After (ms) | Additional peak allocation, before/after | | --- | ---: | ---: | ---: | | Contiguous | 0.898 | 0.885 | 2.50 / 2.50 MiB | | Both transposed | 0.761 | 0.991 | 2.50 / 2.50 MiB | No speedup is claimed. Correct strided accesses are slower in the transposed case; the baseline skips much of the logical gradient while repeatedly touching overlapping addresses. Full model convergence, compilation, and other device backends were not evaluated. - [x] Run the targeted FLCE suite and independent CUDA regression checks. - [x] Run Ruff lint and formatting checks on both changed files. - [ ] Run full `make test`. - [ ] Run full `make checkstyle`. - [ ] Run `make test-convergence`. AI assistance: OpenAI Codex assisted with implementation and local validation. Signed-off-by: taking-lying-flat <[email protected]> Co-authored-by: Codex <[email protected]> Co-authored-by: Arup De <[email protected]> |
||
|
|
395d6166e8 |
fix(kl_div, tvd): make torch.compile work (NameError, then fp64 loop-carried scalar) (#1509)
## 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]> |
||
|
LazuriteandVaibhav Jindal
|
5b0f210cc4 |
[Megatron] Rebind CE call sites so apply_liger_kernel_to_megatron(cross_entropy=True) takes effect (#1486)
## Summary Fixes #1485. `apply_liger_kernel_to_megatron(cross_entropy=True)` replaces the CE functions only on their defining modules, but Megatron calls copies bound at import time. When Megatron model code is imported before the patch (the usual order, also used by `examples/megatron/run_mode1_monkey_patch.py`), the loss silently stays on Megatron's kernel: - Released megatron-core (0.18 / 0.19): `language_module` binds `fused_vocab_parallel_cross_entropy` with `from ... import` ([L17](https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/models/common/language_module/language_module.py#L17)). The unfused branch calls the `megatron.core.tensor_parallel` re-export ([L183](https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/models/common/language_module/language_module.py#L183)), which is bound before the `cross_entropy` submodule can be patched. Neither path reaches Liger. - Megatron-LM main (`NVIDIA/Megatron-LM#6766`): `backends` binds the fused function as `_fused_ce` ([L19-L21](https://github.com/NVIDIA/Megatron-LM/blob/ec806b20a0a976924c1668486c3554ffb90c000e/megatron/core/models/backends.py#L19-L21)). The fused path misses Liger; the unfused path already works. This PR rebinds those copies too, but only where they still hold the original function. No version check is needed, and a call site that another framework has already overridden is left alone. If no fused call site ends up on Liger, the patch logs a warning. The docs now say to call the patch before building the model, matching the docstring: Megatron-LM main picks the CE function when the model is constructed. Once the minimum supported megatron-core includes the `BackendSpecProvider.vocab_parallel_cross_entropy` slot, this can switch to overriding that slot on the providers and drop the rebinding for released versions. ## Testing Done - Hardware Type: 1x NVIDIA GPU - [ ] run `make test` to ensure correctness: ran `test/megatron` (173 passed, 48 skipped multi-GPU tests) - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence: N/A, no kernel change The new tests in `test_monkey_patch.py` compute the loss through stub call sites laid out like each Megatron version; 4 of 5 fail on current `main`. A real 1-layer `GPTModel` (TP=1) was run for one forward and backward pass, with model code imported before the patch. Calls into Liger CE (fused / unfused): | | before | after | |---|---|---| | megatron-core 0.18.2 | 0 / 0 | 1 / 1 | | Megatron-LM main | 0 / 1 | 1 / 1 | When the patch runs first, both cases also give 1 / 1 after this change (before: 1 / 0 on 0.18.2). Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
EdsonandVaibhav Jindal
|
488acbd9c0 |
Fix llama4 class-level SwiGLU patch for configs with dense layers (#1507)
## Summary
`apply_liger_kernel_to_llama4(swiglu=True)` (the default) sets
`modeling_llama4.Llama4TextMLP = LigerSwiGLUMLP`. Transformers builds
dense (non-MoE) Llama4 decoder layers as `Llama4TextMLP(config,
intermediate_size=config.intermediate_size_mlp)`, but
`LigerSwiGLUMLP.__init__` only accepted `config`. Once the class-level
patch was applied before model init (`apply_liger_kernel_to_llama4()`,
`_apply_liger_kernel("llama4")`, `AutoLigerKernelForCausalLM`), building
any Llama4 model with a dense layer raised:
```
TypeError: LigerSwiGLUMLP.__init__() got an unexpected keyword argument 'intermediate_size'
```
This affects Llama-4-Maverick (`interleave_moe_layer_step=2`) and any
config whose `moe_layers` does not cover every layer. The `mini_llama4`
convergence config uses the default `interleave_moe_layer_step=1`, so
every layer there is MoE and the path was never exercised. #907 fixed
the instance-level path only.
Changes:
- `swiglu.py`: `LigerSwiGLUMLP.__init__(self, config,
intermediate_size=None)` now falls back to `config.intermediate_size`,
which is the same pattern as `LigerQwen3MoeSwiGLUMLP` and HF's
`Llama4TextMLP`. Existing callers that pass only `config` behave exactly
as before.
- `test_monkey_patch.py`: adds
`test_apply_liger_kernel_to_llama4_class_patch_with_dense_layers`. The
test applies the class-level patch (swiglu only), builds a 2-layer
`Llama4ForCausalLM` with `moe_layers=[1]`, and checks two things:
- the dense layer's MLP is a `LigerSwiGLUMLP` of width
`intermediate_size_mlp`;
- the MoE `shared_expert` keeps width `intermediate_size`.
It restores `modeling_llama4.Llama4TextMLP` via `monkeypatch`.
## Testing Done
The new test fails without the fix and passes with it:
```
# main @
|
||
|
EdsonandVaibhav Jindal
|
77624a9aa7 |
Fix qwen2_vl / qwen2_5_vl class-level RMSNorm patch on transformers >= 5.2 (#1511)
## Summary
Calling `apply_liger_kernel_to_qwen2_vl(rms_norm=True)` or
`apply_liger_kernel_to_qwen2_5_vl(rms_norm=True)` without a model
instance replaces `modeling_<model>.Qwen2RMSNorm`. Since transformers
5.2.0, `modeling_qwen2_vl` and `modeling_qwen2_5_vl` define and use
their own `Qwen2VLRMSNorm` / `Qwen2_5_VLRMSNorm`, so that assignment
only creates an unused attribute. Every RMSNorm in a model built after
the patch stays the HF module, with no warning. This is the path
`AutoLigerKernelForCausalLM` / `_apply_liger_kernel` and the multimodal
convergence tests use. The instance path is not affected.
| transformers | qwen2_vl uses | qwen2_5_vl uses |
|---|---|---|
| <= 5.1 | `Qwen2RMSNorm` | `Qwen2RMSNorm` |
| >= 5.2 | `Qwen2VLRMSNorm` | `Qwen2_5_VLRMSNorm` |
Changes:
- `monkey_patch.py`: also assign `modeling_qwen2_vl.Qwen2VLRMSNorm` and
`modeling_qwen2_5_vl.Qwen2_5_VLRMSNorm` to `LigerRMSNorm`. The old
`Qwen2RMSNorm` assignment stays for transformers < 5.2. I used both
names instead of an `IS_TRANSFORMERS_V5_OR_LATER` branch because 5.0 and
5.1 still use `Qwen2RMSNorm`. The constructor signature `(hidden_size,
eps)` matches.
- `test_monkey_patch.py`: adds
`test_apply_liger_kernel_to_qwen2_vl_family_class_patch_rms_norm`,
parametrized over qwen2_vl and qwen2_5_vl. It applies the class-level
patch with only `rms_norm=True`, builds the text model from a tiny
config, and checks that `norm`, `input_layernorm` and
`post_attention_layernorm` are `LigerRMSNorm`. The module attributes are
restored via `monkeypatch`.
On 5.2+ the Qwen2.5-VL vision-block norms become `LigerRMSNorm` too,
since they share `Qwen2_5_VLRMSNorm`. That is the same as the behaviour
before 5.2, where `Qwen2RMSNorm` covered both.
Not in this PR: on recent transformers (5.18) neither module defines
`apply_multimodal_rotary_pos_emb` any more, so `rope=True` on the
class-level path is also a no-op there. That one is not a rename,
because mrope moved into the rotary embedding, so it should be handled
separately.
## Testing Done
The new test fails without the fix and passes with it:
```
# main @
|
||
|
Stefan WangandVaibhav Jindal
|
8e698102ae |
fix(glm4v): version-aware hidden_size/vocab_size for Glm4vConfig (#1214)
Reading `config.hidden_size` and `config.vocab_size` directly in `glm4v.py` and `glm4v_moe.py` fails on transformers v5+, where `Glm4vConfig` and `Glm4vMoeConfig` are composite (`text_config` + `vision_config`) and the top-level fields are no longer present. Reported in #1152: ``` File ".../liger_kernel/transformers/model/glm4v.py", line 130, in lce_forward hidden_size=self.config.hidden_size, AttributeError: 'Glm4vConfig' object has no attribute 'hidden_size' ``` Same shape PR #1062 already fixed for Qwen2-VL / Qwen2.5-VL — mirror that helper (`_get_hidden_size`, `_get_vocab_size`) gated on `transformers >= 5.0.0`, applied to both `glm4v.py` and `glm4v_moe.py`. Pre-v5 still works because the helpers fall back to the top-level attribute. Helper signatures are typed (`config: PretrainedConfig -> int`) and the version flag is annotated `bool`. ## Reproducer (transformers 5.7.0) ```python from transformers import Glm4vConfig cfg = Glm4vConfig(text_config={"hidden_size": 1024, "vocab_size": 32000}) cfg.text_config.hidden_size # 1024 cfg.hidden_size # AttributeError: 'Glm4vConfig' object has no attribute 'hidden_size' ``` Same shape reproduces for `Glm4vMoeConfig`. Identical to the trace in #1152. After the helper change, both `_get_hidden_size(cfg)` and `_get_vocab_size(cfg)` route through `text_config` on v5+ and fall back to the top-level attribute pre-v5: ``` transformers=5.7.0 v5_or_later=True _get_hidden_size(Glm4vConfig) = 1024 _get_vocab_size(Glm4vConfig) = 32000 _get_hidden_size(Glm4vMoeConfig) = 1024 _get_vocab_size(Glm4vMoeConfig) = 32000 ``` ## Tests End-to-end coverage already in place in `test/convergence/bf16/test_mini_models.py` and `test/convergence/fp32/test_mini_models.py` — both `mini_glm4v` and `mini_glm4v_moe` configs construct via the `text_config` sub-dict (lines 1245–1308 / 1311–1368), so CI exercises the fixed path on v5+ runners. No new test file added — matches the precedent set by PR #1062 (qwen2_vl fix), which also relied on the existing convergence tests rather than introducing a helper-level test layer. Closes #1152. Signed-off-by: 1fanwang <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
cherry77-cloudandVaibhav Jindal
|
e3e6daa889 |
Fix sequence-level vLLM correction in VESPO loss (#1493)
## Summary With `loss_type="vespo"` and `importance_sampling_level="token"`, passing a sequence-level `vllm_is_ratio` of shape `[B, 1]` to the Triton loss currently expands it to `[B, T]` before `get_gamma_weights` sums its logarithm. This applies a sequence correction `R` as `R**T`, potentially suppressing the policy-gradient term. Preserve `[B, 1]` corrections and only unsqueeze `[B]` inputs. The shared helper already handles both sequence-level and token-level corrections, as used by the chunked loss. For 64 valid tokens, identical current/old policies, positive advantage, and a sequence correction of 1.05, the reproduced policy loss changes from `4.3685e-26` to `1.5272448`, matching the chunked path. The logits-gradient norm changes from `3.0347e-27` to `0.10609366`. This fixes direct Liger API/custom integrations supplying sequence corrections. Current [TRL validation](https://github.com/huggingface/trl/blob/cc5bb3f8c94b7f9cf5c24029d7812e78849ccace/trl/trainer/grpo_trainer.py#L925-L931) only permits token-level vLLM correction modes with VESPO, so supported native TRL configurations do not currently reach this case. ## Testing Done - Hardware Type: NVIDIA RTX A1000 Laptop GPU; PyTorch 2.13.0+cu130, Triton 3.7.1. - Passed 11 local FP32 CUDA forward/backward comparisons of the actual fixed Triton entry point against the chunked loss (`compiled=False`): sequence lengths 1/4/64, `[B, 1]` and `[B]` corrections, `[B, T]` corrections, unit correction, padding, mixed positive/negative advantages, KL, and omitted old log probabilities. All losses and gradients agree within FP32 tolerance. - Passed the repository's Ruff check and format hooks via `pre-commit run --files src/liger_kernel/transformers/grpo_loss.py`, plus `git diff --check`. - This PR contains only the production-code change; the regression harness was run locally. - [ ] `make test` (full suite not run) - [ ] `make checkstyle` (targeted repository hooks passed) - [ ] `make test-convergence` (not run) Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Yulong CaoandVaibhav Jindal
|
8343b3cf7b |
[Feature] Support apply_liger_kernel_to_deepseek_v2 (#1501)
## Summary Adds Liger Kernel support for DeepSeek-V2 (`deepseek_v2` model type). Closes #129 (on the maintainer 2024 Q4 roadmap) and supersedes the stale PR #499, which predates the `deepseek_v2` module in transformers. - `model/deepseek_v2.py`: new `lce_forward`. In training with labels, `skip_logits` defaults to `True`, so logits are never materialized (FLCE path via `LigerForCausalLMLoss`); otherwise logits are materialized and the HF-native `self.loss_function` is used. Returns `LigerCausalLMOutputWithPast` with `token_accuracy` / `predicted_tokens`. - `monkey_patch.py`: `apply_liger_kernel_to_deepseek_v2` with class-level and instance-level patching, registered in `MODEL_TYPE_TO_APPLY_LIGER_FN` (`AutoLigerKernelForCausalLM` picks it up automatically; previously `deepseek_v2` models silently got no Liger kernels). - RMSNorm: `DeepseekV2RMSNorm` → `LigerRMSNorm`, including the MLA attention `q_a_layernorm` / `kv_a_layernorm`. - SwiGLU: dense MLPs and shared experts (`DeepseekV2MLP`) → `LigerQwen3MoeSwiGLUMLP` on all supported versions. Routed experts: batched `DeepseekV2Experts` → fused `LigerExperts` on transformers v5+; per-expert patch on transformers v4 (the `deepseek_v2` module landed in v4.54). - RoPE: unsupported — DeepSeek-V2 uses interleaved partial RoPE, incompatible with `liger_rotary_pos_emb`. `rope=True` emits a `warning_once` and skips (no silent behavior change); RoPE is intentionally not listed in the README table. - `mlp_bias`: `LigerQwen3MoeSwiGLUMLP` is bias-free, matching the `DeepseekV2Config` default (`mlp_bias=False`); `mlp_bias=True` checkpoints are not supported and noted in the docstring. - Tests: instance monkey-patch test (asserts both v5 and v4 branch outcomes), 4 convergence tests (bf16/fp32 × FLCE/logits) driven by a mini `deepseek_v2` config, `revert_liger_kernel_to_deepseek_v2` util, README entry. ## Details - transformers v5 renames the V2 MoE block `DeepseekV2MoE` → `DeepseekV2Moe` (V3 kept `DeepseekV3MoE`), so the apply function resolves the class with a try/except import. - Shared-code change: `LigerExperts`' expert-count lookup previously handled `num_experts` (qwen3_moe/qwen3_next) and `num_local_experts` (mixtral/deepseek_v3). `DeepseekV2Config` on early transformers v5 only exposes `n_routed_experts` (the `num_experts` alias was added later), so a third `elif` branch was added. Existing branches are untouched; regression coverage for both neighboring branches ran in the same validation (see below). - The transformers v4 branch (per-expert patching) mirrors the existing mixtral / qwen3_moe / deepseek_v3 patterns; it is exercised by CI's v4 matrix where the `deepseek_v2` module is available (≥ 4.54) and skipped via `is_deepseek_v2_available()` otherwise. - On transformers v5+, `LigerExperts`' batched expert weights (`(N, 2*I, H)` / `(N, H, I)`) match the HF `DeepseekV2Experts` layout, so pretrained checkpoints load unchanged. ## Testing Done | Check | Status | |---|---| | Instance patching (CPU, transformers 5.16.1) | PASS | | Convergence bf16 FLCE (Kaggle T4) | PASS | | Convergence bf16 with_logits (Kaggle T4) | PASS | | Convergence fp32 FLCE (Kaggle T4) | PASS | | Convergence fp32 with_logits (Kaggle T4) | PASS | | Regression: mini_deepseek_v3 + mini_qwen3_moe (Kaggle T4) | PASS (pre-rebase base; see Post-rebase re-verification below) | | Instance patching deepseek_v3 (Kaggle T4) | PASS | | deepseek_v4 sibling tests | SKIP (module not in transformers 5.0.0) | | Checkstyle (ruff) | PASS | Exact commands run on this change: - `pytest test/transformers/test_monkey_patch.py -k deepseek` → 3 passed (transformers 5.16.1, CPU) - `pytest test/transformers/test_monkey_patch.py test/convergence/{bf16,fp32}/test_mini_models.py test/convergence/{bf16,fp32}/test_mini_models_with_logits.py -k 'deepseek or qwen3_moe'` on Kaggle T4 (torch 2.10.0+cu128, transformers 5.0.0) → **13 passed, 7 skipped, 0 failed** ([kernel log](https://www.kaggle.com/code/cauliyulong/liger-dsv2-129-test)) ### Post-rebase re-verification The branch was rebased onto `main` at `6ad077c` (it was 17 commits behind). Everything below was re-run on the rebased commit: - `pytest test/transformers/test_monkey_patch.py -k deepseek` → **3 passed** (CPU, transformers 5.16.1) - `pytest test/convergence/{bf16,fp32}/test_mini_models{,_with_logits}.py -k deepseek_v2` → **4 passed** (RTX 4060, torch 2.11.0+cu128) - `... -k deepseek_v3` → **4 passed** - `... -k qwen3_moe` → 3 passed, 1 failed: `bf16/test_mini_models_with_logits.py::test_mini_model[mini_qwen3_moe-...]`, `[Loss]` mismatch at index `(0, 27)` (`5.145391464` vs `4.135959148`). **Not caused by this PR**: the identical mismatch (same index, same values) reproduces on a clean `origin/main` checkout at `6ad077c` in the same local environment (transformers 5.16.1 / torch 2.11.0), and this PR does not touch the qwen3_moe path — `LigerExperts` takes the `hasattr(config, "num_experts")` branch for qwen3_moe both before and after the change. It does not reproduce on the Kaggle T4 run above (transformers 5.0.0), so it looks environment-specific rather than a regression from this change. - `ruff check` on all changed files → all checks passed - Hardware Type: NVIDIA T4 (Kaggle, sm_75) + RTX 4060 (sm_89); CPU for the instance-patching tests - [x] run `make checkstyle` to ensure code style (`ruff check` clean on all changed files) - [ ] run `make test` to ensure correctness - [ ] run `make test-convergence` to ensure convergence Note on the two unchecked boxes: `make test` / `make test-convergence` were not run end-to-end on this change. The targeted subsets listed above are what was exercised; I can run the full targets on request. Happy to add a memory/throughput table on A100 for `DeepSeek-V2-Lite` dims if reviewers find it useful. Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
cherry77-cloud
|
7d88a77860 |
fix(grpo): correct sequence-level KL and masked loss gradients (#1491)
## Summary
Fix two incorrect gradients in the sequence-level GRPO backward path
while leaving the forward loss unchanged. This PR changes only
`src/liger_kernel/ops/grpo_loss.py` (+10/-4).
- With `importance_sampling_level="sequence"`,
`use_bias_correction_kl=True`, and nonzero `beta`, every token
influences the KL loss through the shared sequence importance ratio. The
previous backward used the token-level shortcut and omitted the
cross-token derivative, which can reverse a parameter's gradient
direction. Save the already-computed importance-corrected KL and reduce
it with the masked upstream gradients to recover that term.
- With `reduce=False`, padding outputs are zero, but operations such as
`loss.sum()` still produce nonzero upstream gradients at those
positions. Mask upstream gradients before the sequence reductions so
padding cannot amplify the gradients of valid tokens.
For valid token `j`, let `g` include the upstream gradient, mask, and
loss normalization; `r` is the unclamped sequence importance ratio; `k_t
= exp(ref_t - logp_t) - (ref_t - logp_t) - 1`; and `n` is the
valid-token count. The corrected KL contribution is:
```text
dKL/dlogp_j = beta * (r * (1 - exp(ref_j - logp_j)) * g_j
+ sum_t(g_t * r * k_t) / n)
```
The KL reduction uses the pre-delta-clamp ratio and is separate from
vLLM's policy-loss correction. The token-level implementation is
unchanged.
Two deterministic reproductions:
| Case | Before | After / independent reference |
| --- | --- | --- |
| Two valid tokens, selected probabilities 0.1, old=current, ref-logp
gaps [0.1, 1.0], beta=1, advantage=0: first token's log-probability
gradient | -0.05000000 | +0.12827773 |
| One valid token, selected probability 0.2, old=current, advantage=1,
beta=0, padded to 4 positions: valid token's logits gradient | [-3.2,
+3.2] | [-0.8, +0.8] |
## Testing Done
- Hardware Type: NVIDIA RTX A1000 Laptop GPU (4 GB)
- Environment: PyTorch 2.13.0+cu130, Triton 3.7.1
- [ ] `make test` — the entire repository suite was not run; targeted
GRPO results are below.
- [x] `make checkstyle` — passed with the repository-pinned Ruff
0.14.11; `git diff --check` also passed.
- [ ] `make test-convergence` — not run.
Validation used local regression tests and standalone reference scripts;
these are not included in this production-only PR.
```text
Targeted regressions before: 41 failed, 5 passed
Targeted regressions after: 46 passed
Original-size GRPO file: 136 passed, 17 skipped, 12 CUDA OOM
Memory-adjusted GRPO file: 148 passed, 17 skipped
Independent autograd matrix: 864 passed
Token gradients vs old code: 432 bitwise-equal comparisons
Forward outputs vs old code: 540 bitwise-equal comparisons
Chunked autograd comparison: 392 passed
Linear-head FP64 gradcheck: passed
```
The 12 OOM cases are existing large-vocabulary GRPO/CISPO/SAPO tests
whose PyTorch references exceed this GPU's memory. For those cases only,
a local pytest plugin reduced sequence length from 1024 to 256,
retaining vocabulary size 151936, dtype, and algorithm parameters. The
original sizes remain unvalidated on this GPU. The 17 skips follow
existing test conditions.
The 864-case matrix covers FP32/BF16/FP16, token/sequence importance
sampling, absent/full/left/right/interior/all-zero masks,
reduced/unreduced loss, KL correction, delta clipping, `old_logp=None`,
nonuniform and noncontiguous upstream gradients, vLLM ratio shapes, and
inplace/out-of-place backward. Supported loss combinations follow the
existing tests. Cross-checks against
`LigerFusedLinearGRPOFunction.ppo_loss_fn` exclude LUSPO because its
existing padding-forward semantics differ from the ops path. Maximum
FP32 absolute gradient error was 4.30e-6. The reported linear-head
gradient changes from -0.045 to approximately +0.11545, agreeing with
independent autograd and finite differences.
A small local BF16 benchmark (B=2, V=8192, sequence importance sampling,
corrected KL, with padding) measured forward+backward latency of
0.160→0.165 ms at T=128 and 0.375→0.380 ms at T=512. The correction
retains an O(B*T) FP32 KL tensor and adds the upstream-weighted
reduction only when sequence-level corrected KL is enabled.
|
||
|
Alexandre KohlerandVaibhav Jindal
|
d2965de5de |
fix(ppo): warn instead of raising when ref_per_token_logps and ref_input are both passed (#1506)
## Summary
Root cause: in `LigerFusedLinearPPOBase.forward`, the case where both
`ref_per_token_logps` and `ref_input` are passed is handled with `raise
Warning("Both ref_per_token_logps and ref_input are provided. Using
ref_per_token_logps.")`. `Warning` is an exception class, so this line
raises and the GRPO loss call aborts. The rest of the function already
does what the message says: `_ref_input_chunks` is only built when
`ref_per_token_logps is None`, and the reference forward only runs when
`ref_per_token_logps_chunk is None`. So the fallback exists but is never
reached.
Fix: use `warnings.warn(...)` so the call continues and uses
`ref_per_token_logps`, as the message describes.
## Testing Done
Added `test_ref_per_token_logps_preferred_over_ref_input` in
`test/chunked_loss/test_grpo_loss.py`. It passes both
`ref_per_token_logps` and `ref_input`/`ref_weight`, expects a
`UserWarning`, and checks that the loss equals the loss computed with
`ref_per_token_logps` alone.
- Before the fix it fails with `Warning: Both ref_per_token_logps and
ref_input are provided. Using ref_per_token_logps.`
- After the fix it passes, together with the non-compiled
`test_grpo_frozen_lm_head` cases (9 passed).
- Hardware Type: NVIDIA GeForce GTX 1060 6GB, PyTorch 2.5.1 (eager only;
torch.compile needs a newer GPU)
- [x] `make checkstyle` checks: `ruff check .` and `ruff format --check
.` pass (ruff 0.14.11, as pinned in `.pre-commit-config.yaml`)
- [ ] `make test` (not run; Triton is not available on this Windows
host)
- [ ] `make test-convergence` (not run; this change only affects
argument validation)
Co-authored-by: Vaibhav Jindal <[email protected]>
|
||
|
Piyush UmateandJari Kolehmainen
|
b297821787 |
feat(ops): support partial RoPE in Triton kernel (#1451)
## Summary This PR adds support for **Partial Rotary Position Embedding (partial RoPE)** to Liger-Kernel's Triton kernel and enables it for Hugging Face models using partial RoPE (such as Qwen3-Next). ### Context & Motivation - **Related HF Issue**: [huggingface/transformers#48622](https://github.com/huggingface/transformers/issues/48622) ("*Kernels for partial RoPE*") Several modern architectures rotate only a subset of each head's dimensions (`rotary_dim < head_dim`) and leave the remaining dimensions unchanged: - **Qwen3-Next / Qwen3.5**: 25% partial RoPE (`rotary_dim = 32`, `head_dim = 128`) - **GLM-4 / GLM-4v**: 50% partial RoPE - **GPT-NeoX**: 25% partial RoPE (`rotary_dim = 32`, `head_dim = 128`) - **Phi-2 / Phi-1.5**: 25% partial RoPE (`rotary_dim = 20`, `head_dim = 80`) - **StableLM**: 25% partial RoPE (`rotary_dim = 16`, `head_dim = 64`) Previously, Liger's Triton RoPE kernel assumed `rotary_dim == head_dim` and hardcoded `hd // 2` offsets. When used with partial RoPE models, it either threw `NotImplementedError("liger_rotary_pos_emb is not available for Qwen3Next models.")` or accessed out-of-bounds on `cos`/`sin`. ## Details 1. **`src/liger_kernel/ops/rope.py`**: - Updated `_triton_rope` to accept `pad_rd: tl.constexpr` and `rd: tl.constexpr`. - Offsets for loading/storing only address `[0 : rd//2]` and `[rd//2 : rd]`. - The pass-through slice `[rd : hd]` is untouched in-place. - Normalized `cos` and `sin` inputs across 2D `(seq_len, dim)`, 3D `(bsz, seq_len, dim)`, and 4D `(bsz, 1, seq_len, dim)`. - In `rope_forward` and `rope_backward`, automatically infer `rotary_dim = min(cos.shape[-1], head_dim)`. - Preserves 100% backward compatibility with full RoPE (`rotary_dim == head_dim`). 2. **`src/liger_kernel/ops/backends/_cutedsl/rope.py`**: - Added automatic fallback to Triton when `cos.shape[-1] < q.shape[-1]` (partial RoPE), preserving existing CuTe DSL routing for full RoPE. 3. **`src/liger_kernel/transformers/monkey_patch.py`**: - Enabled `modeling_qwen3_next.apply_rotary_pos_emb = liger_rotary_pos_emb`. 4. **Unit Tests**: - Added `test_partial_rope_correctness` in `test/ops/test_rope.py` testing across FP16, BF16, FP32, and representative partial RoPE ratios (Phi-2 20/80, NeoX 32/128, StableLM 16/64, 50% 64/128). - Added `test_partial_rope_transformers_correctness` in `test/transformers/test_rope.py` verifying both forward and backward pass gradients against Hugging Face's reference implementation. ## Testing Done - Hardware Types Tested: **NVIDIA H100 80GB HBM3** (sm_90) and **NVIDIA A10G** - [x] `make test`: `pytest test/ops/test_rope.py test/transformers/test_rope.py` passed 100% on both H100 and A10G (96 passed, 2 skipped) - [x] `make checkstyle`: `ruff check` and `ruff format --check` (0 errors) - [x] End-to-end model forward & backward loss and gradient parity check: - Qwen3-Next / NeoX (25%, D=128, R=32, bf16): `Loss Diff = 0.00e+00`, `Attn Grad Diff = 0.00e+00` - Phi-2 (25%, D=80, R=20, bf16): `Loss Diff = 0.00e+00`, `Attn Grad Diff = 0.00e+00` - 50% Partial RoPE (D=128, R=64, fp16): `Loss Diff = 0.00e+00`, `Attn Grad Diff = 2.86e-06` - [x] Multi-step training loop with AdamW: converged smoothly on H100 (loss 6.34 $\to$ 1.24) with zero NaNs/Infs. - [x] Autoregressive token generation: 100% exact token-for-token parity with reference. ### Benchmark on NVIDIA H100 80GB HBM3 (Forward + Backward, `bfloat16`) | Workload | Shape ($B, H, S, D, R$) | PyTorch Reference | Liger Triton RoPE | Speedup | |---|---|---|---|---| | **Qwen3-Next (S=2048, B=4)** | $4, 32, 2048, 128, 32$ | 1.50 ms | **0.56 ms** | **2.67x** | | **Qwen3-Next (S=4096, B=4)** | $4, 32, 4096, 128, 32$ | 2.91 ms | **1.08 ms** | **2.68x** | | **Qwen3-Next Large (S=2048, B=8)** | $8, 32, 2048, 128, 32$ | 2.91 ms | **1.08 ms** | **2.68x** | | **Phi-2 Style (S=2048, B=8)** | $8, 32, 2048, 80, 20$ | 2.04 ms | **0.78 ms** | **2.61x** | Co-authored-by: Jari Kolehmainen <[email protected]> |
||
|
|
8866e4ee79 |
Warn and skip native configuration on unsupported devices (#1502)
## Summary Follow up on #1499 so consumers can call Liger's public configuration API without duplicating GPU eligibility checks. This unblocks removing those checks from verl-project/verl#7945. - Keep LCK strict: `liger_cute_kernels.configure` raises the exported `UnsupportedDeviceError` when it cannot support the requested device, without loading or initializing native resources. Successful LCK setup retains its `None` return. - Catch that specific exception in **`liger_kernel.ops.configure.configure`**, emit a warning, and return `False`. Warning/fallback policy lives in Liger-Kernel, not LCK or downstream consumers. - Apply the same contract to FLSCE and MoE configuration. On unsupported NVIDIA architectures, retain collective plan validation before LCK raises the typed error. - Return `True` from the Liger-Kernel frontend after successful native setup. Older native packages without the new exception export remain usable; their genuine setup errors still propagate. - Keep actual setup failures visible: invalid capacities, mismatched collective plans, dependency failures, and conflicting NVSHMEM ownership are not converted into fallback results. Missing optional LCK already warns and returns `False`; this extends that behavior to unsupported hardware. Operator frontends continue to own implementation selection. Skipping setup does not make an explicitly selected native-only operator usable, and this PR does not change CUDA kernels or operator fallback algorithms. ## Validation - 143 focused tests passed on each of H200 and B300: 89 Liger configuration contract tests and 54 VeRL integration regressions. - Coverage includes supported Hopper/Blackwell setup, direct LCK exceptions without warnings, public Liger exception-to-warning conversion, both operator sections, no native initialization on skipped setup, collective hardware disagreement, invalid capacity rejection, legacy native return compatibility, and VeRL delegation without hardware probes. - Runtime, value, and OS errors remain unchanged through the public boundary with both old and new native package interfaces. - Unsupported architectures and non-NVIDIA devices are covered with host-side mocks, not physical hardware runs. - Repository Ruff checks and formatting passed. Compatible Python and native packages containing this change must be released before downstream users can rely on the new unsupported-device behavior. The broader VeRL performance campaign was not rerun. --------- Co-authored-by: kolehma8 <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
0043f43309 |
Add shared LCK configure API for FLSCE and MoE (#1499)
## Summary Add a public, runtime-wide `liger_cute_kernels.configure(...)` API for shared NVSHMEM initialization and preparation of **multiple TP and EP groups**. This supplies the upstream state-management prerequisite discussed in https://github.com/verl-project/verl/pull/7945 without making runtime ownership specific to the output head. **Draft: single-node B300 and H200 runtime coverage is recorded below. Multi-host execution and downstream VeRL integration remain pending.** Rebased onto upstream main at `6ad077c`; the already merged CUDA-extras change is not included again. ```python from liger_kernel.ops.configure import FusedLinearCrossEntropyConfig, MoEConfig, configure native_configured = configure( process_groups={"tp": pg_collection.tp, "ep": pg_collection.ep}, bootstrap_group=world_group, device=device, flsce=FusedLinearCrossEntropyConfig( max_tokens=8192, hidden_size=4096, local_vocab_size=32000, group="tp", ), moe=MoEConfig( max_tokens=8192, hidden_size=4096, num_experts=128, top_k=8, num_hosts=ep_num_hosts, gpus_per_host=ep_gpus_per_host, max_inflight=8, # Reserve for all simultaneously retained MoE forwards. group="ep", ), ) ``` ## Configuration and group ownership - One canonical runtime implementation lives in LCK. `liger_kernel.ops.configure` is a thin optional entry point with dependency-free configuration dataclasses. Importing it and constructing configuration data do not require LCK. When LCK is absent, `configure` warns and returns `False`; successful setup returns `True`, and errors from an installed LCK propagate. - Bootstrap defaults to torch WORLD, not the first operator's group. Named groups accept Megatron's `pg_collection` without a Megatron dependency. TP and EP may use different uniform strided partitions of the same bootstrap domain, including a non-WORLD bootstrap domain. - Both sections can be configured together, added later, or repeated to prepare additional partitions and group sizes within their original process-wide capacities. Additional names are additive; an existing name retains its ordered membership. Equal/smaller requests reuse reservations rather than clearing another consumer's state. - Setup collectively validates capacities, topology, membership, cache agreement, CUDA hardware, and runtime ownership. Partial cache hits still participate in required team splits and symmetric reservations. Equivalent ordered memberships reuse a context. - The Python execution path caches `ProcessGroup -> team_handle`; execution never creates a team. Native FLSCE and MoE contexts select the prepared group, and queued kernels/captured graphs retain their own pointers and topology. - Both SM90 and SM100 forward/backward paths use the selected context. MoE visits peers in team-local round-robin order, `(rank + offset) % team_size`, local first. It has no mutable device-global schedule tables or schedule-specific host-topology fields in its kernel parameters/iterators. Non-communicating kernels need no group workspace; Triton vocab-parallel CE already uses caller-owned process groups and per-call tensors. ## Memory and lifecycle - Keep the existing buffer pool; there is no TP- or EP-specific pool wrapper. Fixed-size symmetric payload scratch is shared at the process-wide maximum. TP signaling/mapping state is context-specific. MoE device workspace is context-specific so another group's growth cannot invalidate captured pointers. - Symmetric slots follow bootstrap-wide allocation order, not local team-handle values. Cache-hit contexts retain their original allocations and epochs. Reservations remain resident through model offload and inference-engine sleep/wake. - MoE retained forward tensors stay in the symmetric stack, not the buffer pool. `MoEConfig.max_inflight` defaults to 1 and reserves a bound **per EP context**; count all layers and microbatches awaiting backward. Separate contexts have independent stacks; equivalent groups share their limit. Pop in LIFO order within a context. Overflow is rejected before launching or partially mutating stacks. - Low-level MoE metadata lookup/pop select the forward's team; the autograd wrapper forwards it automatically. Destroying a MoE team with live forward intermediates is rejected. Team destruction, pool clear, and finalization invalidate the relevant contexts and Python aliases. - Setup reserves symmetric storage before independent subgroup execution. Legacy lazy allocation now synchronizes preceding device collectives before a fresh WORLD heap allocation, fixing the stacked-forward deadlock found by the full suite. - Shared-scratch entry/exit synchronization prevents a later group from reusing data while an earlier group still consumes it. ## Additional fixes found during review and full-suite execution - Implement SM100 backward's node-local direct-peer path using the existing FP32 direct-peer all-reduce, including context-local epochs and full-tile scatter. The previous failure with `NVSHMEM_DISABLE_NVLS=1` is addressed rather than excluded. - Size Python MoE tile-ID storage for the minimum supported tile size. - Fix the native test build's missing SM90 contract include and exclude the standalone cuBLAS benchmark from the core library source glob. - Replace the obsolete skipped MoE roundtrip placeholder with a real single-PE eager/CUDA-graph forward/backward test. - Remove obsolete CUDA schedule symbols, host permutation builders/uploads, flat schedule-setup exports, and non-RDC module table plumbing. Preserve positive-topology validation and use a 64-bit product to reject overflow. Non-RDC setup initializes the CUDA context explicitly rather than relying on a table upload's side effect. Binding tests now bootstrap NVSHMEM before configuration, matching the native API contract. - Correct the MLP1 numerical reference to use independent cuBLAS BF16-input/FP32-output GEMMs followed by CPU activation arithmetic. The old sequential CPU GEMM disagreed near cancellation, including with FP64 accumulation; an independent cuBLAS probe reproduced the kernel's worst outputs exactly. **Original random inputs, shapes, 1% mean tolerance, and 5% maximum tolerance are unchanged.** Both fused and activation variants now have maximum relative error below 0.4% on B300. Matching native extensions must be rebuilt for the changed host/TVM binding interfaces. Dependency floors and the CUDA 12.9 release workflow are unchanged. ## MoE schedule performance and spilling The topology-arithmetic implementation was rejected after controlled B300 comparisons showed forward latency increases of 1.02-3.03% versus symbolic tables. The final implementation uses simple round-robin and removes the obsolete table infrastructure. Compared against symbolic-table controls on separate eight-GPU B300 and H200 nodes: EP2/EP8, four shapes, three counterbalanced process trials, 15 samples/trial, 20 native calls per captured batch. CUDA-event timings take the maximum across ranks; results below are medians of trial medians. The first trial checks outputs and gradients against Torch references outside timing. Clocks were not locked. H200 comparisons use **non-RDC for both variants**, with `LIGER_CUTE_ENABLE_SM90_NONRDC_MOE=ON`, `LIGER_CUTE_SM90_NONRDC_ALL_CONFIGS=ON`, and `LIGER_MOE_SM90_NONRDC=1`. All 66 local/IB main-kernel specializations are present in each matching cubin. Default-RDC Hopper measurements are not used to evaluate this change. | Single-node latency change vs. tables (negative is faster) | B300 | H200 non-RDC | | --- | --- | --- | | Forward | -2.30% to -0.39% | -1.47% to +0.29% | | Backward | -0.60% to +0.12% | -0.61% to -0.03% | | Full forward + backward | -0.76% to -0.19% | -0.57% to +0.08% | H200 is not uniformly faster: the largest EP2 forward shape has a +0.29% median increase (paired trials +0.29% to +1.44%), accepted by the author as within the performance-noise tolerance for this cleanup. One EP2 full-pass median increases +0.08%, with paired trials crossing zero. All EP8 H200 medians improve. B300's largest EP8 backward case has a small consistent +0.05% median increase; every B300 forward and full-pass paired trial improves. Compiler-reported spill changes for the final kernels: | Path | Spill-store / spill-load bytes vs. tables | | --- | --- | | B300 forward, 8 configs | 24/24 -> 20/20 | | B300 backward, 3 configs | Unchanged at 172/292 | | H200 non-RDC local forward, 10 configs | Unchanged at zero | | H200 non-RDC local backward, 23 configs | Unchanged | | H200 non-RDC IB-capable forward, 10 configs | -8 / -24 bytes | | H200 non-RDC IB-capable backward, 23 configs | -4 / -8 bytes | Spill counts are static ptxas metrics, not runtime memory traffic. `REG:168` is the kernel-level allocation, not the post-`setmaxnreg` MMA-warp limit: Hopper dynamically uses producer/consumer counts of 64/216 forward and 88/208 backward. Single-node timings exercise local transport; lower IB-capable spill counts do **not** establish multi-host HCA performance. ## Execution contract and remaining limits Use one CUDA device per process. Configure serially at a coordinated setup boundary, with every bootstrap rank participating, outside graph capture. Pipeline stages cannot invoke WORLD-collective setup only from post-process chunks; choose an appropriate common bootstrap domain. **Serialize all native launches and graph replays on each process, including across different TP/EP groups**, using one CUDA stream or explicit synchronization. Participating ranks must issue collectives in consistent order. Warm up shapes/configurations before capture. Finish dependent launches, backward passes, and graphs before destroying teams or clearing pools. Live capacity growth remains unsupported; reserve process-wide maxima up front. Legacy low-level MoE setup retains WORLD-collective lazy allocation; use public `configure` for independent EP subgroup execution. LCK rejects an already initialized visible unmanaged NVSHMEM runtime. DeepEP V1 RDMA ownership sharing remains unsupported in either initialization order. Detection does not guarantee interoperability with hidden runtimes or prevent another library from initializing after LCK. Skipping optional setup does not change backend selection or enable an explicitly selected native-only operator. ## Validation All implementation and review work was performed directly, without subagents. Environment: supplied eight-GPU B300 and H200 nodes, CUDA compiler **12.9.86**, Torch **2.9.1.6+cu129**. B300 uses Triton **3.5.1**, system NVSHMEM **3.6.0**, and an isolated system-site-packages venv. This does not validate the release-pinned NVSHMEM 3.6.5 runtime. H200 runs use non-RDC MoE, node-local NVSHMEM transport (`NVSHMEM_DISABLE_NCCL=1`, `NVSHMEM_REMOTE_TRANSPORT=none`, `NVSHMEM_SYMMETRIC_SIZE=6G`), and disable unrelated system pytest plugin autoload. | Final-source run | Result | | --- | --- | | SM100f native core and complete C++ test build | Passed | | SM90a non-RDC all-config cubin, matching native core, and C++ test build | Passed | | B300 complete LCK Python suite and Liger `test/cute` | 180 passed, no skips | | B300 complete C++ CTest suite | 66 passed, 26 hardware/opt-in benchmark skips, 0 failures | | Liger FLSCE frontend contracts | 21 passed | | B300 corrected NVSHMEM-bootstrap binding fixture | 12 passed | | H200 non-RDC bindings plus forward graph targeted retry | 13 passed | | H200 non-RDC schedule and all-dispatch native tests | 4 passed | | H200 non-RDC MoE binding, graph, multi-EP, and frontend suite | 27 passed, 1 architecture-gated frontend skip | | H200 complete C++ CTest suite, non-RDC MoE enabled | 55 passed, 30 hardware/opt-in benchmark skips, 0 failures | | Changed-file Ruff check/format and `git diff --check` | Passed | The complete suite commands are: ```bash PYTHONPATH=liger_cute_kernels:src uv run --no-project --python .venv/bin/python \ python -m pytest --confcutdir=test/transformers -v -o log_cli=false \ --tb=short --disable-warnings --durations=10 \ liger_cute_kernels/test test/cute ctest --test-dir build/sm100f --output-on-failure --timeout 180 -j 1 ``` Multi-group coverage includes 4/6/8 GPUs, overlapping groups, different group sizes, unequal prior launch counts, partial cache hits, split bootstrap domains, retained forwards across context switches, smaller workloads under reserved maxima, overflow rejection, graph replay, device-workspace growth in another EP group, team destruction/recreation, and pool-clear/reconfiguration. Existing mixed-operator coverage exercises distinct MoE EP and FLSCE TP groups. The first H200 Python run exposed missing CUDA-context initialization after table-upload removal and a binding fixture that configured before NVSHMEM initialization. Both were corrected without weakening assertions. Another case stalled in NVSHMEM bootstrap before entering kernels; rerunning with the explicit node-local transport settings above resolved it. These initial failures are not counted as successful runs. Multi-host communication, release-pinned NVSHMEM 3.6.5, DeepEP interoperability, and downstream VeRL changes are not claimed complete. - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Implementation and recorded runs were AI-assisted with GitHub Copilot. No human-review claim is made. --------- Co-authored-by: kolehma8 <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
6ad077c363 |
Cast tl.program_id to int64 in llama4_rope and qwen2vl_mrope (#1342)
## Summary `llama4_rope.py` and `qwen2vl_mrope.py` index the flattened `batch * seq` token dimension with a raw `tl.program_id(0)`. Triton materializes that as int32, so the pointer arithmetic built on top of it wraps once the offset passes `2**31`. This is the same out-of-bounds class that #803 reported and #804 fixed in `rope.py` and `rms_norm.py`, and neither of these two files was touched by that PR. `llama4_rope.py` was added afterwards as a fresh kernel body rather than reusing `_triton_rope`, so it never inherited the fix, and `qwen2vl_mrope.py` predates #804. Fixes #1335 ## Details Where it overflows. `llama4_rope.py` launches on a `(batch_size * seq_len, n_heads_max)` grid and passes `q.stride(1)` as `q_row_stride`, which for the contiguous `(B, S, H, D)` layout is `H * D`. The kernel then computes: ``` base_offset = batch_idx * seq_len + seq_idx # max B*S - 1 q_base = q_ptr + base_offset * q_row_stride ``` so the largest element offset is `(B*S - 1) * H * D`, one row short of `numel(q)`. The int32 product wraps once that passes int32 max, which as a token budget for a single call means `total_tokens > int32_max / (H * D)`: | q width (`H * D`) | overflows at total tokens | | --- | --- | | 4096 (32 heads, head_dim 128) | 524,289 | | 5120 (40 heads, head_dim 128) | 419,432 | | 8192 (64 heads, head_dim 128) | 262,145 | Using #803's own long-context motivation, `B=4, S=420000, H=40, D=128` gives `1679999 * 5120 = 8601594880`, four times past int32 max. Nothing that extreme is required though. A 64-head model at `B=8, S=65536` reaches `4294959104`, already twice over. `qwen2vl_mrope.py` has the same shape of arithmetic, `q_ptr + pid * (n_qh * hd)` with `pid` up to `B*S - 1`, so the same token thresholds apply. The per-head term is safe in both and is left alone. `pid_h * q_head_stride` is bounded by `n_heads * head_dim`, and widening the batch term promotes the sum anyway. The fix is the exact pattern used by #804 and by the Ascend port of `llama4_rope` from #1053, which already carries `tl.program_id(0).to(tl.int64)`. Ascend backend. `backends/_ascend/ops/qwen2vl_mrope.py` has the same gap, which #1335 also flags. That kernel derives its token index inside a loop rather than straight from the program id, so I widened at `pid`, the value actually used for addressing, and left `start_row`, `rows_per_program` and `actual_rows` as int32 so the loop bound and `tl.range` are unchanged. Cost is one `.to(tl.int64)` on a single scalar per program, the same tradeoff #804 accepted for the two hottest kernels in the repo. Related, but deliberately not included. `_block_rms_norm_forward_kernel` in `rms_norm.py`, the blocked forward added after #804, computes `row_idx = tl.program_id(0) * BLOCK_ROW + tl.arange(0, BLOCK_ROW)` with no cast, while `_block_rms_norm_backward_kernel` in the same file does carry `tl.program_id(0).cast(tl.int64)`. That path is gated to `n_cols <= 256`, so it needs on the order of 8.4M rows to wrap, far beyond what the rope kernels need. I kept it out to hold this PR to the scope of #1335. Glad to fold it in here or send it separately, whichever you prefer. ## Testing Done I do not have GPU access, so I have not run the Triton suites, and I would rather say that than tick boxes I cannot honestly tick. What I did verify: - `ruff check` and `ruff format --check` pass on all three files. - The overflow arithmetic above was computed in plain Python against the real strides read from each kernel's launch site, not from the issue text. - Audited every `tl.program_id` under `src/liger_kernel/ops/` to confirm these are the token-indexing kernels still missing the cast. The change is a pure type widening. It is the identity for every offset below `2**31`, so existing shapes stay bit-identical, and reproducing the actual illegal memory access needs a q tensor above 4 GiB, which is not CI-runnable. #804 landed under the same constraint and added no test for the same reason. - Hardware Type: none, no GPU access - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
ed1f856d5b |
fix: request float32 logits from the fused JSD projection GEMM (#1433)
## Summary `fused_linear_jsd_forward` documents that it computes the logits in FP32: ```python # For anything starting from logits to the final JSD loss, we do computation # in FP32 to avoid losing numerical stability. student_logits_chunk = (student_input_chunk @ student_weight.t()).to(torch.float32) ``` It does not. The matmul runs in the input dtype, so the logits are rounded to 8 (bf16) or 11 (fp16) mantissa bits before the cast executes — the cast changes the container, not the contents. cuBLAS already accumulates a bf16 GEMM in FP32; the kernel discards those bits on the store, then casts the rounded result back up. The JSD scalar hides this (≤0.04% error) because it is dominated by large terms. `dx` does not: it is a difference of nearly-equal distributions, so the rounding cancels catastrophically and **the gradients are off by 4–23% in bf16**. `LigerFusedLinearJSD` is the fused path for bf16 distillation, so this lands on distillation training runs. This asks the projection GEMM for an FP32 output instead of casting after the fact. Closes #1432 ## Why `out_dtype` rather than upcasting the operands `torch.mm(x, w.t(), out_dtype=torch.float32)` is *numerically equivalent* to casting both operands to FP32 and running an FP32 GEMM: a bf16×bf16 or fp16×fp16 product is exactly representable in FP32 (8+8 and 11+11 mantissa bits both fit in 24) and the accumulation is FP32 either way. Unlike an explicit upcast it keeps the low-precision tensor cores and needs no FP32 copy of the head. Isolated projection GEMM, `BT=2048, H=4096, V=128256`: | projection | time | peak transient | max abs logit error vs FP64 | |---|---|---|---| | `(x @ w.t()).to(fp32)` (before) | 29.4 ms | 1504 MiB | 1.0e+00 | | `torch.mm(..., out_dtype=fp32)` (this PR) | **25.2 ms** | **1002 MiB** | **2.1e-03** | | upcast both operands to fp32 | 196.2 ms | 1802 MiB | 8.7e-04 | ## Accuracy Relative error of `max|Δ|` for the gradients against a reference differing from the kernel *only* in projecting in FP32. `BT=256, H=1024, V=32000`, `randn` inputs, weight scaled to the stated logit std: | dtype | beta | logit std | `grad_input` before → after | `grad_weight` before → after | |---|---|---|---|---| | bfloat16 | 0.0 (FKL) | 30 | 13.93% → **0.50%** | 12.20% → **0.61%** | | bfloat16 | 0.5 | 10 | 22.90% → **0.38%** | 23.60% → **0.62%** | | bfloat16 | 1.0 (RKL) | 30 | 12.35% → **0.60%** | 10.28% → **0.71%** | | float16 | 0.0 | 30 | 1.25% → **0.06%** | 1.09% → **0.08%** | | float16 | 1.0 | 30 | 2.18% → **0.08%** | 1.53% → **0.04%** | ## Performance End-to-end fwd+bwd through `LigerFusedLinearJSD`, bf16, mean of 10 iterations after 3 warmups: | BT | H | V | `accum_dtype` | before | after | speedup | peak memory | |---|---|---|---|---|---|---|---| | 4096 | 4096 | 128256 | `None` | 662.21 ms | 640.30 ms | 1.034x | 13186 MiB → 13186 MiB | | 2048 | 4096 | 32000 | `None` | 82.05 ms | 79.36 ms | 1.034x | 3096 MiB → 3096 MiB | | 1024 | 2048 | 32000 | `None` | 31.61 ms | 30.47 ms | 1.038x | 1548 MiB → 1548 MiB | | 4096 | 4096 | 128256 | `float32` | 705.77 ms | 681.43 ms | 1.036x | 14689 MiB → 14689 MiB | | 2048 | 4096 | 32000 | `float32` | 89.88 ms | 86.47 ms | 1.039x | 3346 MiB → 3346 MiB | | 1024 | 2048 | 32000 | `float32` | 35.96 ms | 33.60 ms | 1.070x | 1673 MiB → 1673 MiB | Faster, and memory-neutral end-to-end. (The isolated GEMM saves the bf16 logits buffer, but the caching allocator absorbs it at the kernel level — so no memory win is claimed here.) ## This changes bf16/fp16 numerics by default Deliberately, and with no new flag: there is no trade-off to expose, since the FP32 projection is more accurate, faster, and memory-neutral. Unlike `accum_dtype` — which is opt-in because FP32 accumulation costs a real `V x H x 4` buffer — this costs nothing, so gating it would leave every user on the inaccurate path. The kernel's own comment already promises this behavior. Compatibility: - **All existing tolerances pass unchanged**, including the bf16 `test_correctness`, `test_correctness_with_ignore_index`, `test_correctness_functional`, `test_correctness_all_ignored` and `test_amp` cases, which compare against the bf16-projection `TorchLMHeadJSD` oracle. - `out_dtype` is honored under `torch.autocast(bfloat16)` (verified) and compiles under `torch.compile(fullgraph=True)` (verified). - Platforms without `out_dtype` support (torch < 2.8, non-CUDA, sm < 80) keep today's expression, so there is no regression and no new OOM risk there. This mirrors how the `accum_dtype` grad-weight path in this same function already guards its `torch.addmm(..., out_dtype=...)` fast path. - `grad_input`, the `accum_dtype` grad-weight path, `fused_linear_jsd_backward`, the `dispatch("jsd_loss_and_grad", ...)` inner kernel and the autograd signature are all untouched, so the cuTile / CuTeDSL JSD backends are unaffected. The existing `_ADDMM_SUPPORTS_OUT_DTYPE` gate is renamed `_MM_SUPPORTS_OUT_DTYPE` since it now guards both `torch.mm` and `torch.addmm` — same torch 2.8 feature, same value. Happy to drop the rename if you would rather keep the diff to the projection alone. ### What this does *not* fix: AMP with a float32 head The guard requires the operands to share a dtype, so it deliberately skips the autocast case. `torch.amp.custom_fwd` is applied without `cast_inputs`, so autocast stays enabled inside the forward and harmonizes a bf16 hidden state with an fp32 `lm_head` itself. That path has the same magnitude of error (`grad_input` 13.18%, `grad_weight` 11.40% at the shapes above) and this PR does not improve it. That is a scope decision, not an oversight: unlike the same-dtype case there is no free fix. Rounding the fp32 head ourselves once per chunk duplicates autocast's weight cache and still loses the head's mantissa (13.18% → 3.23%); refusing the downcast for an fp32 GEMM gets 13.18% → 0.50% but overrides the low-precision matmul the user explicitly asked autocast for and gives up tensor cores. Both are trade-offs worth their own discussion, so they are recorded in #1432 rather than folded in here. (Also worth knowing: `aten::mm.dtype` is not on autocast's promotion list, so passing `out_dtype` inside an autocast region raises rather than promoting.) `test_amp` is unaffected and still passes. ## Testing New regression test `test_logits_projection_is_not_rounded_to_input_dtype`, parametrized over `dtype x beta`. It compares against an FP32-projection oracle rather than the existing `TorchLMHeadJSD`, which cannot detect this defect because it performs the same inert cast (a bf16 `nn.Linear` followed by `.to(torch.float32)`) and so matches the buggy kernel exactly. The existing cases also draw both operands from `torch.rand`, clustering the logits at mean ≈ 127 / std ≈ 5; `randn` at a realistic spread is what exposes the rounding. Verified as a real regression test: with the kernel change reverted, all 5 active parametrizations fail at 1.25%–22.90%; with it, all pass under 1%. `TorchLMHeadJSD` is intentionally left alone — it still passes and legitimately documents parity with an eager bf16 `lm_head`. The `float16, beta=0.5` combination is explicitly skipped with its reason: there the reference gradient peaks at ~4e-6, below float16's smallest normal (6.1e-5), so quantization of the output buffer dominates and the projection cannot be isolated. The FP32 oracle itself misses the 1% bar there, so asserting it would be testing float16's exponent range, not this change. `make checkstyle` is clean. `make test`: | | passed | failed | skipped | xfailed | |---|---|---|---|---| | this branch | 4457 | 2 | 980 | 15 | | `main` (control, same machine) | 4452 | 2 | 979 | 15 | The +5 passed is exactly this PR's new parametrizations. **The 2 failures are pre-existing on `main`**: I ran the same suite on an unmodified checkout and it fails the identical two tests, `test/transformers/test_grpo_loss.py::test_grpo_loss_vs_trl[2-128-1000-0.04-{False,True}-luspo-sequence-1.5]`, with byte-identical NaN gradient tensors. They also pass in isolation and at file scope (104 passed, 17 skipped), so they are order-dependent within the full run. `test_grpo_loss.py` and `chunked_loss/grpo_loss.py` contain no reference to JSD or `fused_linear_jsd`. Happy to file that separately if it is not already tracked. Every JSD-adjacent suite is green on this branch: ``` pytest test/transformers/test_fused_linear_jsd.py test/ops/test_fused_linear_jsd.py \ test/ops/test_jsd.py test/transformers/test_jsd.py \ test/chunked_loss/test_jsd_loss.py test/ops/test_jsd_chunk_ignore_regression.py # 245 passed, 34 skipped ``` Tolerances in `test/ops/test_fused_linear_jsd.py` are deliberately left alone: that file parametrizes over every registered backend, whose inner JSD kernels differ. --------- Co-authored-by: Claude Opus 5 (1M context) <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Michaël DefferrardandVaibhav Jindal
|
5b828eef2b |
Fix JSD underflow for low-probability vocabulary entries (#1397)
## Summary Fix NaNs in generalized JSD when both student and teacher probabilities underflow to zero in FP32. The previous implementation constructed the mixture distribution in probability space. For log probabilities around `-104` or lower, both probabilities could underflow to zero, resulting in `log(0)` and NaN loss terms. This change: - computes `log(M)` using an element-wise log-sum-exp formulation - evaluates the loss as weighted KL terms to avoid `0 * -inf` - applies the fix consistently to the Triton, Ascend, and cuTile backends - updates the PyTorch reference to use the equivalent stable formulation `torch.logaddexp` - adds standalone and fused-linear regression tests Both new regression tests return NaN on `main` and pass with this change. ## Testing Done Hardware Type: A100 - [x] run `make test` to ensure correctness (18 unrelated chunked GRPO LUSPO failures) - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence (does not exercise JSD) Ascend and cuTile received the equivalent backend change but were not runtime-tested as I don't have the hardware. Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Truong VuandVaibhav Jindal
|
b5bcace9b1 |
Fix group norm backward for inputs with more than three dimensions (#1376)
## Summary `group_norm_backward` reads `hidden_size = dY.shape[-1]` before flattening the gradient while `group_norm_forward` flattens first and then measures, so for an `(N, C, H, W)` input the backward pass uses `W` where it means the per-channel element count `H*W`: it allocates the input-gradient buffer too small, addresses it with the input's real strides (so the kernel writes past the allocation), reduces `dW`/`dB` over one row per channel, and finally raises `RuntimeError: shape '[N, C, H, W]' is invalid for input of size ...` on the closing `view`. `LigerGroupNorm` accepts any input with `dim() >= 3` and its forward pass is correct for 4-D and 5-D, so this hits the standard convolutional case at the first backward pass; 3-D inputs are unaffected, since there `dY.shape[-1]` already equals the per-channel element count. Fixes #1375 ## Details - `src/liger_kernel/ops/group_norm.py`: move the `hidden_size` computation below the flatten and divide by `channels_per_group`, matching the forward pass. For 3-D inputs the kernel receives exactly the same arguments as before, so that path is unchanged; this is host-side scalar arithmetic, no kernel change and nothing to benchmark. - `test/transformers/test_group_norm.py`: factor the existing test body into `_test_liger_group_norm(shape, ...)` and add `test_liger_group_norm_spatial_dims` covering `(2, 6, 4, 8)`, `(4, 32, 16, 16)` and `(2, 8, 2, 4, 4)`. The existing cases keep their parameters. The NPU backend's own `group_norm_backward` (`ops/backends/_ascend/ops/group_norm.py`) already flattens first and then derives its per-channel count as `hidden_size // channels_per_group`, so this brings the default path in line with it. `(2, 6, 4, 8)`, 3 groups, fp32 on H100, against `torch.nn.GroupNorm`: | | forward | dX | dW | dB | |---|---|---|---|---| | before | 2.4e-07 | `RuntimeError` | `RuntimeError` | `RuntimeError` | | after | 2.4e-07 | 4.8e-07 | 1.9e-06 | 0.0 | Under `compute-sanitizer --tool memcheck` with `PYTORCH_NO_CUDA_MEMORY_CACHING=1`, the same 4-D backward reports 28 invalid global writes at the `DX` store (30 errors total) before the fix and 0 after. ## Testing Done - `test/transformers/test_group_norm.py`: 8 passed (5 pre-existing + 3 new). - Reverting only the `group_norm.py` change makes the 3 new cases fail and leaves the 5 pre-existing ones passing. - `test_group_norm.py` is the only file in the suite that exercises this op, so nothing else in `make test` is affected. I ran `ruff check .` and `ruff format --check .` over the repo, but not the full suite or the convergence suite locally — boxes below reflect that. - Hardware Type: H100 NVL - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence ## Environment - Liger-Kernel `a5d795efd2c1436549e70118ef519134e9c27833` (main) + this branch, editable install - GPU: NVIDIA H100 NVL - torch 2.6.0+cu124, triton 3.2.0, CUDA 12.4, transformers 5.14.1, Python 3.10.20 Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Truong VuandVaibhav Jindal
|
559d33e7b1 |
Fix TVD gradient at p == q (#1374)
## Summary `_tv_distance_kernel` computes the gradient as `tl.where(p > q, 0.5 * scale, -0.5 * scale)`, a two-way split of a three-way sign: `p == q` falls into the `p < q` branch and gets `-0.5 * scale` where the derivative of `0.5 * |p - q|` is `0`, so the loss is unaffected but the gradient is wrong on exact ties. Ties are reachable with default settings — self-distillation (student and teacher identical at step 0) makes every element a tie, and padded vocabulary slots are exactly `0.0` in both distributions — and the suite's own reference, `TorchTVDLoss` (`torch.abs(p - q) / 2.0`), already returns `0` there, but the tests never draw a tie because both distributions are random. Fixes #1373 ## Details - `src/liger_kernel/ops/tvd.py`: wrap the existing `tl.where` in a zero branch for `p == q`. Nothing else changes — untied elements take the same path as before. - `test/transformers/test_tvd.py`: new `test_correctness_with_ties`, which ties the trailing slots of the target (fraction `1.0` = self-distillation, `0.5` = padded vocabulary) and compares against `TorchTVDLoss` across all four reductions and both dtypes. Gradient at a tie, `p == q` everywhere, `reduction="batchmean"`, `(8, 512)` fp32 on H100: | | gradient at a tied element | wrong elements | |---|---|---| | before | `-0.0625` | 4096 / 4096 | | after | `0.0` | 0 / 4096 | | `TorchTVDLoss` reference | `0.0` | — | The extra comparison is one predicate and one select per element in a bandwidth-bound kernel. Forward op, `(4096, 32000)` bf16, `do_bench`, three interleaved runs on one H100 NVL: 0.2337 / 0.2338 / 0.2335 ms before, 0.2337 / 0.2339 / 0.2354 ms after — within run-to-run noise. ## Testing Done - `test/transformers/test_tvd.py`: 232 passed (200 pre-existing + 32 new). - Reverting only the kernel change makes 24 of the 32 new cases fail; the 8 that still pass are `reduction="none"`, which returns before the backward pass since the loss itself is correct. - `test_tvd.py` is the only file in the suite that exercises this kernel, so nothing else in `make test` is affected. I ran `ruff check .` and `ruff format --check .` over the repo, but not the full suite or the convergence suite locally — boxes below reflect that. - Hardware Type: H100 NVL - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence ## Environment - Liger-Kernel `a5d795efd2c1436549e70118ef519134e9c27833` (main) + this branch, editable install - GPU: NVIDIA H100 NVL - torch 2.6.0+cu124, triton 3.2.0, CUDA 12.4, transformers 5.14.1, Python 3.10.20 Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Aditya Singh
|
39adfeefab |
Fix AttnRes silently dropping blocks past 32 (#1344)
`_get_max_blocks` capped `MAX_BLOCKS` at 32 for any `n_blocks`:
```python
for mb in [4, 8, 16, 32]:
if n_blocks <= mb:
return mb
return 32
```
`MAX_BLOCKS` is the static trip count of the `tl.static_range(0,
MAX_BLOCKS)` loops in both AttnRes kernels, and the width of the
register-held `scores` and `alpha` vectors. So with more than 32 blocks
the forward normalized the depth softmax over only the first 32 blocks
and dropped the remainder from the weighted sum, and the backward left
`dV[32:]` as the uninitialized `torch.empty_like` buffer. Both are
silent, there is no shape error and no warning.
The fix rounds up to the next power of two instead. `MAX_BLOCKS` still
has to be a power of two because of `tl.arange(0, MAX_BLOCKS)`, and it
still has a floor of 4.
This is behavior-preserving for every `n_blocks` at or below 32, since
rounding to the next power of two reproduces the old ladder exactly: 1,
2, 3 and 4 give 4, then 5 to 8 give 8, 9 to 16 give 16, and 17 to 32
give 32. Only the previously broken range above 32 changes.
Added `N=40` to the forward and backward correctness parametrizations,
which is the smallest case that exceeds the old ceiling.
Verification note, stated plainly: these tests need CUDA and I do not
have a GPU, so I did not run them. The argument here is by reading, plus
the arithmetic above showing the mapping is unchanged at or below 32.
Please run the AttnRes suite before merging.
|
||
|
Aryan DhawanandVaibhav Jindal
|
f01ffb1f5d |
docs: fix input/target shapes in DPO and ORPO forward docstrings (#1406)
## Summary The `forward` docstrings of `LigerFusedLinearDPOFunction` and `LigerFusedLinearORPOFunction` describe `_input` (and DPO's `ref_input`) as `(batch_size * seq_len, hidden_size)` and `target` as `(batch_size * seq_len,)`. The implementation requires `(batch_size, seq_len, hidden_size)` and `(batch_size, seq_len)`: - `LigerFusedLinearPreferenceBase` chunks dim 0 by sequence (`torch.chunk(_input[:len_chosen], dim=0)`), reduces per-token log-probs over dim 1, and documents these shapes itself (`_compute_loss` describes `input_chunk` as `(2 * chunk_size, sequence_length, hidden_size)`). - `test/chunked_loss/test_dpo_loss.py` and `test/chunked_loss/test_orpo_loss.py` construct 3D inputs (`torch.randn(B, T, H, ...)`). - Following the docstrings as written would silently split tokens instead of sequences into chosen/rejected halves. Docstring-only change; `make checkstyle` passes; no kernel or convergence impact. The same wording also appears in `cpo_loss.py` (lines 62, 64), `simpo_loss.py` (lines 70, 72), and `kto_loss.py` (lines 94, 96, 99, including `ref_input`). Happy to fold those into this PR if you prefer one sweep. ## Testing Done - Hardware Type: n/a (docstring-only) - [x] run `make checkstyle` to ensure code style - [ ] run `make test` to ensure correctness (unaffected) - [ ] run `make test-convergence` to ensure convergence (unaffected) Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
cherry77-cloud
|
0c9e8ad056 |
Preserve FP32 LayerNorm statistics for backward (#1498)
## Summary LayerNorm computes its forward statistics in FP32 but saves them in the input dtype in the Triton and legacy cuTile implementations. BF16/FP16 rounding then changes the statistics used by backward, producing incorrect input and affine gradients for near-constant inputs. For a BF16 `[1, 1024]` input with halves `1` and `1.0078125`, unit weight, zero bias, `eps=1e-5`, and an all-ones upstream gradient, Triton returns `dX=-240` in the second half instead of zero. After this change, `dX` is zero and `dWeight` matches an independent FP64 reference. Allocate both Mean and RSTD in FP32 and store the FP32 values directly in both affected implementations. This adds four bytes of saved statistics per row for BF16/FP16 inputs. The dispatcher cuTile, CuTe DSL, and Ascend implementations already preserve FP32 statistics. The diff contains only the two production files. ## Testing Done - Hardware: NVIDIA RTX A1000 Laptop GPU; PyTorch 2.13.0+cu130; Triton 3.7.1. - Existing LayerNorm transformer and op suites: 60 passed. - External regression harness: 12 passed across BF16/FP16/FP32, shapes `[1, 1024]` and `[3, 1000]`, and separate mean/rstd rounding cases. The original implementation fails all eight low-precision cases. The reference differentiates an independent FP64 expression from the same quantized inputs. This harness is not included in the diff. - Legacy cuTile: FP32 allocation/launch checks passed for three dtypes and two shapes; all 18 forward, input-only backward, and combined backward variants compiled offline for `sm_100`. - cuTile numerical execution was not available on this `sm_86` GPU. The installed compiler also does not support `sm_90`, so Hopper compilation was not validated. - [x] `make checkstyle` with the repository's Ruff 0.14.11 version. - [ ] Full `make test` (scoped LayerNorm suites run instead). - [ ] `make test-convergence` (not run). |
||
|
|
e064108ef6 |
Add fused CE + TV distance kernel (#1384)
## Summary Adds `LigerFusedCETVDFunction`: per-token cross-entropy and total variation distance against a teacher, computed in one streaming pass over the vocabulary. Distillation losses that pair a hard-label term with a distributional distance currently materialize two `(BT, V)` float32 softmax temporaries. This fuses both terms and keeps only `O(BT)` state for backward, rebuilding the distributions from logits the caller already holds. `ops/tvd.py` is unchanged — it consumes probabilities, so it cannot avoid those temporaries. Both terms are returned unreduced, so the caller keeps ownership of masking, per-token weighting, and normalization. ## Testing Done `test/transformers/test_fused_ce_tvd.py`: forward and backward parity against eager autograd in fp32 and bf16, over vocab sizes 128–131072 (multi-tile) and non-power-of-two tails, plus `ignore_index`, teacher-receives-no-gradient, and shape validation. Exact `p == q` ties are covered in the backward pass, including self-distillation: the sign uses a zero subgradient at ties so the kernel matches `torch.abs` backward. Swapping it for a two-way `p > q` split fails both tie tests — the same bug class #1374 fixes in `ops/tvd.py`. At `V=151936` bf16, against a chunked and activation-checkpointed eager reference: marginal memory falls from ~0.87 MB/token to ~0 (890 MiB at 1024 tokens), and HBM traffic ~5.9x. - Hardware Type: AMD Radeon 890M (ROCm) - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence `make test` reached 2741 passed / 858 skipped before stopping on `test_fused_moe.py::test_correctness`, which is unrelated and simply too slow on this integrated GPU; tests ordered after it did not run. The timings above come from that same box, so the traffic ratio is the transferable number rather than the wall-clock. Please run the full suite and convergence in CI. --------- Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
Nilesh PatilandVaibhav Jindal
|
72e9159d86 |
Load unsupported model types instead of raising KeyError in AutoLigerKernelForCausalLM (#1404)
## Summary
`AutoLigerKernelForCausalLM` raises `KeyError` for any model type
without a Liger patching function, instead of loading the model
unpatched.
```python
AutoLigerKernelForCausalLM.from_pretrained("gpt2") # KeyError: 'gpt2'
```
## Details
`_apply_liger_kernel` already handles this case deliberately:
```python
if model_type not in MODEL_TYPE_TO_APPLY_LIGER_FN.keys():
logger.info(f"There are currently no Liger kernels supported for model type: {model_type}.")
return
```
But `auto_model.py` then indexes the same dict unconditionally, to strip
the kwargs that function consumed:
```python
_apply_liger_kernel(model_type, **kwargs)
apply_fn = MODEL_TYPE_TO_APPLY_LIGER_FN[model_type] # KeyError for anything unsupported
apply_fn_signature = inspect.signature(apply_fn)
```
so the graceful path above is unreachable from the wrapper. Both
`from_pretrained` (line 36) and `from_config` (line 55) do it.
This contradicts the documented contract. The class docstring says it is
"a drop-in replacement for AutoModelForCausalLM that applies the Liger
Kernel to the model **if applicable**", and the README says "**If the
model type is supported**, the modeling code will be automatically
patched." Both read as: unsupported models still load.
The intended usage is a one-line swap in a training script:
```python
model = AutoLigerKernelForCausalLM.from_pretrained(cfg.model_path) # was AutoModelForCausalLM
```
which now hard-fails the moment that script is pointed at a model Liger
does not patch. It also means each new model type in `transformers` is a
crash here until a patch function is added.
## Change
An unsupported model type consumed no kwargs, so all of them belong to
the underlying `AutoModel` call. A small helper returns the kwargs
unchanged in that case, and both classmethods use it.
It also fixes `from_config` returning `None` when `model_type` is falsy:
it logged "No Liger kernels will be applied" and then returned no model
at all, so callers got an `AttributeError` on `None` downstream. The
redundant `model_type = config.model_type` on the very next line shows
the block was meant to fall through rather than bail.
## Testing
Three cases added to `test/transformers/test_auto_model.py`, following
the existing mock style. All are CPU-only.
Without the change:
```
src/liger_kernel/transformers/auto_model.py:36: KeyError: 'unsupported_model_type'
src/liger_kernel/transformers/auto_model.py:55: KeyError: 'unsupported_model_type'
test/transformers/test_auto_model.py:141: AssertionError: Expected 'from_config' to be called once. Called 0 times.
3 failed, 2 passed
```
With it:
```
5 passed
```
Also checked against a real config, no mocks:
```
AutoLigerKernelForCausalLM.from_config(GPT2Config()) -> GPT2LMHeadModel
same class as AutoModelForCausalLM: True
```
Regression, comparing a clean checkout with the patched tree over
`test_auto_model.py`, `test_monkey_patch.py`, `test_utils.py` and
`test_trainer_integration.py`:
```
clean: 89 passed, 8 skipped
patched: 92 passed, 8 skipped
```
The delta is exactly the three new tests; the 8 skips are GPU-gated and
identical in both runs. `ruff check` and `ruff format --check` both
pass.
**Hardware Type:** CPU only. This change touches no kernel code, so I
ran the CPU-runnable suites above rather than the full GPU matrix. Happy
to have CI cover the rest.
Co-authored-by: Vaibhav Jindal <[email protected]>
|
||
|
|
998ac696fa |
fix(llama4): honour skip_logits in lce_forward (#1445)
## Summary `llama4.lce_forward` is the only one of the 42 forwards under `transformers/model/` that has no `skip_logits` parameter. It still decides on its own: ```python if self.training and (labels is not None or shift_labels is not None): ``` #704 added `skip_logits` precisely so a caller could override that, and swept every model that existed on 2025-05-23. llama4 arrived a month later in #740 and never picked it up; everything added since has it. Because the signature does not name it, a `skip_logits` argument lands in `**kwargs` and is silently ignored, so `skip_logits=False` during training still takes the fused path and returns `logits=None`. That is the one thing the flag exists to prevent, and the convergence harness relies on it, setting `eval_batch["skip_logits"] = False` before reading `eval_output.logits`. The change mirrors `mllama`, which has the identical `LigerForCausalLMLoss` shape. `skip_logits=None` evaluates to the previous condition, so nothing moves for callers that do not pass it. ## Testing Done New `test/transformers/test_llama4.py`. Both cases stay on the non-fused branch, so they run without a GPU: - `skip_logits=False` in training mode returns logits and a loss. On `main` this fails: the flag is ignored, the fused path runs, and on a CPU box Triton raises `RuntimeError: 0 active drivers`. On a GPU box it fails on `logits is None` instead. - `skip_logits=True` with no labels raises `ValueError`, as it does on every other model. On `main` this is `DID NOT RAISE`. - `pytest test/transformers/test_monkey_patch.py test/transformers/test_llama4.py`: 65 passed / 1 skipped (63 + the 2 new; the skip is `muse_glimmer` missing from the installed transformers). - `ruff check .` and `ruff format --check .` with the v0.14.11 pinned in `.pre-commit-config.yaml`: clean. Hardware Type: CPU. Ran on transformers 5.9.0 / torch 2.11.0. --------- Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
cherry77-cloudandVaibhav Jindal
|
946a3693c1 |
Fix softcap padding in cross entropy (#1490)
## Summary With softcap enabled, masked vocabulary-tail lanes change from `-inf` to the finite value `-softcap` and incorrectly contribute to the online softmax denominator. For three zero logits with `softcap=1`, this produces a loss of `1.21428323` instead of `log(3) = 1.09861229`, and changes the gradients. Restore `-inf` for out-of-range lanes immediately after softcapping, before the max/sum reductions. This one-line fix applies to the shared Triton kernel used by cross entropy and fused linear cross entropy. ## Testing Done Hardware: NVIDIA RTX A1000 Laptop GPU, 4 GiB; Triton 3.7.1. Existing tests used PyTorch 2.13.0+cu130; standalone checks used PyTorch 2.12.1+cu130. - Existing CE tests: **94 passed**, 91 deselected. Selected cases fit the 4 GiB GPU: ```bash python -m pytest test/transformers/test_cross_entropy.py -q -x \ -k '423 or true_class_grad_confident_predictions or structured_output' --disable-warnings ``` - Standalone GPU validation: **10 checks passed**, comparing loss and gradients against PyTorch through the public CE and fused-linear-CE APIs. Covered the minimal reproduction, nonzero logits, a multi-block tail (`V=32769`), large softcap with negative logits, FP32/FP16/BF16, and no-padding/no-softcap controls. Validation scripts are outside the production-code diff. - CUDA-graph timing of input copy plus the loss/gradient kernel, `[128, 32000]`, BF16, `softcap=30`: median **0.2194 ms before / 0.2210 ms after** across five alternating measurements. Measured full-pass incremental peak allocation was unchanged at **7.815 MiB**. - [x] `make checkstyle` passed (Ruff 0.15.20). - [ ] Full `make test` was not run; the targeted suite above was run. - [ ] `make test-convergence` was not run. Signed-off-by: taking-lying-flat <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
znzrybandVaibhav Jindal
|
70fe0333fe |
fix(triton): widen softmax/sparsemax row index to int64 (#1497)
## Summary Widen the softmax and sparsemax program IDs to int64 before row-offset address arithmetic. With an int32 `tl.program_id(0)`, `row * row_stride` wraps once `n_rows * n_cols` crosses 2^31, and the kernels hit an illegal memory access (e.g. `n_cols=32768`, `n_rows > 65536`). This applies the same safeguard as #251, #804 and #1482 to the four softmax kernels and the two sparsemax kernels. Fixes #1496. ## Details - The softmax multi-block kernels are currently unreachable (see #1252), but they get the same cast so they stay safe once that path is enabled. - The sparsemax regression test launches `_sparsemax_forward_kernel` / `_sparsemax_backward_kernel` directly: at this size the `torch.sort` in the autograd path would need ~26 GB extra for sorted values and int64 indices. - The Ascend backends (`ops/backends/_ascend/ops/{softmax,sparsemax}.py`) are left unchanged since I have no NPU to test on. ## Testing Done - Hardware Type: NVIDIA L4 (24 GB) - [x] `python -m pytest -v test/transformers/test_softmax.py test/transformers/test_sparsemax.py` (36 passed, including both new large-tensor tests, which need 17.2 GB free VRAM) - [x] The two new large-tensor tests fail on `main` with `illegal memory access was encountered` - [x] `ruff check` and `ruff format --check` on the changed files - [ ] `make test` (only the softmax and sparsemax suites were run) - [ ] `make checkstyle` (ran the equivalent `ruff` commands above on the changed files) - [ ] `make test-convergence` (not run; the change only affects address arithmetic) Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Alexandre Kohler
|
d76315bd77 |
fix(kto): sum chosen/rejected rewards per sample when chunk_size > 1 (#1495)
## Summary Root cause: in `LigerFusedLinearKTOFunction.preference_loss_fn`, `rewards` has shape `(chunk_size,)`, but it is multiplied by `preference_labels_chunk.unsqueeze(1)`, which has shape `(chunk_size, 1)`. Broadcasting turns the product into a `(chunk_size, chunk_size)` matrix, so `chosen_rewards_sum` becomes `num_chosen_in_chunk * sum(rewards)` and `rejected_rewards_sum` becomes `num_rejected_in_chunk * sum(rewards)`. With the default `chunk_size=1` the matrix is 1x1 and the result is right, which is why the existing tests did not catch it. With `chunk_size > 1` and mixed labels in a chunk, both reward metrics are wrong. The loss and the gradients are not affected. Example (B=4, labels `[True, False, True, False]`, with a reference model): `chunk_size=1` returns chosen 0.193 and rejected -1.435, while `chunk_size=4` returns -2.485 for both (2 x the sum of all four rewards). Fix: multiply by the 1D mask. This is the same change #1441 made for `chosen_logps_sum` / `rejected_logps_sum` in `fused_linear_unpaired_preference.py`; the two reward lines in `kto_loss.py` still had the old `unsqueeze(1)`. ## Testing Done Added `test_kto_rewards_sum_with_chunk_size`, which checks that the chosen and rejected reward sums with `chunk_size` 2 and 4 match `chunk_size=1`. On main it fails (chosen sum -85.45 with `chunk_size=2` and -170.90 with `chunk_size=4`, vs -38.35 with `chunk_size=1`) and it passes with the fix. - Hardware Type: NVIDIA GeForce GTX 1060 6GB, PyTorch 2.5.1 (the new test was also run CPU only) - [x] `python -m pytest test/chunked_loss/test_kto_loss.py` (54 passed) - [x] `make checkstyle` checks: `ruff check .` and `ruff format --check .` pass (ruff 0.14.11, as pinned in `.pre-commit-config.yaml`) - [ ] `make test` (not run; Triton is not available on this Windows host) - [ ] `make test-convergence` (not run; KTO is not covered by the convergence tests) |
||
|
Yulong Cao
|
40a9d8a61b |
[Feature] Support apply_liger_kernel_to_deepseek_v3 (#1424)
## Summary Adds Liger Kernel support for DeepSeek-V3 (`deepseek_v3` model type), closing the gap left open by #623. - `model/deepseek_v3.py`: new `lce_forward`. In training with labels, `skip_logits` defaults to `True`, so logits are never materialized (FLCE path via `LigerForCausalLMLoss`); otherwise logits are materialized and the HF-native `self.loss_function` is used. Returns `LigerCausalLMOutputWithPast` with `token_accuracy` / `predicted_tokens`. - `monkey_patch.py`: `apply_liger_kernel_to_deepseek_v3` with class-level and instance-level patching, registered in `MODEL_TYPE_TO_APPLY_LIGER_FN` (`AutoLigerKernelForCausalLM` picks it up automatically). - RMSNorm: `DeepseekV3RMSNorm` → `LigerRMSNorm`, including attention `q_a_layernorm` / `kv_a_layernorm`. - SwiGLU: dense MLPs and shared experts (`DeepseekV3MLP`) → `LigerQwen3MoeSwiGLUMLP` on all supported versions. Routed experts: batched `DeepseekV3Experts` → fused `LigerExperts` on transformers v5+; per-expert patch on transformers v4. - RoPE: unsupported — DeepSeek-V3 uses interleaved partial RoPE, incompatible with `liger_rotary_pos_emb`. `rope=True` emits a `warning_once` and skips (no silent behavior change); RoPE is intentionally not listed in the README table. - Tests: instance monkey-patch test (asserts both v5 and v4 branch outcomes), 4 convergence tests (bf16/fp32 × FLCE/logits) driven by a mini `deepseek_v3` config, `revert_liger_kernel_to_deepseek_v3` util, README entry. ## Details - The transformers v4 branch (per-expert patching) mirrors the existing mixtral / qwen3_moe patterns but could not be exercised locally (only transformers 5.16.1 is installed here); it is covered by the CI v4 matrix. - On transformers v5+, `LigerExperts` accepts `DeepseekV3Config` (`num_local_experts` aliases `n_routed_experts`), and its batched expert weights (`(N, 2*I, H)` / `(N, H, I)`) match the HF layout, so pretrained checkpoints load unchanged. ## Testing Done | Check | Status | |---|---| | Import | PASS | | Instance Patching | PASS | | Convergence bf16 FLCE | PASS | | Convergence bf16 with_logits | PASS | | Convergence fp32 FLCE | PASS | | Convergence fp32 with_logits | PASS | | Convergence multimodal | SKIP (text-only model) | | Checkstyle | PASS | Exact commands run on this change: - `pytest test/transformers/test_monkey_patch.py -k deepseek_v3` → 1 passed - `pytest test/convergence/{bf16,fp32}/test_mini_models.py test/convergence/{bf16,fp32}/test_mini_models_with_logits.py -k deepseek_v3` → 4 passed - `ruff check` + `ruff format --check` clean on all changed files ### Memory Single A100-SXM4-80GB, bf16, forward + backward, no optimizer step. Randomly initialised (no public checkpoint). Real dims: `vocab_size=129280`, `hidden_size=7168`, `num_attention_heads=128`, `intermediate_size=18432`, `moe_intermediate_size=2048` and the MLA projections. Slimmed and disclosed: 2 layers with `first_k_dense_replace=1` (real: 61 / 3), 16 routed experts with `num_experts_per_tok=2`, `n_group=2`, `topk_group=1`, `routed_scaling_factor=1.0` (real: 256 / 8, 8 / 4, 2.5). | batch | seq | peak mem HF | peak mem Liger | saved | |------:|----:|------------:|---------------:|------:| | 1 | 2048 | 13.1 GiB | 12.61 GiB | 4% | | 1 | 4096 | 16.24 GiB | 13.02 GiB | 20% | | 1 | 8192 | 26.18 GiB | 16.4 GiB | 37% | | 1 | 16384 | 46.06 GiB | 23.15 GiB | 50% | | 1 | 32768 | OOM | 36.64 GiB | | | 4 | 8192 | OOM | 36.67 GiB | | Savings grow with the token count: HF's loss pipeline materializes the full `[tokens, 129280]` logits tensor and its fp32 upcast/softmax (~2.4 GiB per additional 1K tokens), while FLCE replaces it with chunked transients plus a fixed 3.45 GiB fp32 weight-gradient accumulator — 4% at 2K, 50% at 16K, and at 32K tokens (batch 1 x 32K or 4 x 8K) HF no longer fits the 80 GB GPU while Liger holds ~36.7 GiB. ### Throughput Same hardware and dtype, batch 1, seq 4096, varying depth: | layers | lm_head share | tok/s HF | tok/s Liger | ratio | |-------:|--------------:|---------:|------------:|------:| | 2 | 27% | 8,457 | 8,172 | 0.97x | | 4 | 18% | 5,117 | 5,039 | 0.98x | | 8 | 10% | 2,871 | 2,862 | 1.00x | The ratio converges to 1.00x as depth dilutes the loss segment. The shipping 61-layer model puts `lm_head` at ~0.14% of parameters, so the end-to-end cost is ~1.00x. (On a 15 GB T4 in fp16 — sm_75 has no bf16 tensor cores — slimmer variants show the same pattern with the fixed accumulator dominating until ~4K tokens; happy to share those numbers too.) - Hardware Type: RTX 4060 Laptop GPU (8GB) — transformers 5.16.1, torch 2.11.0+cu128; benchmarks: NVIDIA A100-SXM4-80GB — transformers 5.8.1, torch 2.11.0+cu130 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence Closes #623 |
||
|
|
55bc58a842 |
fix(triton): widen blocked RMSNorm row index (#1482)
## Summary Widen the blocked RMSNorm forward program ID before row-offset address arithmetic. The default blocked path (`n_cols <= 256`, `n_rows >= 32768`) otherwise overflows a 32-bit row address at 2^31 elements. This completes the same safeguard introduced for the non-blocked RMSNorm kernels in #804. Fixes #1483. ## Testing Done - Hardware Type: NVIDIA RTX 2000 Ada Generation Laptop GPU (8 GB) - [x] `python -m pytest -q test/transformers/test_rms_norm.py -k block_rms_norm -x` (1 passed; large CUDA regression skipped because it requires 8.6 GB free VRAM) - [x] `ruff check src/liger_kernel/ops/rms_norm.py test/transformers/test_rms_norm.py` - [x] `ruff format --check src/liger_kernel/ops/rms_norm.py test/transformers/test_rms_norm.py` - [ ] `make test` (make is unavailable on this host) - [ ] `make checkstyle` (make is unavailable on this host; its ruff commands above passed) - [ ] `make test-convergence` (not run) Co-authored-by: Srijan Keshri <[email protected]> Co-authored-by: Jari Kolehmainen <[email protected]> |
||
|
cherry77-cloudandVaibhav Jindal
|
9358939c4d |
Fix ignored-token masking in cosine distillation (#1478)
## Summary Cosine distillation currently sums soft loss over every token, while the shared distillation base divides by the number of non-ignored targets. Consequently, padding or prompt positions labeled with `ignore_index` still change the loss and receive soft-loss gradients. For example, appending one ignored token to a perfectly aligned valid token changes the soft loss from zero to approximately 0.292893. Mask the per-token cosine loss before summing, matching the existing valid-token denominator and the JSD implementation's masking semantics. Calls without masking arguments retain their existing behavior. Clarify the mask and reduction contract in the cosine and base-class docstrings; the shared normalization logic is unchanged. ## Testing Done - Hardware: NVIDIA RTX A1000 Laptop GPU, 4 GB; PyTorch 2.13.0+cu130, Triton 3.7.1. - All 40 existing tests in `test/chunked_loss/test_cosine_loss.py` passed, including a separate compiled run with Dynamo's recompile limit raised to 128 and failure on limit exhaustion enabled, to avoid falling back to eager execution across parameter combinations. - 48 additional local regression cases passed, covering the FP64 counterexample, FP32/BF16, eager/compiled execution, default/custom ignore indices, mixed/all-ignored inputs, fully ignored chunks, student input/weight/bias gradients, and invariance when ignored tokens are appended or changed. Four counterexample cases failed before the fix. These local regression tests are not included in this production-code-only patch. - [ ] `make test` (full suite not run; 88 focused cases passed) - [x] `make checkstyle` - [ ] `make test-convergence` (not run; focused loss and gradient validation above) Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Yulong CaoandJari Kolehmainen
|
aed0dfbb58 |
refactor(kl_div): route fused-linear KL inner kernel through the dispatcher (#1476)
## Summary Follow-up promised during the #1423 review (@kolehma8): expose the fused-linear KL's inner computation as a dispatchable primitive, mirroring what `jsd_loss_and_grad` already does for `fused_linear_jsd`. `LigerFusedLinearKLDivFunction` used to launch `_kl_div_kernel` directly from its chunk loop, bypassing impl selection — the same class of bug that the `jsd_loss_and_grad` registration guards against (composed ops silently bypass the dispatcher; cf. #1228). After this PR: - `kl_loss_and_grad` is a registered non-autograd primitive in `ops/backends/_triton/kl_div.py`, wrapping `_kl_div_kernel` (the kernel stays in `ops/fused_linear_kl_div.py`, mirroring the FLCE backend layout); - the composed op dispatches the primitive per chunk and **rebinds the returned dx**, so a future impl that returns a fresh tensor instead of writing in place needs no changes to the composed op; - `functional.py` declares the primitive's op locations. No user-facing API or numerical behavior changes: same kernel, same chunk geometry, same GEMMs. ## Why this matters beyond tidiness - The elementwise `kl_div` op already ships a CuTe DSL backend (#1418). A CuTe DSL / cuTile `kl_loss_and_grad` primitive is the natural next step (cuTile registrations are actively landing, e.g. #1456, tracked in #1205) — and after this PR it only needs a `@register_op` under the same primitive name; the composed op picks it up via auto-select. - Small contract detail surfaced while wiring this: the kernel early-returns for `ignore_index` rows without writing the loss, so the primitive pre-zeroes its per-row loss buffer (the composed op's old `loss_1d` used to provide that zeroing implicitly). ## Validation - New `test/ops/test_fused_linear_kl_div.py` (14 cases): primitive numerics vs a PyTorch reference (2 shapes × 3 dtypes × with/without `ignore_index`), `available_backends` sanity, and a tracking-dispatch test asserting the composed op's default path routes through `dispatch("kl_loss_and_grad")`. - Regression (RTX 4060 Laptop): `test/transformers/test_fused_linear_kl_div.py` (full behavioral net incl. AMP and `detect_anomaly` paths) + `test/ops/test_kl_div.py` + `test/ops/test_fused_linear_jsd.py` + `test/backends` → **184 passed, 19 skipped** (env-conditional skips). - `ruff check` + `ruff format --check` clean on all touched files. GPU CI workflows need maintainer approval to run on this fork PR (same as #1423/#1424) — happy to provide additional numbers if useful. Co-authored-by: Jari Kolehmainen <[email protected]> |
||
|
Daniel LentonandJari Kolehmainen
|
ec10acbbe6 |
perf(ppo): skip dW/dB calculation when lm_head parameters are frozen (#1488)
## Summary `LigerFusedLinearPPOBase`, which backs `LigerFusedLinearGRPOLoss` and its variants (GRPO, DAPO, Dr. GRPO, BNPO, CISPO, SAPO, ...), always makes the `lm_head` weight (and bias) differentiable (`weight.detach().requires_grad_(True)`). As a result, `_ChunkedSelectiveLogProbFunction.backward` allocates a float32 `[V, H]` gradient buffer and runs a `grad_logits.T @ hidden` GEMM on every sequence × vocab chunk, even when the weight does not require grad. With LoRA/PEFT the `lm_head` is frozen and that gradient is discarded. This PR differentiates the weight and bias only when they require grad, as #1441 did for the preference losses. `_ChunkedSelectiveLogProbFunction` now computes only the gradients autograd asks for (`ctx.needs_input_grad`). The loss and the input gradients are unchanged, and trainable parameters take the same path as before. ## Details Forward + backward in bf16, `LigerFusedLinearGRPOLoss(beta=0.0, use_ref_model=False)` with its other defaults (compiled), B = 8, T = 1024, H = 4096. Each figure is the median of 10 timed iterations after 3 warm-ups, then the median over separate processes, with main and this branch run alternately. | V | lm_head | main | this PR | |---|---|---|---| | 128,256 | frozen | 471.6 ms, 6.14 GiB | **263.7 ms, 3.25 GiB (−44%)** | | 151,936 | frozen | 557.4 ms, 7.22 GiB | **311.7 ms, 3.79 GiB (−44%)** | | 128,256 | trainable | 472.2 ms | 471.6 ms | | 151,936 | trainable | 558.8 ms | 558.0 ms | #1294 also edits `fused_linear_ppo.py`, changing the GEMM operand dtypes in the same backward. The two changes are independent, but whichever lands second needs a rebase. ## Testing Done - Hardware Type: NVIDIA H100 80GB HBM3 (torch 2.14.0+cu130, triton 3.8.0) - [ ] run `make test` to ensure correctness: I ran the test file for the changed base rather than the full suite. `test/chunked_loss/test_grpo_loss.py`: 994 passed and 614 skipped on this branch (976 passed and 614 skipped on main, plus 18 new). - [x] run `make checkstyle` to ensure code style (ruff 0.15.22 `check` and `format --check` on the changed files) - [ ] run `make test-convergence` to ensure convergence: not run; the convergence tests do not exercise the chunked losses. New tests: - A frozen vs a trainable `lm_head`, with and without bias, with and without a reference model, and at a vocab size above the 4096 vocab chunk, so both chunk loops run. - A trainable weight with a frozen bias. - A check that `_selective_logprob_backward` is asked for the weight and bias gradients only when they require grad. The frozen-head test calls `torch.compiler.reset()` first, as `test_correctness_large_seq_exercises_chunking` does. Without it, starting from a cold compile cache on CPU, two of its compiled cases hit a dynamo guard error ("sources must not be empty for symbol"). The error reproduces identically on main with the same test, so it is not related to this change. Co-authored-by: Jari Kolehmainen <[email protected]> |
||
|
Rain
|
fbdea0ff3b |
fix(npu): Replace Ascend A5 comments with 950 (#1484)
## Summary Follow https://github.com/linkedin/Liger-Kernel/issues/969#issuecomment-5773453873 Replace the machine-type name `A5` with `950` in Ascend comments and docstrings, matching the `A2` / `A3` / `950` naming rule. - Touched files: `ub_manager.py`, `group_norm.py`, and `grpo_loss.py`. - Kernel logic, dispatch, and numerics are unchanged. ## Testing Done - Hardware Type: N/A (comment-only) - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Rain
|
09f7febf1a |
fix(npu): Ascend A3/A5 UB tile cap and A5 workarounds for group norm, GRPO, fused MoE (#1464)
## Summary This PR lands two commits that unblock Ascend **A5** regression coverage for several transformers ops while **triton-ascend** fixes are pending upstream. 1. **`ub_manager`**: Hard-cap block tiles using detected UB and the default **auto-multi-buffer (~4/3)** live-set factor, so margin-only sizing cannot pick tiles that only fit mis-reported UB (notably some **A5 GLU** paths). Uses **A3 ≈ 192 KiB / A5 ≈ 256 KiB** from runtime detection—not SoC-specific branches in kernels. 2. **Kernel workarounds (A5)**: Targeted changes in **group norm**, **GRPO / CISPO / SAPO**, and **fused MoE** so the same Triton semantics pass on A5. Several paths are **temporary** until the linked **triton-ascend** issues are fixed; we intend to revert or simplify once upstream ships. ## Upstream triton-ascend issues (filed) These are **compiler/backend bugs on the A5 pipeline**, not application logic errors. This PR works around them; long-term fix belongs in [triton-lang/triton-ascend](https://github.com/triton-lang/triton-ascend/issues). | Issue | Symptom on A5 | Workaround in this PR | |-------|----------------|------------------------| | [#2239](https://github.com/triton-lang/triton-ascend/issues/2239) | BiShengHIR `memref.extract_strided_metadata` when **dynamic `task_base` + `tl.sum(..., axis=1)`** on padded 2D tiles (e.g. group norm `(1,1,1,3)`) | `_group_norm_forward_single_task_kernel`: **`tl.sum(X_vals)`** / **`tl.sum(X_vals * X_vals)`** (equivalent for `cpg=1`) | | [#2246](https://github.com/triton-lang/triton-ascend/issues/2246) | Compile-time **UB overflow** for large-vocab scans with **`tl.static_range`** + bf16→fp32 online softmax (`V=151936`) | Inner vocab loops: **`tl.static_range` → `tl.range(0, N, BLOCK_N)`** in `grpo_loss.py` (fwd / softmax / bwd) | | [#2247](https://github.com/triton-lang/triton-ascend/issues/2247) | **2-D reduction** over `[TPT, K_POW2]` drops lanes when **`K_POW2 < 8`** (MoE router histogram undercount) | `_moe_router_histogram_kernel`: pad inner dim to **`HIST_K_POW2 = max(8, K_POW2)`** with unchanged `k < K` mask | Additional **fused MoE** A5 fixes in the same commit (gather **`TILE_H` padding**, backward **one row per program** for `tl.dot`, **H on grid** for `dX_expanded`, bf16 K-reduce via mask weights) are documented internally and may get separate upstream issues; they are **not** covered by #2239/#2246/#2247. ## Details ### `ub_manager` (`8ba7d90`) - Cap: `max_block_size = min(margin_solve, ub_capacity // (bits_per_block * 4/3))`. - Does **not** use `multiplier * 2` (would unnecessarily shrink A5 GLU tiles that legitimately fit 256 KiB UB). - On A3, plain CE/LCE fwd vocab tile may step down (e.g. 16384→8192) where 4/3 straddles a power of two; SwiGLU / RMSNorm / RoPE tiles unchanged in our A/B checks. ### Group norm (`d30508b`) - Single-task forward: avoid **`axis=1`** reduce on A5 (see #2239). ### GRPO loss (`d30508b`) - Replace inner **`tl.static_range`** with **`tl.range`** for vocab dimension (see #2246). - Keeps **`multibuffer=True`**; shrinking `BLOCK_N` alone does not fix A5 compile UB for `static_range`. ### Fused MoE (`d30508b`) - Histogram: **`HIST_K_POW2`** (see #2247). - Token gather: **`TILE_H = BLOCK_H + BLOCK_H // 2`**; K-reduction uses explicit weights where bf16 **`tl.sum(axis=0)`** was unsafe on A5. - Backward down-proj / dX: row-expanded grid, **`pid_base`**, **`row_mask = (m_offs == 0) & …`**; dX launch **`(num_rows, n_h_tiles)`** instead of nested H+I **`tl.range`** in one program. ## Testing done - **Hardware**: Ascend **A3** (192 KiB UB) <img width="1039" height="115" alt="A3 第二次执行" src="https://github.com/user-attachments/assets/6a8208ae-ff70-49c4-a245-597e6cc16ec1" /> Ascend **A5** (256 KiB UB) <img width="925" height="109" alt="A5 全量测试用例执行情况" src="https://github.com/user-attachments/assets/3c9ccb90-7b85-4214-a045-7876c0c3a603" /> The difference in the execution results of the 76 test cases between A3 and A5 is due to different `transformers` versions. - **`make test`** (transformers): full suite green on A5 after above (run **`test_fused_moe.py` in a subprocess** if combining with other tests—MoE correctness can poison the device in-process). - **`make checkstyle`** - **`ub_manager` A/B** (same machine, old vs new): SwiGLU / RMSNorm / RoPE tile sizes and latency unchanged where expected; CE flat on V=32k/128k; FLCE spot check per prior notes on this PR. ## Follow-up - Drop or narrow workarounds when [#2239](https://github.com/triton-lang/triton-ascend/issues/2239), [#2246](https://github.com/triton-lang/triton-ascend/issues/2246), and [#2247](https://github.com/triton-lang/triton-ascend/issues/2247) are resolved in a released **triton-ascend** build. - Track remaining MoE A5 items against future upstream issues as filed. |
||
|
|
c5b9f31d3e |
Fix CUTEDSL FLCE functional wrapper arguments (#1475)
## Summary - forward `chunk_mem_const` only when the selected FLCE autograd Function advertises the extended inner-dispatch signature - preserve the standalone CUTEDSL and Ascend Function argument contracts - restore the intended SM100 hardware skip for CUTEDSL FLCE functional-wrapper tests on H200 instead of raising an argument-count `TypeError` ## Root cause `chunk_mem_const` was appended unconditionally by `liger_fused_linear_cross_entropy`. The standalone CUTEDSL Function accepts the original 15 user arguments, so the wrapper supplied one extra positional argument before the Function could reach its SM100 capability gate. ## Testing Run in pod `f45151a224e134137869-n0-0` on H200/SM90: - `make checkstyle` - `make test`: 5067 passed, 915 skipped, 3 xfailed - `make test-cutedsl`: 488 passed, 246 skipped - `make test-cutile`: 211 passed, 113 skipped - `make test-cute`: 4 passed, 10 skipped - `make test-convergence`: 151 passed, 15 skipped, 3 xfailed, 1 xpassed, 8 failed The eight convergence failures are unrelated CUTEDSL LayerNorm failures at `src/liger_kernel/ops/backends/_cutedsl/layer_norm.py:615`, where `view()` rejects non-contiguous multimodal inputs. A targeted rerun against untouched `origin/main` reproduced all eight failures, confirming this change introduces no additional regression. Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
bad6c3ec42 |
Keep Triton FLCE C=1 by default and use C=8 for GPT-OSS (#1472)
## Summary - keep the public Triton fused linear cross entropy chunk-memory default at `C=1` - add an explicit per-call `chunk_mem_const` override while preserving backward compatibility for existing autograd callers - use `C=8` only in the GPT-OSS patched training path, where C=1 causes a deterministic BF16 convergence mismatch - keep every other model and direct/public FLCE caller at C=1 - retain the reusable-logits, direct-dX, and direct-dW optimizations from #1468 - lock the default and override chunk geometry in regression coverage ## Why the override is scoped The exact mini GPT-OSS convergence case fails with C=1 and passes with C=4, C=8, and C=16. However, setting C=8 globally causes the BF16 Qwen3.5-MoE convergence case to fail. The model-scoped override preserves both trajectories: - public/default FLCE: C=1 - GPT-OSS patched loss: C=8 - fused linear KL: unchanged independently tuned policy ## H200 memory and latency BF16, `BT=8192`, `H=4096`, `V=128256`: | Policy | Chunk size | Chunks | Time | Incremental peak allocation | |---|---:|---:|---:|---:| | Public default C=1 | 256 | 32 | 53.039 ms | 1128.66 MiB | | GPT-OSS override C=8 | 2048 | 4 | 37.748 ms | 1568.04 MiB | | Previous global C=16 | 4096 | 2 | 35.074 ms | 2068.04 MiB | The default therefore retains the minimum memory budget, while GPT-OSS uses the smallest requested release override with stable convergence. ## Testing done On H200 with PyTorch 2.9.1.6+cu129 and Triton 3.6.0: - **144 passed**: FLCE ops suite - **144 passed**: FLCE ops suite with `LIGER_KERNEL_STRICT=1` - **141 passed**: transformer-facing FLCE suite - **71 passed, 3 dependency-gated skips**: FP32 + BF16 text-model convergence suites - exact BF16 mini GPT-OSS convergence passes with the scoped C=8 override - BF16 mini Qwen3.5-MoE convergence passes because its default remains C=1 - Ruff lint/format and `git diff --check` pass The broader six-file convergence attempt also exposed four PaliGemma multimodal failures in the ad hoc Transformers 5.15 environment. Those configurations explicitly set `fused_linear_cross_entropy=False`, so they are unrelated to this change. No owned B300 pod was available for an additional Blackwell run. --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
Vaibhav JindalandCopilot
|
de77856a01 |
perf(flce): reuse logits buffers and write input gradients directly (#1468)
## Summary Reduce peak allocated memory in fused linear cross entropy by reusing chunk-logits storage and writing input gradients directly into their destination. On B300 with BF16/FP16, the reference 8,192-token workload uses **33.3% less incremental peak memory**; the 65,536-token workload uses **46.5% less**. Latency is effectively unchanged. This is a buffer-management optimization, not a new kernel or chunk-size policy. ## Details - Allocate one reusable logits buffer on eligible NVIDIA CUDA paths, and use `torch.mm(..., out=...)` to fill each chunk. - Write dX directly into its preallocated slice instead of allocating a GEMM result and copying it. - Release old logits, gradient/transpose views, CE argument tuples, and token-scaling temporaries once consumed. Release reusable workspace before the final gradient-dtype conversion. - Preserve the original allocating matmul and bias-promotion behavior when direct output would change dtype or autocast semantics. The direct-output path requires no-grad execution, matching FP16/BF16/FP32 operands, and disabled or dtype-matching CUDA autocast. Logits reuse additionally requires absent/same-dtype bias; direct dX requires a matching destination dtype. Other cases retain the original matmul behavior. The older autocast-query API remains supported. ## B300 measurements Full forward + backward, mean reduction, no bias, default accumulation, H=4096 and V=128256. Memory is incremental peak allocated MiB above the resident input/parameter baseline, not reserved or total process memory. | Dtype | Tokens | Before ms | After ms | Before peak MiB | After peak MiB | Memory reduction | |---|---:|---:|---:|---:|---:|---:| | BF16 | 8,192 | 19.193 | 19.067 | 3,102.04 | 2,068.04 | 33.3% | | BF16 | 16,384 | 38.301 | 38.420 | 5,202.08 | 3,134.08 | 39.8% | | BF16 | 65,536 | 146.720 | 148.472 | 17,802.31 | 9,530.31 | 46.5% | | FP16 | 8,192 | 16.120 | 16.164 | 3,102.04 | 2,068.04 | 33.3% | | FP16 | 16,384 | 31.688 | 31.556 | 5,202.08 | 3,134.08 | 39.8% | | FP16 | 65,536 | 124.956 | 123.641 | 17,802.31 | 9,530.31 | 46.5% | Across 23 cases covering BF16, FP16, FP32, FP32 accumulation, AMP, and optional features, observed latency changes ranged from **2.2% faster to 1.4% slower**. No material speedup is claimed. The measured losses and gradients matched the baseline with zero numerical difference in these cases. Methodology: - NVIDIA B300 SM103; PyTorch 2.11.0.1+cu130; Triton 3.6.0. - Baseline: FLCE source at `f2ac965f619aa4aa090c34252f87349a968b7597`, unchanged by the two subsequent main commits used as this PR's base. - Baseline and candidate loaded in the same process; identical seeded tensors, chunk sizes, precision, and upstream gradient 1.0. - Fresh forward/backward graph per iteration; both leaf gradients cleared. No optimizer step. - Compilation excluded; eager CUDA-event `triton.testing.do_bench`, 100 ms warmup, 400 ms repetition budget, four alternating-order rounds; median of round medians. - Memory: median of three separate warmed peak-allocation probes. The improvement removes approximately one overlapping logits slab plus a temporary dX chunk. It does **not** make memory independent of token count: the unchanged C=16 policy still grows the chunk size with total tokens. ## Hopper and other GPUs There are no new SM100/SM103-specific instructions or kernels. The optimization uses ordinary PyTorch/cuBLAS output-buffer operations, so the same allocation-lifetime benefit is expected on eligible Hopper and other NVIDIA CUDA workloads. **Only B300 has been measured.** H100/H200/Ampere latency, exact peak allocations, and real older-PyTorch behavior still need hardware coverage; identical performance percentages are not claimed. ROCm and non-CUDA devices do not enter the new direct-output path. The HIP eligibility simulation below is not a ROCm hardware run. ## Testing Done - **322 passed:** transformer FLCE suite, ops FLCE suite, and CuTe shape-fallback contracts. - **170 passed:** ops FLCE suite with `LIGER_KERNEL_STRICT=1` (overlaps the preceding coverage). - **58 passed / 19 NVIDIA-only skips:** targeted tests with only the HIP eligibility flag simulated on B300. - Regression coverage includes actual aten dispatch, storage lifetime through views, ragged/noncontiguous chunks, AMP and bias promotion, FP32 accumulation, token scaling, forward-only execution, and deferred `reduction="none"` backward. - Ruff lint, format checks, and `git diff --check` pass. - Full-repository `make test` and convergence suites were not run. Only the FLCE implementation and its existing regression-test file are included. Local dependency pins, benchmark artifacts, and unrelated files are excluded. Co-authored-by: Copilot <[email protected]> |
||
|
|
f2ac965f61 |
Add SM100 tensor-parallel fused scaled cross entropy (#1458)
## Summary Adds a native Blackwell-family (`sm_100f`, including B200/B300) tensor-parallel fused scaled linear cross-entropy implementation. - Fuses the classifier GEMM with NLL/LSE and optional entropy. - Adds a persistent SM100 backward that executes dZ, dX, dW, token-wave scheduling, node-local NVLS, and two-host reduction in one cluster-2 launch. - Supports TP1/2/4/8 node-local execution and two-host hierarchical TP with 1/2/4/8 selected GPUs per host, including TP subgroups of a larger NVSHMEM world. - Uses the shared forward-verified QP put-signal ring for inter-host communication. - Distributes the remote FP32 dX merge over warp 1 from every resident CTA, synchronized with an HBM atomic completion counter. - Adds reusable workspace configuration/query APIs, public API documentation, correctness validators, phase benchmarks, and topology contract tests. ## Code entry points | Layer | Forward | Backward | |---|---|---| | PyTorch autograd | [`LigerFusedLinearScaledCrossEntropyNativeTPFunction.forward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/src/liger_kernel/ops/cute/fused_linear_scaled_cross_entropy_tp.py#L139) | [`LigerFusedLinearScaledCrossEntropyNativeTPFunction.backward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/src/liger_kernel/ops/cute/fused_linear_scaled_cross_entropy_tp.py#L186) | | TVM FFI | [`fused_linear_scaled_cross_entropy_forward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp#L535) | [`fused_linear_scaled_cross_entropy_backward`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/liger_cute_kernels/tvm_ffi_bindings.cpp#L610) | | SM100 host launcher | [`fused_linear_scaled_cross_entropy_forward_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_forward_sm100.cu#L703) | [`fused_linear_scaled_cross_entropy_backward_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/fused_linear_scaled_cross_entropy_backward_sm100.cu#L445) | | Device kernel | [`forward_gemm_tp_kernel_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_kernel_sm100.cuh#L26) | [`backward_gemm_tp_kernel_sm100`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm100.cuh#L1109) | | Role implementations | [`forward_gemm_roles_sm100.cuh`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/forward_gemm_roles_sm100.cuh) | [`backward_gemm_mainloop_sm100.cuh`](https://github.com/kolehma8/Liger-Kernel/blob/jkolehm/tp-fslce-sm100/liger_cute_kernels/csrc/core/src/fused_scaled_linear_cross_entropy/backward_gemm_mainloop_sm100.cuh) | ## Multi-host subgroup topology Multi-host setup is relative to the configured TP team rather than the full NVSHMEM world. `nvshmem_team_split_2d` derives: - one local NVLS row containing only the selected GPUs on the current host; - one matching-rank remote column containing the corresponding selected GPU from each host. For a 16-rank world on two eight-GPU hosts, TP group `{0,4,8,12}` becomes local rows `{0,4}` / `{8,12}` and remote pairs `{0,8}` / `{4,12}`. SM100 remote kernels use a parent-team stream barrier followed by a cooperative clustered CUDA launch, so disjoint TP groups do not require world-synchronous kernel entry. Validated concurrently as four TP4 groups `{0,4,8,12}`, `{1,5,9,13}`, `{2,6,10,14}`, and `{3,7,11,15}`: - 28/28 numerical cases pass across the four groups. - Covers entropy, ignored targets, ragged vocabulary, full waves, and two-wave tails. - Cross-rank dX equality remains exact. - Deliberately staggering group entry by 500 ms also passes for forward and backward, confirming subgroup-scoped progress. - Two-host local widths 1, 2, 4, and 8 GPUs per host pass. - `NVSHMEM_IB_ENABLE_IBGDA=1` was forced for the final campaign. Runtime initialization reported `Successfully initialized the transport: IBGDA. It will be used for device-side APIs over IB`, created DCI/RC QPs, and selected each configured `mlx5_*` device. Random BF16, `M=4096`, `H=4096`, global `V=131072`, 10 warmups, and 30 measured iterations: | TP group root | Native forward ms | Native forward TF/GPU | Forward speedup | Native F+B ms | Native F+B TF/GPU | F+B speedup | |---:|---:|---:|---:|---:|---:|---:| | 0 | 1.080 | 1017.9 | 2.82x | 5.193 | 846.9 | 1.69x | | 1 | 1.072 | 1025.8 | 2.85x | 5.166 | 851.3 | 1.69x | | 2 | 1.058 | 1039.5 | 2.86x | 5.128 | 857.6 | 1.69x | | 3 | 1.090 | 1009.2 | 2.78x | 5.175 | 849.9 | 1.68x | Native versus fallback forward maximum absolute difference was `5.45e-4`. ### SM90 regression validation Validated the shared topology, workspace, and binding changes on two eight-GPU H200 hosts with a clean `sm_90a` build: - Complete native core build passed. - 12/12 shared reduction and parent-relative topology tests passed. - 8/8 SM90 FSLCE configuration and online-softmax contract tests passed. - 12/12 standalone SM90 forward GEMM numerical cases passed. - Public FFI world-team and process-group subteam tests passed, including forward, dX, dW, entropy, and CUDA graph replay. - A two-host TP16 world team passed forward, dX, dW, and entropy validation. - Four concurrent strided TP4 groups `{0,4,8,12}`, `{1,5,9,13}`, `{2,6,10,14}`, and `{3,7,11,15}` all passed on H200. - NVSHMEM reported IBGDA as the active device-side InfiniBand transport on all 16 ranks. ## Implementation ### Forward - 2SM UMMA, `M256 x N256 x K64`, cluster `2 x 1`. - Five-stage TMA mainloop. - Four source slots and a default wave width of 64 N256 vocabulary tiles. - Warp 0 performs local reduction, warp 1 runs the matching-rank host ring, warp 2 owns TMA, warp 3 issues UMMA, and warps 4-11 drain TMEM/epilogue work. - The `__global__` entry point is isolated in `forward_gemm_kernel_sm100.cuh` and dispatches those roles explicitly; templated implementations live in `forward_gemm_roles_sm100.cuh`. ### Backward - One persistent 384-thread cluster-2 kernel. - dZ: five-stage `M256 x N256 x K64`. - dX: four-stage 2SM N512 schedule, publishing two M128xN256 communication tiles. - dW: six-stage paired-N schedule with TMA store for wave 0 and TMA reduce-add for later waves. - A single 512-column `Allocator2Sm` TMEM allocation is reused across all phases. - TP16: 1. every CTA's warp 0 performs node-local NVLS reduce-scatter; 2. CTA 0 warp 1 performs the existing inter-host ring transfer; 3. warp 1 from all resident CTAs merges a global-strided portion of the remote FP32 shard; 4. one HBM atomic arrival per CTA releases CTA 0; 5. warp 0 performs node-local all-gather and BF16 scatter. The measured TP16 remote stage is 0.285-0.328 ms: | Component | Latency | |---|---:| | IBRC transfer, 8 MiB/GPU | 0.119-0.156 ms | | All-CTA FP32 merge | 0.160-0.176 ms | | Total remote stage | 0.285-0.328 ms | | In-kernel dW overlap window | ~2.41 ms | | Exposed remote tail | 0 ms | ## Correctness - TP1/2/4/8/16: 35/35 distributed numerical cases pass. - Covers aligned and ragged shapes, entropy, ignored targets, 4K waves, 4097-token boundaries, and 5000-token two-wave tails. - Cross-rank dX equality is exact. - Maximum observed relative error: - dX: 0.262% - dW: 0.498% - Forward native versus Verl maximum absolute difference: `1.09e-3`. - C++ tests: - shared reduction/topology: 11/11 - SM100 backward contracts: 12/12 - `make checkstyle` passes. ## Performance methodology - GPU: NVIDIA B300. - BF16 inputs, FP32 accumulation. - Hidden size: 4096. - Fixed global vocabulary: 131072. - Tokens: 4096, 8192, 16384. - TP: 1, 2, 4, 8, 16. - 10 warmups, 30 measured iterations, median CUDA-event latency. - Verl baseline: repository 512-token chunked TP fallback. - Forward+backward throughput counts four classifier GEMMs: forward logits, backward logits recomputation for dZ, dX, and dW. - Efficiency is per-GPU TFLOP/s relative to TP1 for the same shape. - TP16 uses two eight-GPU hosts and explicit one-HCA-per-local-GPU mapping. ### M=4096 forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 2.852 | 1542.0 | 100.0% | 8.489 | 518.1 | 100.0% | 2.98x | | 2 | 1.457 | 1508.9 | 97.9% | 4.653 | 472.6 | 91.2% | 3.19x | | 4 | 0.764 | 1438.3 | 93.3% | 2.856 | 384.9 | 74.3% | 3.74x | | 8 | 0.428 | 1283.6 | 83.2% | 2.134 | 257.6 | 49.7% | 4.98x | | 16 | 0.465 | 591.0 | 38.3% | 2.096 | 131.1 | 25.3% | 4.51x | ### M=4096 forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 10.950 | 1606.5 | 100.0% | 29.898 | 588.4 | 100.0% | 2.73x | | 2 | 6.170 | 1425.6 | 88.7% | 15.570 | 564.9 | 96.0% | 2.52x | | 4 | 3.562 | 1234.7 | 76.9% | 8.451 | 520.4 | 88.5% | 2.37x | | 8 | 2.253 | 976.0 | 60.8% | 5.426 | 405.2 | 68.9% | 2.41x | | 16 | 2.171 | 506.5 | 31.5% | 4.472 | 245.9 | 41.8% | 2.06x | ### M=8192 forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 5.744 | 1531.4 | 100.0% | 17.107 | 514.2 | 100.0% | 2.98x | | 2 | 2.867 | 1533.8 | 100.2% | 9.376 | 469.1 | 91.2% | 3.27x | | 4 | 1.469 | 1496.9 | 97.7% | 5.638 | 390.1 | 75.9% | 3.84x | | 8 | 0.793 | 1385.8 | 90.5% | 4.185 | 262.7 | 51.1% | 5.28x | | 16 | 0.762 | 721.6 | 47.1% | 4.077 | 134.8 | 26.2% | 5.35x | ### M=8192 forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 22.264 | 1580.3 | 100.0% | 58.375 | 602.8 | 100.0% | 2.62x | | 2 | 11.600 | 1516.5 | 96.0% | 31.618 | 556.4 | 92.3% | 2.73x | | 4 | 6.463 | 1360.9 | 86.1% | 17.348 | 507.1 | 84.1% | 2.68x | | 8 | 4.219 | 1042.5 | 66.0% | 10.653 | 412.8 | 68.5% | 2.52x | | 16 | 4.020 | 546.9 | 34.6% | 8.659 | 254.0 | 42.1% | 2.15x | ### M=16384 forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 12.645 | 1391.2 | 100.0% | 35.020 | 502.4 | 100.0% | 2.77x | | 2 | 6.399 | 1374.5 | 98.8% | 19.560 | 449.7 | 89.5% | 3.06x | | 4 | 3.206 | 1371.9 | 98.6% | 11.426 | 384.9 | 76.6% | 3.56x | | 8 | 1.674 | 1313.3 | 94.4% | 8.197 | 268.3 | 53.4% | 4.90x | | 16 | 1.449 | 759.0 | 54.6% | 9.809 | 112.1 | 22.3% | 6.77x | ### M=16384 forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 50.387 | 1396.5 | 100.0% | 115.214 | 610.8 | 100.0% | 2.29x | | 2 | 23.270 | 1512.0 | 108.3% | 62.505 | 562.9 | 92.2% | 2.69x | | 4 | 12.579 | 1398.5 | 100.1% | 35.498 | 495.6 | 81.1% | 2.82x | | 8 | 8.281 | 1062.1 | 76.1% | 21.188 | 415.2 | 68.0% | 2.56x | | 16 | 7.898 | 556.8 | 39.9% | 16.706 | 263.2 | 43.1% | 2.12x | ## Memory Inputs, weights, and final outputs are excluded. Native reports its exact reusable C++/NVSHMEM forward+backward workspace. Verl reports incremental PyTorch peak memory. ### M=4096 | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 1810.40 | 1024.10 | 2720.11 | | 2 | 1202.42 | 512.10 | 1376.11 | | 4 | 898.46 | 256.10 | 704.11 | | 8 | 746.53 | 128.10 | 368.11 | | 16 | 682.67 | 64.10 | 200.11 | ### M=8192 | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2004.79 | 1024.19 | 2752.21 | | 2 | 1300.81 | 512.19 | 1408.21 | | 4 | 948.84 | 256.19 | 736.21 | | 8 | 772.91 | 128.19 | 400.21 | | 16 | 709.06 | 64.19 | 232.21 | ### M=16384 | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2393.56 | 1024.35 | 2816.40 | | 2 | 1497.58 | 512.35 | 1472.40 | | 4 | 1049.61 | 256.35 | 800.40 | | 8 | 825.68 | 128.35 | 464.40 | | 16 | 761.83 | 64.35 | 296.40 | ## Additional weak-scaling production result With fixed `V_local=65536`, `M=4096`, and `H=4096`: | TP | Native forward TF/GPU | Verl forward TF/GPU | Native F+B TF/GPU | Verl F+B TF/GPU | |---:|---:|---:|---:|---:| | 1 | 1524.5 | 488.6 | 1607.1 | 570.3 | | 2 | 1505.3 | 471.8 | 1424.4 | 560.7 | | 4 | 1513.3 | 459.2 | 1541.9 | 543.2 | | 8 | 1490.8 | 446.8 | 1580.4 | 537.9 | | 16 | 1323.0 | 412.2 | 1434.5 | 533.7 | ## Model-shape campaign The following production model shapes use `M=8192`, random BF16 inputs, 10 warmups, and 30 measured iterations. TP16 spans the same two eight-GPU hosts as the fixed-shape campaign. ### Qwen3-30B-A3B (`H=2048`, `V=151936`) #### Forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 3.692 | 1380.8 | 100.0% | 15.885 | 320.9 | 100.0% | 4.30x | | 2 | 1.937 | 1315.9 | 95.3% | 9.192 | 277.3 | 86.4% | 4.75x | | 4 | 1.040 | 1225.4 | 88.7% | 5.581 | 228.4 | 71.2% | 5.37x | | 8 | 0.592 | 1076.2 | 77.9% | 4.136 | 154.1 | 48.0% | 6.99x | | 16 | 0.664 | 479.8 | 34.7% | 4.012 | 79.4 | 24.7% | 6.04x | #### Forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 13.692 | 1489.3 | 100.0% | 47.689 | 427.6 | 100.0% | 3.48x | | 2 | 7.480 | 1363.2 | 91.5% | 26.363 | 386.8 | 90.4% | 3.52x | | 4 | 4.100 | 1243.3 | 83.5% | 14.834 | 343.7 | 80.4% | 3.62x | | 8 | 2.507 | 1016.5 | 68.3% | 9.720 | 262.3 | 61.3% | 3.88x | | 16 | 2.545 | 500.8 | 33.6% | 8.346 | 152.7 | 35.7% | 3.28x | #### Memory | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 1975.79 | 1187.19 | 1962.08 | | 2 | 1286.31 | 593.69 | 996.65 | | 4 | 941.84 | 296.94 | 515.33 | | 8 | 769.41 | 152.19 | 276.94 | | 16 | 695.56 | 74.37 | 154.58 | ### Llama 4 (`H=5120`, `V=202048`) #### Forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 11.727 | 1445.3 | 100.0% | 28.506 | 594.6 | 100.0% | 2.43x | | 2 | 5.795 | 1462.3 | 101.2% | 15.228 | 556.5 | 93.6% | 2.63x | | 4 | 2.880 | 1471.2 | 101.8% | 8.489 | 499.1 | 83.9% | 2.95x | | 8 | 1.479 | 1433.0 | 99.1% | 5.339 | 396.8 | 66.7% | 3.61x | | 16 | 1.103 | 960.3 | 66.4% | 7.578 | 139.8 | 23.5% | 6.87x | #### Forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 50.180 | 1351.1 | 100.0% | 100.330 | 675.7 | 100.0% | 2.00x | | 2 | 22.677 | 1494.8 | 110.6% | 54.966 | 616.7 | 91.3% | 2.42x | | 4 | 11.785 | 1438.3 | 106.4% | 29.663 | 571.5 | 84.6% | 2.52x | | 8 | 6.939 | 1221.3 | 90.4% | 16.230 | 522.1 | 77.3% | 2.34x | | 16 | 6.026 | 703.2 | 52.0% | 23.633 | 179.3 | 26.5% | 3.92x | #### Memory | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2655.29 | 1578.69 | 5015.46 | | 2 | 1626.31 | 792.19 | 2547.99 | | 4 | 1111.84 | 394.81 | 1315.52 | | 8 | 854.41 | 200.19 | 698.15 | | 16 | 756.06 | 98.84 | 389.87 | ### DeepSeek-V3 (`H=7168`, `V=129280`) #### Forward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 10.923 | 1390.0 | 100.0% | 22.352 | 679.2 | 100.0% | 2.05x | | 2 | 5.317 | 1427.6 | 102.7% | 11.923 | 636.7 | 93.7% | 2.24x | | 4 | 2.638 | 1439.0 | 103.5% | 6.943 | 546.7 | 80.5% | 2.63x | | 8 | 1.356 | 1399.8 | 100.7% | 4.556 | 416.5 | 61.3% | 3.36x | | 16 | 1.041 | 911.7 | 65.6% | 4.862 | 195.2 | 28.7% | 4.67x | #### Forward + backward | TP | Native ms | Native TF/GPU | Efficiency | Verl ms | Verl TF/GPU | Efficiency | Speedup | |---:|---:|---:|---:|---:|---:|---:|---:| | 1 | 43.588 | 1393.3 | 100.0% | 81.751 | 742.9 | 100.0% | 1.88x | | 2 | 22.167 | 1369.9 | 98.3% | 44.494 | 682.4 | 91.9% | 2.01x | | 4 | 12.411 | 1223.3 | 87.8% | 24.582 | 617.6 | 83.1% | 1.98x | | 8 | 7.600 | 998.8 | 71.7% | 13.475 | 563.3 | 75.8% | 1.77x | | 16 | 7.024 | 540.4 | 38.8% | 10.843 | 350.0 | 47.1% | 1.54x | #### Memory | TP | Native workspace MiB | Verl forward peak MiB | Verl F+B peak MiB | |---:|---:|---:|---:| | 1 | 2278.79 | 1010.19 | 4279.46 | | 2 | 1437.81 | 505.19 | 2196.71 | | 4 | 1017.34 | 256.19 | 1156.21 | | 8 | 807.41 | 128.19 | 634.08 | | 16 | 744.56 | 64.19 | 373.04 | ## Validation commands ```bash make checkstyle ``` On B300 (`-DLIGER_CUTE_CUDA_ARCH=100f`): ```bash cmake --build build --target \ libliger_cute_kernels.so \ test_fslce_shared_reduction \ test_fused_scaled_linear_cross_entropy_backward_sm100 ./build/tests/cpp/test_fslce_shared_reduction ./build/tests/cpp/test_fused_scaled_linear_cross_entropy_backward_sm100 ``` Distributed numerical validation: ```bash python -m torch.distributed.run --standalone --nproc_per_node=<1|2|4|8> \ liger_cute_kernels/tests/python/validate_fused_scaled_linear_cross_entropy_backward_sm100.py ``` TP16 uses two eight-GPU hosts with: ```bash export NVSHMEM_REMOTE_TRANSPORT=ibrc export NVSHMEM_IB_ENABLE_IBGDA=1 export NVSHMEM_ENABLE_NIC_PE_MAPPING=1 unset NVSHMEM_HCA_LIST export NVSHMEM_HCA_PE_MAPPING='mlx5_0:1:1,mlx5_2:1:1,mlx5_3:1:1,mlx5_4:1:1,mlx5_5:1:1,mlx5_6:1:1,mlx5_8:1:1,mlx5_9:1:1' ``` --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
sunyi0505
|
bcb5962dc0 |
[NPU] Fix Ascend JSD UT failures and register LayerNorm dispatch (#1452)
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
Unit tests for JSD / FusedLinearJSD on Ascend failed, mainly due to
signature misalignment with the NVIDIA frontend, and invalid kernel
launches when all samples are ignored. The ascend-triton dispatcher
registration for LayerNorm is also added so that the NPU can
automatically select the UB-aware implementation.
JSD / FusedLinearJSD signature alignment: NVIDIA's LigerJSDFunction /
LigerFusedLinearJSDFunction accepts jsd_impl / jsd_mode (apply(...,
None, None) is used in tests). The Ascend kernel is self-contained and
does not use the inner dispatcher, so it simply accepts and ignores
these two parameters, and the number of backward return values is
supplemented accordingly; supports_inner_impl_dispatch is set to False
for JSD.
Short-circuit for all-ignored case: When n_non_ignore == 0, skip kernel
launch and directly return zero loss and zero gradients. This covers
cases such as test_correctness_all_ignored to prevent division-by-zero /
invalid grid.
LayerNorm registration: Add ops/backends/_ascend/layer_norm.py and
register it in declare_op_locations("layer_norm") within functional.py.
Otherwise, auto-selection only detects the NVIDIA-triton implementation
for CUDA and cannot use _ascend.ops.layer_norm.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->
## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->
pytest tests/transformers
<img width="1255" height="383" alt="aa9edfdb912726cb88d8905e1b40f870"
src="https://github.com/user-attachments/assets/db596c84-8a94-480c-82a9-2a98351b4510"
/>
<!--
Replace BLANK with your device type. For example, A100-80G-PCIe
Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them.
-->
- Hardware Type: Atlas 800T A3(X86)
- [x] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [x] run `make test-convergence` to ensure convergence
|
||
|
Yulong Cao
|
f6385e909b |
feat: add FusedLinearKLDivLoss (fused linear + KL divergence for distillation) (#1423)
## Summary
Add `FusedLinearKLDivLoss`, which fuses the vocabulary-head linear
projection with
KL(Q || P) loss against an explicit target distribution, for knowledge
distillation
workflows where teacher probabilities are available directly (e.g.
cached / offline
soft labels), complementing `LigerKLDIVLoss` (element-wise, no fusion)
and
`LigerFusedLinearJSD` (teacher side needs its own projection).
The kernel follows the chunked gradient-in-forward structure of
`fused_linear_jsd.py`,
so the full `BT x V` logits tensor is never materialized, and adopts the
`_CHUNK_MEM_CONST = 16` token-chunk geometry of
`fused_linear_cross_entropy.py`.
## Details
- Supports `reduction` (`"batchmean"` / `"mean"` / `"sum"`, same
semantics as
`torch.nn.KLDivLoss`), `ignore_index` via `shift_labels`, softmax
`temperature`,
and an `eps` clamp so that `0 * log(0)` contributes 0.
- Supports `accum_dtype` for fp32 weight-gradient accumulation across
chunked
low-precision GEMMs (same semantics as `LigerFusedLinearJSD` /
`LigerFusedLinearCrossEntropyLoss`), including the
`torch.addmm(out_dtype=fp32)`
fast path on torch >= 2.8 / sm80+ and a fallback elsewhere.
- The log-softmax backward is computed in place on the saved log-probs
and the
`chunk_size x V` fp32 temporaries are freed eagerly, so the peak during
the
trailing GEMMs stays low.
- `reduction="none"` is intentionally unsupported (gradient-in-forward
requires a
scalar loss); a `ValueError` points users to `LigerKLDIVLoss` instead.
- bf16 + `reduction="sum"` test tolerances follow the existing
convention in
`test_fused_linear_cross_entropy.py`.
- Peak-memory improvement over the unfused PyTorch baseline (full
fwd+bwd,
`max_memory_allocated`, bf16, llama-3 head shape H=4096 / V=128256, RTX
4060
Laptop GPU): BT=2048 6537 -> 4310 MB (1.52x), BT=4096 12072 -> 6594 MB
(1.83x);
at BT=1024 the benchmark harness reports 4793 -> 4244 MB. Runtime is at
parity
with the baseline at BT=1024 (0.89x).
## Testing Done
- Hardware Type: NVIDIA GeForce RTX 4060 Laptop GPU (sm89); NVIDIA Tesla
T4 (sm75)
- [x] run `make test` to ensure correctness (32 pre-existing failures
unrelated to
this change: 18 `test_mlp.py` shared-memory OOR on sm89 and 14
`test_grpo_loss.py`
LUSPO numerics; both reproduce on `main` without this change)
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence (does not
exercise the new
standalone loss)
<details>
<summary>Test logs (RTX 4060 Laptop, torch 2.11.0+cu128, triton
3.7.1)</summary>
```
make test: 32 failed, 4152 passed, 914 skipped, 15 xfailed in 1050.50s (0:17:30)
test/transformers/test_fused_linear_kl_div.py: 172 passed, 0 failed, 0 skipped
failures (all pre-existing on this GPU, reproduced on main without this change):
18 x test/transformers/test_mlp.py
triton.runtime.errors.OutOfResources: shared memory
Required: 131072, Hardware limit: 101376 (sm89 has 100 KB usable smem)
14 x test/chunked_loss/test_grpo_loss.py
chunked GRPO LUSPO numerics (also seen by other contributors, see #1397)
make checkstyle: ruff check and ruff format --check pass on all files in this PR
```
</details>
<details>
<summary>Test logs (Kaggle Tesla T4, torch 2.10.0+cu128, triton
3.6.0)</summary>
```
test/transformers/test_fused_linear_kl_div.py: 51 passed, 25 skipped in 10.29s
(bf16 cases skipped by design: T4 is sm75 without native bf16 support)
```
</details>
|
||
|
xvxuopop
|
c64878355b |
Add DTensor support for Ascend SwiGLU (#1448)
## Summary Add DTensor support to the Ascend SwiGLU implementation. SwiGLU now runs on local DTensor shards and preserves the original device mesh and placements in forward and backward. Also fix `infer_comm_backend()` to return `hccl` on Ascend NPU. ## Testing Done Environment: 4 x Ascend 910 NPU, CANN 9.1.0, PyTorch/torch-npu 2.9.0. - Non-DTensor SwiGLU tests: `67 passed, 18 skipped` - Four-device FP32/BF16 SwiGLU DTensor tests: passed - Four-device RMSNorm DTensor test for HCCL regression: passed - `make checkstyle`: passed - Hardware Type: 4 x Ascend 910 NPU - [ ] run `make test` to ensure correctness (targeted tests run instead) - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence (not applicable to this DTensor change) |
||
|
Vaibhav JindalandCopilot
|
95b01e9402 |
perf(flce): accumulate low-precision Triton dW with addmm (#1454)
## Summary Optimize Liger's Triton fused linear cross entropy (FLCE) weight-gradient accumulation for BF16/FP16 parameters on NVIDIA SM80+ GPUs. The previous default-accumulator path computes a parameter-sized matrix, converts it to FP32, then adds it into the BF16/FP16 running gradient buffer for **every token chunk**: ```python grad_weight += torch.mm(grad_logits_chunk.t(), input_chunk).float() ``` This PR instead accumulates directly into the existing dtype-matched gradient buffer: ```python torch.addmm(grad_weight, grad_logits_t, input_chunk, out=grad_weight) ``` This removes the separate parameter-sized product/conversion temporary and extra addition pass. The stored gradient-buffer dtype is unchanged; the old `.float()` temporary was not an FP32 running accumulator. The fused accumulation can change rounding relative to the old standalone product, so bitwise equivalence is not claimed. The existing guarded FP32 accumulation fast path, AMP operand alignment, and legacy dtype/device fallbacks are preserved. The CE kernel, its launch configuration, the default token-chunk heuristic, and the forward/backward scheduling are unchanged. **Scope:** this is the first stack layer only: Triton accumulation plus regressions in the existing FLCE test file. It does **not** add a public `chunk_size` argument, change cuTile, or include the later three-way benchmark implementation. ## Performance improvements attributable to this PR ### Unchanged production defaults BF16 inputs/weights, `accum_dtype=None`, and the existing `_CHUNK_MEM_CONST=16` heuristic on both sides. The number of token chunks is identical before and after. At the Llama-3-sized reference shape, this is **1.26x faster** with **38.9% lower incremental peak allocated memory** without changing any user settings. | Tokens / hidden / vocab | Chunk rows | Before ms | After ms | Speedup | Before / after peak MiB | Memory reduction | |---|---:|---:|---:|---:|---:|---:| | 8,192 / 4,096 / 128,256 | 4,096 | 25.002 | 19.897 | **1.26x** | 5,074.04 / 3,102.04 | **38.9%** | | 4,096 / 2,048 / 32,000 | 4,096 | 1.774 | 1.366 | **1.30x** | 766.02 / 407.02 | **46.9%** | ### Controlled smaller chunks: benefit per accumulation step These are additional controlled experiments, **not the production defaults**. For the reference shape, the existing private chunk-budget constant was set to 4 (1,024-row chunks) or 1 (256-row chunks) on **both** versions. No new chunk-size API or later stack-layer code was needed. | Tokens / hidden / vocab | Chunk rows | Before ms | After ms | Speedup | Before / after peak MiB | Memory reduction | |---|---:|---:|---:|---:|---:|---:| | 8,192 / 4,096 / 128,256 | 1,024 | 43.167 | 21.836 | **1.98x** | 4,322.54 / 1,575.04 | **63.6%** | | 8,192 / 4,096 / 128,256 | 256 | 123.443 | 31.686 | **3.90x** | 4,134.66 / 1,193.29 | **71.1%** | The larger gain at smaller chunks is expected: the old implementation repeatedly allocated/converted/added a full `[vocab, hidden]` matrix, and this PR removes that work on each iteration. At the reference shape, one such FP32 matrix alone is approximately 2,004 MiB. ### FP32-accumulator control The existing `accum_dtype=torch.float32` fast path is intentionally preserved. At 8,192 tokens / hidden 4,096 / vocab 128,256, with the unchanged default 4,096-row chunk: - Before: **20.473 ms**, **4,104.04 MiB** incremental peak. - After: **20.471 ms**, **4,104.04 MiB** incremental peak. This control is effectively unchanged; the improvement targets the default low-precision gradient-buffer path. ### Measurement methodology - **Hardware:** NVIDIA B200, SM100, driver 580.105.08. - **Software:** PyTorch 2.11.0.1+cu130, Triton 3.6.0. - **Before / after source:** base FLCE implementation `96233e3` versus `6a5077f`, loaded against the same installed dependencies. Current head `d5f3593` changes only a test comment relative to `6a5077f`; its kernel implementation is identical. - Identical seeded BF16 inputs, weights, targets, and token-chunk geometry; both inputs and weights require gradients; no bias or optional loss features in the timing cases. - Measure **full forward + backward**, mean reduction, fixed upstream gradient 1.0. Each iteration uses a fresh graph and clears both leaf gradients. No optimizer step or model-training-throughput claim. - Compilation plus five warmup iterations excluded. Each run uses 30 CUDA-event timed iterations; tables report the median of three runs, alternating before/after order in the same process. - Memory is **incremental peak allocated MiB above the live input/weight baseline**, with both gradients cleared before reset and a dedicated memory-probe iteration. It is not total process memory or reserved memory. - These compare the same Triton-backed implementation before/after the portable cuBLAS accumulation improvement. They are not CuTe-versus-Triton or DSL/compiler speedup claims. ## Testing Done - **Hardware:** NVIDIA B200 / SM100. - Latest relevant suites: **470 passed, 1 environment-gated skipped** across `test/ops/test_fused_linear_cross_entropy.py`, `test/transformers/test_fused_linear_cross_entropy.py`, and `test/cutedsl/test_fused_linear_cross_entropy.py`. - All **70 addmm regression cases** now live in the existing ops FLCE test file alongside its existing 23 cases. Their chunk-budget fixture is class-scoped and does not alter the pre-existing tests. - Coverage includes real aten dispatch checks for direct `addmm.out` and no parameter-shaped upcasts, preserved FP32 `dtype_out`, AMP and FP32-bias promotion, BF16/FP16/FP32, mean/sum/nonuniform per-token gradients, ignored chunks, frozen weights, noncontiguous tensors, and old-Torch/SM capability fallbacks. - Review requests addressed: restored the original inline FP32 guard, removed the extra boolean aliases, consolidated tests, and renamed the marker to `requires_sm80_or_higher`. - [x] `make checkstyle` passes locally; GitHub Checkstyle is also green. - [ ] Full-repository `make test` was not run; the targeted suites above were run. - [ ] `make test-convergence` was not run; no convergence claim is made. --------- Co-authored-by: Copilot <[email protected]> |