mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
23ba1065a70ff8cd445a2077d607b17ac11df7cb
979
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Vaibhav Jindal
|
23ba1065a7 |
Update README.md (#1522)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## 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. ---> <!-- 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: <BLANK> - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
Vaibhav JindalandCopilot
|
d5f281789a |
fix(test): stop passing removed device kwarg to Llama4 rotary embed… (#1516)
## Summary
`test/transformers/test_llama4_rope.py` constructs the rotary embedding
directly as
`Llama4TextRotaryEmbedding(config=config, device=device)`. Newer
HuggingFace `transformers`
deprecated the `device` kwarg on config-based
`*RotaryEmbedding.__init__`
(`@deprecate_kwarg("device", version="5.18")`) and **removes it at
5.18**. Because CI installs
the latest `transformers` (`setup.py` pins `transformers>=4.52.0`, no
upper bound), every test
in the file fails at setup with:
```
TypeError: Llama4TextRotaryEmbedding.__init__() got an unexpected keyword argument 'device'
```
This fixes all of the currently-failing `test_llama4_rope.py` cases by
dropping the removed kwarg.
## Why dropping `device=` is correct (not just silencing)
`Llama4TextRotaryEmbedding.forward` relocates the frequency buffer to
the input tensor's device:
```python
inv_freq_expanded = self.inv_freq[...].to(x.device)
```
and the test tensors are already created on the target device, so the
rotary module does not
need to be constructed on a specific device. Verified on
`transformers==5.15.0` that `cos`/`sin`
are **bitwise identical** (`torch.equal`) to the previous
`device=`-based construction.
## Relationship to #1514
#1514 fixed the same `device`-kwarg drift, but only for call sites
routed through
`transformers_version_dispatch`. `test_llama4_rope.py` constructs the
module **directly**, so it
was not covered. This PR is the direct-construction counterpart.
## Testing
- `pytest test/transformers/test_llama4_rope.py` → **19 passed**
(transformers 5.15.0, H100).
- Simulated transformers ≥5.18 (strict signature): the old `device=`
call raises the exact CI
`TypeError`; the new call succeeds.
- `ruff check` / `ruff format --check` clean.
## Compatibility
| transformers | Before this PR | After this PR |
|---|---|---|
| < 4.48.0 | n/a (test path unchanged) | n/a |
| 4.48.0 – <5.18 | `device=` accepted (deprecation warning) | no
`device=`, identical result, no warning |
| ≥ 5.18 | **`TypeError` at construction** | works |
Co-authored-by: Copilot <[email protected]>
|
||
|
Vaibhav JindalandCopilot
|
98f6371a9c |
ci: run only Checkstyle in the merge queue; make AMD/Intel manual-only (#1515)
## Summary The AMD and Intel GPU runners are currently unavailable, so those CI jobs fail whenever they run — including inside the **merge queue**, which blocks all merges for some time. This PR makes the merge queue run **only Checkstyle** for now. ## Change Reduce the **AMD GPU** and **Intel GPU** workflow triggers to `workflow_dispatch` only, mirroring the **NVIDIA GPU** workflow's `on:` (manual trigger) **minus the schedule** — we don't maintain a schedule for AMD/Intel. This removes their `push`, `pull_request`, and `merge_group` triggers, so: - **Only Checkstyle runs in the merge queue** (it keeps its `merge_group` trigger). - AMD/Intel no longer run automatically on PRs or pushes. - AMD/Intel can still be launched on demand via the Actions **"Run workflow"** button once the runners are healthy again. `checkstyle.yml` is unchanged. ## Notes - Observed (not changed here): `merge-queue-failure.yml` listens for the **NVIDIA GPU** workflow with `event == 'merge_group'`, but NVIDIA GPU has no `merge_group` trigger, so that handler is currently a no-op. Can be addressed separately. Co-authored-by: Copilot <[email protected]> |
||
|
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]> |
||
|
Vaibhav JindalandCopilot
|
280444a998 |
fix(test): use even head_dim in rope functional-correctness weird shape (#1513)
Transfomers change: https://github.com/huggingface/transformers/pull/48524, leads to failing tests because of odd head_dim. ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> HuggingFace's newer config validation (validate_rope) requires an even rotary dimension, so constructing LlamaConfig(head_dim=41) now raises StrictDataclassClassValidationError. test_functional_correctness builds that config in its fixture, so the weird-shape case [*-9-7-41-41-41] fails during setup (before the kernel runs) on transformers versions that enforce the validator -- this is what breaks the daily H100 correctness CI. Odd head_dim is unsupported across the stack: the CuTeDSL backend rejects it (test_rope_cutedsl_rejects_odd_head_dimension) and the sibling test_correctness block in this same file already documents this and uses an even head_dim. Switch the weird shape from head_dim=41 to head_dim=92, matching that convention while preserving non-power-of-2 masking coverage (bsz=9, seq_len=7, 41 q/kv heads). <!--- ## 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. ---> <!-- 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: <BLANK> - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Co-authored-by: Copilot <[email protected]> |
||
|
Arup DeandCopilot
|
2024b7d83b |
test(cutedsl): gap-op dispatcher tests + fused_add_rms_norm 3-D backward + shape fallbacks (#1457)
## Summary Add CuTe DSL correctness coverage in three test files. Runtime source is unchanged. - Compare supported native operations with both Triton and independent PyTorch references, checking every expected output and gradient. - Cover 2-D/3-D fused-add RMSNorm, including normalization-only, residual-only, and combined backward paths with independent gradient storage. - Reject missing, detached, truncated, or vacuous gradient comparisons. Check JSD endpoint, ignore-mask, aliasing, and global-denominator contracts. - Distinguish native execution from explicit Triton fallbacks, including separate alignment and size-limit cases. No benchmark results or speedup claims are included. ## Testing Done Exact commit `425650d645462aa804ee0e5d42d3559c70a3a3e5` on NVIDIA B300, using an exclusive GPU lease: `python -m pytest test/cutedsl/test_fused_add_rms_norm.py test/cutedsl/test_gap_ops.py test/cutedsl/test_shape_fallbacks.py` **86 passed, 0 failed, 0 skipped:** 60 native CuTe DSL cases, 17 explicitly asserted Triton-fallback cases, and 9 CPU helper/reference cases. For fused-linear JSD, native attribution applies to the inner JSD primitive, not the shared projection or backward-scaling operations. Changed-file Ruff lint and formatting pass. This is targeted correctness coverage, not full model convergence or cross-platform qualification. Hopper, AMD, and Intel execution remain unqualified by these B300 results. Maintainer approval and repository CI gates remain separate requirements. Co-authored-by: Copilot <[email protected]> --------- Co-authored-by: Copilot <[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]> |
||
|
Fuxiang DuandArup De
|
9384e4096f |
bench(flce): add configurable chunk-memory sweep (#1479)
## Summary - Add an opt-in benchmark sweep for FLCE `chunk_mem_const` values. - Encode each chunk budget in the provider name so CSV speed and memory rows remain distinguishable. - Allow benchmark callers to configure Triton `do_bench` warmup and repetition windows. - Use longer measurement windows only when the chunk sweep is enabled. - Reject combining the Triton chunk sweep with `LIGER_KERNEL_IMPL=cutile`. - Preserve existing benchmark providers, timings, and production FLCE defaults when the sweep is disabled. ## Details The shared benchmark repetition window can be shorter than a single FLCE iteration for large-vocabulary shapes. In that case, `triton.testing.do_bench` can effectively collect one sample and report identical p20, p50, and p80 values. Enable the sweep with: ```bash LIGER_FLCE_BENCH_CHUNK_MEM_CONSTS="1,2,4,8,16" \ LIGER_FLCE_BENCH_REP_MS="5000" \ LIGER_BENCH_TARGET="flce_chunk_sweep" \ python benchmark/scripts/benchmark_fused_linear_cross_entropy.py \ --model llama_3_8b \ --overwrite ``` The resulting providers are recorded as `liger-chunk-c1`, `liger-chunk-c2`, and so on. Without `LIGER_FLCE_BENCH_CHUNK_MEM_CONSTS`, the script retains the existing provider set and the shared 25 ms warmup / 10 ms repetition defaults. This is a benchmark-only change. It does not modify the production chunk policy or recommend a hardware-wide default. Local SM89 measurements showed that the speed/memory trade-off varies with token count and vocabulary size. This complements #1439: that PR exposes runtime chunk-size control, while this PR provides reproducible tooling for evaluating chunk-memory policies. ## Testing Done - Hardware Type: NVIDIA GeForce RTX 4070 Laptop GPU (SM89), BF16, WSL2 - [x] run `make checkstyle` to ensure code style - [ ] run `make test` to ensure correctness - The single-process suite was terminated by the local WSL2 host at 91% and 93% on two attempts, in different transformer test files, consistent with cumulative host-memory pressure. - `pytest -q test/ops/test_fused_linear_cross_entropy.py`: 144 passed. - `pytest -q test/transformers/test_rms_norm.py`: 64 passed, 8 expected XFAIL. - [ ] run `make test-convergence` to ensure convergence - Not completed locally. This PR changes benchmark tooling only and does not modify production kernel behavior. Representative benchmark validation: - comma-separated chunk values are preserved in provider order; - duplicate values are removed; - invalid and non-positive values are rejected; - a configurable repetition window produces non-degenerate p20/p50/p80 values for long-running C=1 paths; - ascending and descending large-vocabulary sweeps showed less than 1.1% maximum median order difference; - default execution remains unchanged when the sweep environment variables are absent. 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]>
|
||
|
|
91e51dfd9c |
Allow manual publishing of Liger Kernel releases (#1504)
## Summary Allow the Python `liger-kernel` wheel to be published from corrected `main` without creating another GitHub release event or changing the existing release tag. - Add `workflow_dispatch` with a boolean `publish` input, defaulting to `false`. - Preserve publication on release events; manual runs upload only when `publish=true`. - Reject release tags that disagree with the source project version before building. - Check distribution metadata and the wheel version before upload. After merge, publish the corrected 0.8.4 source with: ```bash gh workflow run publish-release.yml --repo linkedin/Liger-Kernel --ref main -f publish=true ``` ## Validation Built the 0.8.4 sdist and wheel, passed `twine check`, exercised the manual and matching-tag cases, rejected a mismatched tag, and ran the wheel metadata check. This PR changes only the publishing workflow. Experimental-package changes are intentionally excluded at the maintainer's request; distribution metadata validation does not claim that all runtime imports work. Co-authored-by: kolehma8 <[email protected]> Co-authored-by: Copilot <[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]> |
||
|
|
492bfeb791 |
Fix CUDA 13 Blackwell MoE with uniform 1SM UMMA (#1505)
## Summary Make Blackwell MoE homogeneous **1SM UMMA** so the native core builds with CUDA 13. The previous fused backward kernel combined 1SM token-gradient/activation phases with 2SM weight-gradient phases; CUDA 13 rejects mixed `.cta_group::1` and `.cta_group::2` instructions within one kernel. This PR contains **only the 1SM implementation**. Uniform-2SM alternatives were implemented and measured separately, including an NC=8 communication ablation, but are not included in the code change. Qwen is the priority workload: 1SM outperforms both tested 2SM alternatives in forward and backward for every measured Qwen shape. ### Implementation - Convert SM100 MLP3/MLP4 from joined M256xN256 two-CTA tiles to independent M128xN128 tiles, retaining the existing K tile and stage selections. - Use 1SM MMA atoms, CTA-local TMEM allocation, single-CTA TMA loads, local transaction-byte accounting, and one accumulator-release arrival per CTA. - Remove paired-CTA TMEM and accumulator rendezvous from fused MoE backward; retain the required CTA-local phase synchronization. - Keep communication geometry at **TileM=128, NC=4**. Forward math is unchanged. - Retain the existing tuning-table row identifiers for compatibility; map SM100 weight-gradient traits explicitly to the new tile geometry. - Update standalone harnesses and document CUDA 13 native-source support. The cross-entropy dW harness changes only its MoE reference path. **Unchanged:** production cross-entropy remains homogeneous 2SM; Hopper-specific production sources, release workflows, CUDA 12.9 release-toolchain pins, and published wheels are not changed. ## Performance methodology All tables report **TFLOP/s per GPU**, not the sum across GPUs. Measurements use eight NVIDIA B300 SXM6 PC GPUs, 8,192 input tokens per rank, eight microbatches, and 250 ms per-launch cooldown. Workload matching uses model, GPU count, T, D, I, E, and K, not just the model label. | Variant | Toolchain | MoE instruction groups | Communication TileM | NC | Weight-gradient tiles | |---|---|---|---:|---:|---| | Historical baseline | CUDA 12.9 | Original mixed 1SM/2SM backward | 128 | 4 | Joined M256xN256, 2SM | | **1SM: this PR** | CUDA 13.0.88 | Uniform 1SM | 128 | 4 | Independent M128xN128 | | 2SM, NC=4 ablation | CUDA 13.0.88 | Uniform 2SM in forward/backward | 256 | 4 | Joined M256xN256 | | 2SM, NC=8 ablation | CUDA 13.0.88 | Uniform 2SM in forward/backward | 256 | 8 | Joined M256xN256 | The CUDA 13 variants use the same B300 host/image, driver 580.159.04, PyTorch 2.11.0.1+cu130, and existing CUTLASS/NVSHMEM/TVM FFI installations. Their source base is `43293d84d760e402e4fd974b12410a1d47c31761`. The benchmark is MoEBenchmark at `85443d69ab065bfa4d303e1e4af56594c993ddf3`, with the same small adapter updates for all CUDA 13 variants: current team-handle arguments and the existing unique-ID bootstrap for single-node torchrun. Those changes are outside the timed kernel regions. Timing loops, max-rank CUDA-event latency, FLOP conventions, and correctness tolerances are unchanged. Forward FLOPs/GPU: `6*T*K*D*I + 5*T*K*I + 2*T*K*D`. Backward FLOPs/GPU: `16*T*K*D*I`. Percentages below are calculated from the CSV precision before display rounding. ### Workloads Here D is hidden size, I is per-expert intermediate size, E is total expert count, and K is routing top-k. T=8192 and world_size=8 for every row. | Model | D | I | E | K | |---|---:|---:|---:|---:| | Llama 4 Scout | 5120 | 8192 | 16 | 1 | | Mixtral 8x7B | 4096 | 14336 | 8 | 2 | | Mixtral 8x22B | 6144 | 16384 | 8 | 2 | | Qwen3 30B-A3B | 2048 | 768 | 128 | 8 | | Qwen3 235B-A22B | 4096 | 1536 | 128 | 8 | | Qwen3.5 35B-A3B | 2048 | 512 | 256 | 8 | | Qwen3.5 122B-A10B | 3072 | 1024 | 256 | 8 | ## Full comparison against the historical baseline Each ablation cell shows **TFLOP/s (throughput change vs. baseline)**. Positive is faster; negative is slower. ### Forward | Model | Baseline | 1SM: this PR | 2SM, NC=4 | 2SM, NC=8 | |---|---:|---:|---:|---:| | Llama 4 Scout | 886.7 | 851.6 (-3.96%) | 865.9 (-2.35%) | 920.8 (+3.85%) | | Mixtral 8x7B | 1,175.8 | 1,140.1 (-3.04%) | 1,203.1 (+2.32%) | 1,194.4 (+1.58%) | | Mixtral 8x22B | 1,174.3 | 1,155.9 (-1.56%) | 1,099.2 (-6.39%) | 1,107.6 (-5.68%) | | Qwen3 30B-A3B | 467.8 | 465.8 (-0.44%) | 366.8 (-21.60%) | 361.1 (-22.82%) | | Qwen3 235B-A22B | 762.2 | 766.4 (+0.56%) | 663.8 (-12.90%) | 704.3 (-7.60%) | | Qwen3.5 35B-A3B | 273.8 | 272.9 (-0.31%) | 232.5 (-15.09%) | 241.8 (-11.68%) | | Qwen3.5 122B-A10B | 551.1 | 554.9 (+0.68%) | 462.8 (-16.03%) | 489.7 (-11.14%) | ### Backward | Model | Baseline | 1SM: this PR | 2SM, NC=4 | 2SM, NC=8 | |---|---:|---:|---:|---:| | Llama 4 Scout | 714.0 | 614.6 (-13.93%) | 683.4 (-4.29%) | 731.5 (+2.45%) | | Mixtral 8x7B | 956.3 | 796.3 (-16.74%) | 916.4 (-4.18%) | 934.5 (-2.28%) | | Mixtral 8x22B | 1,018.2 | 839.1 (-17.59%) | 960.8 (-5.64%) | 989.1 (-2.86%) | | Qwen3 30B-A3B | 402.5 | 390.9 (-2.89%) | 327.7 (-18.58%) | 353.9 (-12.08%) | | Qwen3 235B-A22B | 646.2 | 577.8 (-10.58%) | 554.0 (-14.27%) | 571.6 (-11.53%) | | Qwen3.5 35B-A3B | 290.0 | 280.7 (-3.22%) | 226.1 (-22.05%) | 234.3 (-19.22%) | | Qwen3.5 122B-A10B | 534.0 | 489.3 (-8.36%) | 426.1 (-20.19%) | 436.4 (-18.27%) | ### Aggregate throughput changes These are unweighted geometric means of per-model throughput ratios, not arithmetic averages of TFLOP/s and not an end-to-end training-speed claim. | Comparison | Forward | Backward | |---|---:|---:| | 1SM vs. historical baseline | -1.17% | -10.64% | | 2SM NC=4 vs. historical baseline | -10.63% | -13.05% | | 2SM NC=8 vs. historical baseline | -8.02% | -9.44% | | 2SM NC=4 vs. 1SM | -9.57% | -2.69% | | 2SM NC=8 vs. 1SM | -6.94% | +1.34% | | 2SM NC=8 vs. 2SM NC=4 | +2.92% | +4.14% | ## Why choose 1SM? Qwen is the priority workload. Even the improved NC=8 2SM candidate is slower than 1SM for every measured Qwen shape in both directions: | Qwen model | 2SM NC=4 vs. 1SM forward | Backward | 2SM NC=8 vs. 1SM forward | Backward | |---|---:|---:|---:|---:| | Qwen3 30B-A3B | -21.26% | -16.15% | -22.48% | -9.46% | | Qwen3 235B-A22B | -13.39% | -4.13% | -8.11% | -1.07% | | Qwen3.5 35B-A3B | -14.83% | -19.46% | -11.41% | -16.54% | | Qwen3.5 122B-A10B | -16.59% | -12.91% | -11.75% | -10.81% | Uniform 2SM does have a useful counterexample: with NC=8, backward throughput exceeds 1SM by 19.03% for Scout, 17.36% for Mixtral 8x7B, and 17.88% for Mixtral 8x22B. Those gains explain its slightly positive all-model backward geometric mean, but do not justify choosing it for the Qwen-focused workload. The NC=8 experiment restores the communication band from 64 rows/CTA (`256/4`) to 32 (`256/8`, matching the original `128/4`). It improves backward over NC=4 for all seven shapes. However, it also changes the number of independent communication groups, ring capacity/adaptive grouping, and the existing NC-aligned launch grid: on this 148-SM GPU, NC=4 launches 148 CTAs/37 groups, while NC=8 launches 144 CTAs/18 groups. Other tuning parameters were not retuned. The proposed choice is therefore **uniform 1SM for CUDA 13 compatibility**, accepting the explicitly reported historical-baseline regressions rather than introducing the larger Qwen regressions and paired scheduling changes of the tested 2SM alternatives. Shape-selective 2SM tuning is outside this PR. ### Interpretation limits - The baseline is the recorded September CUDA 12.9 campaign, not a fresh same-image baseline. It differs in toolchain, dependencies, and intervening code. Its changes cannot be attributed exclusively to the UMMA instruction group. The unchanged-math 1SM forward control itself moves -1.17% geometrically. - The 1SM/2SM comparisons use the same CUDA 13 environment, but compare complete implementations: tile dimensions and required pairing/scheduling change with the instruction group. They do not isolate the causal cost of a single instruction choice. - These are single-campaign point estimates, without confidence intervals. In particular, small differences such as the +1.34% overall backward result are not statistical-significance claims. - The conclusion is scoped to these seven eight-GPU, 8K-token BF16 workloads; no claim is made about every Qwen configuration or multi-node performance. ## Validation - Full native CUDA 13 `sm_100f` core build. - B300 standalone MLP3/MLP4 correctness, MLP1 forward/activation-recompute correctness, and the separate MoE-reference/FSLCE-dW harness. - All seven eight-B300 model shapes compared with the existing Torch reference in forward and backward for the 1SM candidate and both 2SM ablations; no tolerances relaxed. - Retained Hopper MLP3/MLP4 harness correctness on H200 with CUDA 12.9. The benchmark's reference column is labeled "TE", but these runs were built with `HAVE_TE=OFF`: the reference is Torch matmul, not Transformer Engine. Reference comparisons were run separately from the eight-microbatch, cooled Liger-only performance measurements. ### Benchmark invocation Using the same existing MoEBenchmark executable and adapter for each candidate: ```bash export BENCH_NAMED_ONLY=1 BENCH_ONLY_CFG=_t8k export NVSHMEM_DISABLE_NCCL=1 NVSHMEM_REMOTE_TRANSPORT=none export NVSHMEM_SYMMETRIC_SIZE=6G OMP_NUM_THREADS=1 # Run separately for bench_moe_mlp and bench_moe_mlp_bwd. BENCH_BACKENDS=liger,torch_baseline N_MICROBATCHES=1 BENCH_COOLDOWN_MS=0 \ uv run --no-project python -m torch.distributed.run \ --standalone --nproc-per-node=8 --no-python ./build/bench_moe_mlp BENCH_ONLY_BACKEND=liger BENCH_SKIP_CORRECTNESS=1 \ N_MICROBATCHES=8 BENCH_COOLDOWN_MS=250 BENCH_CSV=forward_gpu8.csv \ uv run --no-project python -m torch.distributed.run \ --standalone --nproc-per-node=8 --no-python ./build/bench_moe_mlp ``` Native library SHA256 identifiers for the measured CUDA 13 candidates: ```text 1SM: 6226ae0c033515746eec2bf3bc2430624e10577f9b6ac2b72fd3ea45c90b162e 2SM NC4: 193decd95a38a8e2838379f84660c00ce42c2d3d4c027d167fae9d46c36265f8 2SM NC8: 8bc029baaecdab649a30199146867ef59e6da099801675fe49e7e50132f26b4b ``` Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
43293d84d7 |
Bump shared Liger and LCK release version to 0.8.4 (#1503)
## Summary Bump the single shared project version from 0.8.3 to 0.8.4. Both Python and native package setup scripts derive their versions from this field, and the LCK release workflow compares it with the release tag. This fixes the version mismatch observed in release run 36758816145 and prevents the Python release workflow from building another 0.8.3 wheel. No kernel or workflow changes are included. ## Validation Both the root and native `setup.py --version` entrypoints report `0.8.4`. ## Release sequencing Publish from the corrected commit after merge. Merely rerunning the old release event without correcting its source reference will still build version 0.8.3. The separately identified Python-wheel `ops.experimental` packaging issue also remains to be repaired before shipping a usable wheel. Co-authored-by: kolehma8 <[email protected]> Co-authored-by: Copilot <[email protected]>v0.8.4 |
||
|
|
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]> |
||
|
|
6b2b09ce4f |
Add requires-python >=3.10 to pyproject.toml (#1332)
## Summary Fixes #1302. `pyproject.toml` has no `requires-python`, so pip resolves and installs liger-kernel on any interpreter. Users on unsupported Pythons get runtime failures instead of a clean "requires Python >= X" resolver message. This adds `requires-python = ">=3.10"`, matching the project's own floor: every CI test workflow (nvi-ci, amd-ci, checkstyle, benchmark) and the release-publish workflow run Python 3.10. Picking up the hand-off from #1307 (closed by its author due to inactivity, fix noted as still valid). ## Testing Done Metadata-only change. Verified the line parses and the floor matches CI: ``` $ python -c "import tomllib; print(tomllib.load(open('pyproject.toml','rb'))['project']['requires-python'])" >=3.10 ``` Co-authored-by: yashb98 <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
Kobi HikriandVaibhav Jindal
|
c8c999932f |
docs: sync README triton floor with setup.py and fix contributing install command (#1301)
## Summary Two small documentation accuracy fixes: 1. **README CUDA requirements** said `triton >= 2.3.0`, but `setup.py` resolves `triton>=2.3.1` — issue #682 shows a user hitting exactly this confusion. Updated the README to `>= 2.3.1` so both sources agree. 2. **docs/contributing.md step 2** had `pip install . -e[dev]`, which isn't valid pip syntax (the guide even carries a workaround note for the error it causes). Replaced with `pip install -e ".[dev]"`, matching AGENTS.md and the README. ## Testing Done Doc-only change. Verified `setup.py` line 17 (`triton>=2.3.1`) and that `pip install -e ".[dev]"` is the form AGENTS.md uses. - Hardware Type: n/a (docs) - [x] run `make test` to ensure correctness — n/a, docs only - [x] run `make checkstyle` to ensure code style — n/a, docs only - [x] run `make test-convergence` to ensure convergence — n/a, docs only --- For transparency (per AGENTS.md's AI-assistant support): I used AI assistance to spot and draft this; I verified both lines against setup.py and issue #682 myself. 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) |
||
|
|
be2a4eb14c |
Add LCK CUDA-selectable NVSHMEM extras; retain CUDA 12.9 releases (#1492)
## Summary Make the LCK NVSHMEM dependency selectable from the pip command while keeping the combined Hopper/Blackwell release wheel on **CUDA 12.9.1**. ```bash pip install "liger-cute-kernels[cu12]" # For a locally built CUDA 13 Hopper wheel: pip install "/path/to/liger_cute_kernels-<version>-py3-none-linux_x86_64.whl[cu13]" ``` - Move NVSHMEM out of unconditional dependencies into `cu12` and `cu13` extras, each pinned to `3.6.5`. Plain LCK leaves NVSHMEM provisioning to the caller. - Keep the existing digest-pinned CUDA 12.9.1 release image, SM90a/SM100f targets, and `nvidia-nvshmem-cu12` build prerequisite. Explicitly package `".[cu12]"` and validate that both optional dependencies activate only for their corresponding extra. - Fix the Hopper non-RDC cubin command to prepend `-I` to every CUDA include directory. CUDA 13 exposes an additional CCCL include directory; previously NVCC interpreted that bare path as a second input file. - Update installation documentation and missing-runtime guidance. Extras select runtime dependencies, not different compiled binaries. Select exactly one matching the core's build toolkit; the two NVSHMEM distributions share an installation namespace. Source builds still require NVSHMEM to be installed before compilation. ### Why the release stays on CUDA 12.9 The complete Hopper core and non-RDC cubin build with CUDA 13.0.88 after the include-path fix. However, the complete SM100 build fails in three MoE backward variants because `ptxas` rejects mixed `.cta_group::1` and `.cta_group::2` instructions in one kernel. This PR does not remove those variants or attempt a kernel redesign. NVIDIA reports the same newer-toolkit restriction in [CCCL #8297](https://github.com/NVIDIA/cccl/pull/8297). The CUDA 13 extra supports appropriately built source artifacts; it does not claim that the combined release wheel builds with CUDA 13. ## Testing Done - Hardware Type: **4 x NVIDIA H200**, CUDA compiler **13.0.88**, Torch **2.11.0.1+cu130**. - **21 focused build/loader tests passed**, including optional dependency metadata, missing-runtime guidance, and CMake include-list expansion with single/multiple CUDA include paths. - Ruff lint and formatting checks passed for the changed Python files. - Parsed the release workflow's embedded Python, checked the packaging shell syntax, and exercised its optional-dependency validation against metadata from the built wheel. - Built, packaged, and installed the complete Hopper core plus non-RDC cubin. The test wheel was 46.97 MiB and used `linux_x86_64`, not a manylinux claim, because that test image uses glibc 2.38. - Ran six numerical cases: Qwen3-30B-A3B, Qwen3-235B-A22B, and Llama-4-Scout at 1,024/2,048 tokens per GPU, EP=4, seed 42. Compared CUDA 13 default and non-RDC-enabled execution against the existing CUDA 12 wheel on the same GPUs and with the same Python frontend and inputs. - Covered forward output, input gradients, all three expert-weight gradients, and selected routing-weight gradients against Torch FP32. All 36 maximum-error metrics matched the CUDA 12 baseline exactly; the change in the equally weighted average mean error was below 0.000008% in both CUDA 13 modes. This is a single-seed comparison, not a claim of bitwise-identical tensors or exhaustive coverage. - Complete CUDA 13 SM100 compilation was attempted and failed as described above. No combined CUDA 13 release wheel was produced. A fresh full CUDA 12.9 release build was not run for this PR; its build image and native compilation settings remain unchanged apart from the compatible include-list fix. The full suites below were not run; the targeted checks and numerical coverage are listed above. - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
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]> |