mirror of
https://github.com/linkedin/Liger-Kernel.git
synced 2026-10-12 02:17:47 +00:00
main
190
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
|
|
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 |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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]> |
||
|
|
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 |
||
|
|
d5df610cbf |
Bump Liger packages to 0.8.3 (#1471)
## Summary - bump the shared public version from `0.8.2` to `0.8.3` - update both `liger-kernel` and `liger-cute-kernels`, since the native package resolves its version from the root `pyproject.toml` - add `liger-kernel[lck]` as an optional install path - pin the extra to `liger-cute-kernels==0.8.3` so the Python and native distributions cannot drift across releases - document stable and editable `lck` extra installation ## Resulting installs ```bash pip install liger-kernel pip install "liger-kernel[lck]" ``` The second command installs the normal Liger dependencies plus the exactly matching standalone native wheel. ## Validation - root wheel: `liger_kernel-0.8.3-py3-none-any.whl` - root wheel metadata includes `Provides-Extra: lck` - root wheel metadata includes `liger-cute-kernels==0.8.3; extra == "lck"` - standalone LCK package metadata resolves to `0.8.3` - targeted LCK version tests pass - Ruff check and format checks pass for `setup.py` --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
f6385e909b |
feat: add FusedLinearKLDivLoss (fused linear + KL divergence for distillation) (#1423)
## Summary
Add `FusedLinearKLDivLoss`, which fuses the vocabulary-head linear
projection with
KL(Q || P) loss against an explicit target distribution, for knowledge
distillation
workflows where teacher probabilities are available directly (e.g.
cached / offline
soft labels), complementing `LigerKLDIVLoss` (element-wise, no fusion)
and
`LigerFusedLinearJSD` (teacher side needs its own projection).
The kernel follows the chunked gradient-in-forward structure of
`fused_linear_jsd.py`,
so the full `BT x V` logits tensor is never materialized, and adopts the
`_CHUNK_MEM_CONST = 16` token-chunk geometry of
`fused_linear_cross_entropy.py`.
## Details
- Supports `reduction` (`"batchmean"` / `"mean"` / `"sum"`, same
semantics as
`torch.nn.KLDivLoss`), `ignore_index` via `shift_labels`, softmax
`temperature`,
and an `eps` clamp so that `0 * log(0)` contributes 0.
- Supports `accum_dtype` for fp32 weight-gradient accumulation across
chunked
low-precision GEMMs (same semantics as `LigerFusedLinearJSD` /
`LigerFusedLinearCrossEntropyLoss`), including the
`torch.addmm(out_dtype=fp32)`
fast path on torch >= 2.8 / sm80+ and a fallback elsewhere.
- The log-softmax backward is computed in place on the saved log-probs
and the
`chunk_size x V` fp32 temporaries are freed eagerly, so the peak during
the
trailing GEMMs stays low.
- `reduction="none"` is intentionally unsupported (gradient-in-forward
requires a
scalar loss); a `ValueError` points users to `LigerKLDIVLoss` instead.
- bf16 + `reduction="sum"` test tolerances follow the existing
convention in
`test_fused_linear_cross_entropy.py`.
- Peak-memory improvement over the unfused PyTorch baseline (full
fwd+bwd,
`max_memory_allocated`, bf16, llama-3 head shape H=4096 / V=128256, RTX
4060
Laptop GPU): BT=2048 6537 -> 4310 MB (1.52x), BT=4096 12072 -> 6594 MB
(1.83x);
at BT=1024 the benchmark harness reports 4793 -> 4244 MB. Runtime is at
parity
with the baseline at BT=1024 (0.89x).
## Testing Done
- Hardware Type: NVIDIA GeForce RTX 4060 Laptop GPU (sm89); NVIDIA Tesla
T4 (sm75)
- [x] run `make test` to ensure correctness (32 pre-existing failures
unrelated to
this change: 18 `test_mlp.py` shared-memory OOR on sm89 and 14
`test_grpo_loss.py`
LUSPO numerics; both reproduce on `main` without this change)
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence (does not
exercise the new
standalone loss)
<details>
<summary>Test logs (RTX 4060 Laptop, torch 2.11.0+cu128, triton
3.7.1)</summary>
```
make test: 32 failed, 4152 passed, 914 skipped, 15 xfailed in 1050.50s (0:17:30)
test/transformers/test_fused_linear_kl_div.py: 172 passed, 0 failed, 0 skipped
failures (all pre-existing on this GPU, reproduced on main without this change):
18 x test/transformers/test_mlp.py
triton.runtime.errors.OutOfResources: shared memory
Required: 131072, Hardware limit: 101376 (sm89 has 100 KB usable smem)
14 x test/chunked_loss/test_grpo_loss.py
chunked GRPO LUSPO numerics (also seen by other contributors, see #1397)
make checkstyle: ruff check and ruff format --check pass on all files in this PR
```
</details>
<details>
<summary>Test logs (Kaggle Tesla T4, torch 2.10.0+cu128, triton
3.6.0)</summary>
```
test/transformers/test_fused_linear_kl_div.py: 51 passed, 25 skipped in 10.29s
(bf16 cases skipped by design: T4 is sm75 without native bf16 support)
```
</details>
|
||
|
|
0245181a11 |
Use native CUTLASS + NVSHMEM package terminology (#1434)
## Summary - Remove the internal "LCK" nickname from Liger frontend classes, helpers, messages, tests, and user-facing documentation. - Use `native` terminology for the CUTLASS + NVSHMEM tensor-parallel adapter. - Describe `liger_cute_kernels` as the native CUTLASS + NVSHMEM kernel package rather than a MoE-only package. - Document both current operator families: expert-parallel MoE and tensor-parallel fused scaled linear cross entropy. - Replace all remaining `LigerCommKernels` and `liger_comm_kernels` references with `LigerCuteKernels`. - Keep existing internal build identifiers unchanged; this is a terminology/documentation update with no runtime behavior change. ## Validation - `make checkstyle` - H200 frontend tests: 28 passed --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
cc7f68fc9c |
[CUDA] Add tensor-parallel fused scaled cross entropy frontend (#1428)
## Summary - Add `LigerFusedLinearScaledCrossEntropyTPFunction` and `liger_fused_linear_scaled_cross_entropy_tp`, with `tp_group` passed as a call-time argument. - Route Hopper BF16 inputs to the optional `liger_cute_kernels` implementation when its FSLCE runtime surface is available. - Keep the fallback entirely in Liger using a chunked tensor-parallel implementation adapted from Verl's fused PPO formulas; there is no runtime dependency on the Verl package. - Match both paths on global per-token NLL/entropy, ignored targets, local `dW`, and TP-summed `dX`. - Leave NVSHMEM bootstrap and team preparation application-owned. FSLCE resolves only the call-time `tp_group`, so it coexists with LigerMoE using a different `ep_group` without frontend-global process-group state. - Remove the native NCCL dependency and all Torch communicator internals from the LCK path. Forward and backward communicate exclusively through the prepared NVSHMEM team, and the TVM-FFI core remains Torch-ABI-independent. - Extend the reusable local reduction API with FP32 `MAX`. The forward WGMMA kernel performs split last-arrival merging, local MAX, correction, and local SUM in its epilogue. - For two-host execution, each local GPU sends one unique contiguous token shard over IBRC; the follow-up kernel merges the remote online-softmax state, gathers shards with NVLS, and writes NLL/LSE/entropy. - Bind native FFI calls to the input CUDA device so autograd worker threads have a valid context for TMA descriptor construction. ## Validation - `make checkstyle` - `python -m pytest test/transformers/test_fused_linear_scaled_cross_entropy.py -q` - 15 passed, 1 skipped - H200 distributed frontend matrix using `torch.multiprocessing.spawn` - 10 passed in 77.45s: - WORLD process groups at TP1, TP2, TP4, and TP8 - 4-GPU contiguous subgroups `[0,1]` / `[2,3]` - 4-GPU strided subgroups `[0,2]` / `[1,3]` - Built-in Liger fallback at TP2 WORLD and strided TP4 - LigerMoE on `ep_group=WORLD` followed by FSLCE on independent strided TP groups in the same four processes - DirectPeer fallback with NVLS disabled - Native forward self-test: aligned, entropy, ragged M/H/V, temperature, and large split-N cases passed. - Raw TVM-FFI world/subgroup tests, including forward and backward CUDA Graph replay: 4 passed. - Two-host TP16 IBRC random-data correctness: passed. - `libliger_cute_kernels.so` has no NCCL dynamic dependency. ### Forward performance M=8192, H=4096, Vglobal=131072, entropy disabled: | TP | NVSHMEM p50 | Previous p50 | Delta | |---:|---:|---:|---:| | 1 | 11.227 ms | 11.148 ms | +0.7% | | 2 | 5.639 ms | 5.594 ms | +0.8% | | 4 | 2.832 ms | 2.821 ms | +0.4% | | 8 | 1.442 ms | 1.453 ms | -0.8% | | 16 (two-host IBRC) | 0.829 ms | 0.784 ms | +5.7% | TP1-8 are within the agreed 1% run-to-run noise threshold; the TP16 IBRC gap is accepted. --------- Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
e6a81bb0c3 |
Add Liger Kernel support for Muse Glimmer (#1390)
## Summary Adds Liger Kernel support for Muse Glimmer (`model_type=muse_glimmer`), covering both capabilities requested in #1378: 1. FusedLinearCrossEntropy / skip-logits loss path, avoiding the `[B, S, V]` logits tensor (Muse Glimmer's `vocab_size` is 202,048) 2. `return_token_accuracy=True` → `outputs.token_accuracy`, so TRL's `SFTTrainer` can log `mean_token_accuracy` without a trainer-side fallback Registered in `MODEL_TYPE_TO_APPLY_LIGER_FN`, so `use_liger_kernel=True` picks it up. Patches RoPE, RMSNorm, SwiGLU, vision LayerNorm (instance path only), CrossEntropy and FusedLinearCrossEntropy. ## Details **Softcap.** HF computes `T * tanh(logits * m / T)` where `m = output_multiplier` and `T = final_logit_softcapping`, but Liger's fused softcap only implements `T * tanh(logits / T)`. `m` is folded into the hidden states instead, since `(m * h) @ W.T == m * (h @ W.T)` — scaling `[batch, seq, hidden]` rather than `[vocab, hidden]`. **Two RMSNorm variants.** `MuseGlimmerRMSNorm` scales by `weight` directly (gemma-style fp32 compute, offset 0); `MuseGlimmerTextCenteredRMSNorm` scales by `(1 + weight)` from zero-init (offset 1). The `with_scale=False` instances, `qk_norm`, `embed_norm`, `perception_emb_norm`, have no weight parameter and fall back to a torch implementation matching HF exactly; an assert guards against that fallback silently dropping a scale. **`in_place=False`** on the TextCentered norms is required, not conservative: each decoder layer hands `pre_feedforward_layernorm` the same tensor it retains as the residual. Note that `token_accuracy` is populated on the fused path only; with `skip_logits=False` (eval) it stays `None`, matching existing patched models such as Qwen3. ## Testing Done - `test/transformers/test_muse_glimmer.py` (new) — instance-path numerical parity against the unpatched model: loss plus embedding / SwiGLU / RMSNorm gradients in fp32 and bf16, the non-fused softcap branch, vision LayerNorm, and the `token_accuracy` / `predicted_tokens` contract. The convergence suite patches at module level, so this is what covers `_apply_liger_kernel_to_instance` — the path TRL actually takes. - `test/transformers/test_monkey_patch.py` — instance patching - Convergence in bf16 and fp32: `test_mini_models`, `test_mini_models_multimodal`, `test_mini_models_with_logits` ### Memory H100, bf16, forward + backward, no optimizer step. Randomly initialised (no public checkpoint), real `hidden_size=6656` / `vocab_size=202048`, `num_hidden_layers` reduced from 52 so a 30B-class model fits on one GPU. | seq_len | peak mem (HF) | peak mem (Liger) | saved | |--------:|--------------:|-----------------:|------:| | 1024 | 18.4 GB | 20.0 GB | -9% | | 2048 | 19.6 GB | 21.2 GB | -8% | | 4096 | 30.4 GB | 23.5 GB | 23% | | 8192 | 52.0 GB | 28.2 GB | 46% | | 16384 | OOM | 38.2 GB | — | The crossover sits near 3K: FLCE accumulates `grad_weight` in fp32, a fixed 5.4 GB for `[202048, 6656]` against bf16's 2.7 GB. Below ~3K that outweighs the logits it avoids; above it the saving compounds, and 16K only runs at all with Liger. ### Throughput FLCE trades speed for memory by design. End to end that penalty is diluted by the transformer body, and the reduced depth above overstates it: | layers | lm_head share | tok/s (HF) | tok/s (Liger) | ratio | |-------:|--------------:|-----------:|--------------:|------:| | 4 | 41% | 22,405 | 13,433 | 0.60x | | 8 | 26% | 13,685 | 10,396 | 0.76x | | 16 | 15% | 7,705 | 7,063 | 0.92x | `lm_head share` is measured from the model's parameter split, which tracks the matmul FLOP split. The shipping model's 52 layers put it at ~5%, below the last row here. - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence Co-authored-by: Zhipeng Wang <[email protected]> |
||
|
|
f65ae6576f |
Update the ascend npu test repository link (#1411)
### Summary Transfer the ascend npu test of liger-kernel from the individual warehouse to the Ascend organization. |
||
|
|
0515c5aa2c |
Fix broken star history chart (#1385)
The star history chart on the README and docs was broken due to GitHub stargazer API restrictions. Update the chart to use a working data source that requires no API token so the chart renders correctly again. Co-authored-by: Dessalines39394 <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
|
500866071d |
[NPU] Pin torch/torch_npu=2.9.0, triton-ascend=3.2.2 and cann=9.1.0 for NPU (#1392)
## Summary Adapt three Ascend kernels to the new stack (`torch`/`torch_npu` 2.9.0, `triton-ascend` 3.2.2, CANN 9.1 / BiShengIR 1.2.0). `triton` itself stays 3.2.0; the failures come from compiler/runtime default changes in `triton-ascend` 3.2.2, not from test logic. Each change is scoped to the kernel that actually hits a 3.2.2 constraint. GPU paths are untouched. The same sources still run on the previous stack (`torch`/`torch_npu` 2.7.1 + `triton-ascend` 3.2.1). ## Details Stack delta that matters: | | previous (green) | new (broken) | |---|---|---| | torch / torch_npu | 2.7.1 | 2.9.0 | | triton | 3.2.0 | 3.2.0 | | triton-ascend | 3.2.1 | **3.2.2** | | CANN / BiShengIR | 9.0 / older | **9.1 / 1.2.0** | 3.2.2 defaults `multibuffer=True` and cannot lower some mix cube+vector programs to static UB shapes. Fixes: 1. **Fused MoE** — fp32 `_fused_up_proj_swiglu_kernel` failed compile with `ConvertLinalgRToBinary` / `Failed to obtain op buffer shape size which should be static`. Cause: gather + dual `tl.dot` + SwiGLU in one kernel. Fix: cube kernel writes `pre_act` only; new vector kernel `_swiglu_from_pre_act_kernel` does `silu(gate) * up`. Routing / down-proj / backward unchanged. 2. **mHC** — auto-multi-buffer races fp32 mix cube+vector results and hangs compiling mix + `tl.atomic_add`. Fix: pass `multibuffer=False` only on those launches (`mm_norm_fwd/bwd`, `pre_bwd`, `post_res_bwd`, `coeffs_bwd_assemble`). Vector-only forward kernels keep the 3.2.1 default. 3. **Fused neighborhood attention** — consecutive forwards on the same stream hang until QK+AV retire (`test_fused_neighborhood_attention_deterministic`). Fix: `torch.npu.current_stream().synchronize()` after the AV launch (stream-local, not a device-wide barrier). ## Testing Done - Hardware Type: Ascend NPU (`torch` 2.9.0, `torch_npu` 2.9.0, `triton-ascend` 3.2.2, CANN 9.1) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence Targeted pytest on 3.2.2: Coming soon. |
||
|
|
70c83b2631 |
perf(cutedsl): add Blackwell fast path to RMSNorm (#1388)
## Summary This PR adds a Blackwell fast path to the CuTeDSL RMSNorm op (`ops/cutedsl/ops/rms_norm.py`), the only generic CuTeDSL op with no arch-specific path (e.g. `swiglu.py` has packed-f32x2 math and a TVM-FFI direct call; `rope.py` has a TMA kernel tuned on B200; and `cross_entropy.py` dispatches on `infer_device_arch() == "blackwell"`, while `rms_norm.py` had no arch query at all). **This PR changes how the two vector fast paths — `_launch_fwd_vector` and `_launch_bwd_fused` — are launched, and the instruction mix of the backward. Their tiling, their CTA-level reduction, and the op's public API are unchanged.** ### 1. TVM-FFI direct-call launch (host-side; the bulk of the win) `_launch_fwd_vector` and `_launch_bwd_fused` gain a fast branch that skips the `from_dlpack` / memref marshalling entirely and hands PyTorch tensors straight to the compiled function: - New `_make_fwd_vector_ffi` / `_make_bwd_fused_ffi` closures capture every constexpr, so the compiled signature is just tensors + fp32 scalars. - New `_get_fwd_vector_ffi` / `_get_bwd_fused_ffi` compile those once per (dtype, geometry, device) against abstract tensors via `make_fake_tensor`, cached in a new `_ffi_compile_cache`. `eps` / `offset` stay runtime scalars, so changing either does not trigger a recompile. - The stream parameter is compiled against `make_fake_stream(use_tvm_ffi_env_stream=True)`, so it drops out of the call signature and resolves from the caller's env stream — no per-call stream query from Python. Result: forward host cost **54us -> 15us**, backward **74us -> 26us**. ### 2. packed-f32x2 math in the fused backward (device-side) `_rms_norm_bwd_fused_vector_kernel` gains a `PACKED_MATH: cutlass.Constexpr` branch that computes the dX/dW inner products pairwise with `fma_packed_f32x2` / `mul_packed_f32x2` /`add_packed_f32x2` (two fp32 lanes per instruction), keeping two accumulators that fold before the CTA reduction. Selected by the new `_use_packed_math`, which requires sm_10x **and** `n_cols > 4096` — the 16-warp regime where the kernel is issue-bound rather than memory-bound. Per-width measurements justifying that threshold are in its docstring. ### Behavior when the fast path does not apply - Without `apache-tvm-ffi` the original marshalling launch runs; - Off data-center Blackwell (including Hopper and RTX 50xx) the scalar math runs; - Irregular or unaligned shapes still take the existing scalar/split paths. Both halves are A/B switchable with `LIGER_RMS_FORCE_NO_FFI` / `LIGER_RMS_FORCE_NO_PACKED`. ### Why the op was slow in the first place The kernels were already faster than Triton on *device* time (forward 5.8us vs 7.8us at bf16 4096x2048; fused backward 26.3us vs 32.4us). Wall clock disagreed because each call spent ~48us marshalling tensors into `cute.Tensor` handles before it could enqueue, so the op was host-bound at every shape up to 16384x4096 — measured wall time tracked host time (~54us), not kernel time. ## Results 1x B200, bf16, llama mode, op-level, cutedsl vs Triton (>1.0 = cutedsl faster): | shape | fwd before | fwd after | bwd before | bwd after | |---|---|---|---|---| | 1024x2048 | 0.65x | **2.36x** | 0.81x | **2.36x** | | 4096x2048 | 0.66x | **2.32x** | 0.81x | **2.29x** | | 4096x4096 | 0.69x | **2.33x** | 0.82x | **1.51x** | | 4096x8192 | 0.66x | **1.52x** | 1.08x | 1.17x | | 16384x2048 | 0.67x | **1.65x** | 1.73x | 1.75x | | 16384x4096 | 0.90x | **1.12x** | 1.75x | 1.75x | | 16384x8192 | 1.24x | 1.24x | 1.17x | **1.29x** | | 32768x4096 | 1.11x | 1.12x | 1.84x | 1.83x | fp32 moves the same way: forward 0.66-1.08x -> 0.99-2.29x, backward 0.80-1.73x -> 1.15-2.58x. **No shape regressed in either dtype.** "before" is measured by running the same build with both toggles off, so it is exactly the old code path. The packed-math half accounts for the wide-row backward gains specifically: at bf16 16384x8192 it is 245.8us -> 222.0us (-9.7%) and at fp32 16384x6144 251.3us -> 194.6us (-22.6%), device time. ## Validation - Three new tests in `test/cutedsl/test_rms_norm.py` cover ground the existing parity suite cannot reach — it only goes up to 4096 wide (so it never compiles the packed backward) and always runs on the default stream: - `test_rms_norm_ffi_matches_marshalling_path` — bit-identity between the two launch paths, affine **and** non-affine; the non-affine case is what pins the dummy-W layout. - `test_rms_norm_packed_backward_matches_scalar` — packed vs scalar backward at 8192. - `test_rms_norm_runs_on_the_current_stream` — stream binding via CUDA graph capture. Both comparison tests assert the two variants were really compiled, so they cannot degenerate into comparing a path against itself. Mutation-checked: corrupting the packed dW accumulation or the non-affine dummy-W layout makes them fail. - Full cutedsl suite on B200: 579 passed with an **identical 33-failure set before and after** (pre-existing bf16/fp32 parity gaps in `cross_entropy` and RMSNorm's bf16-llama dW) — no new failures. - `ruff check` / `ruff format` clean. ## Packaging note `apache-tvm-ffi` is added to the `cutedsl` extra. It stays optional at runtime, but the entire forward win depends on it and the existing swiglu kernel already has the same latent dependency. Happy to drop this if you would rather keep the extra minimal. |
||
|
|
30978024fe |
Update CuTe DSL backend docs: op list and DSL version (#1386)
## Summary Two doc fixes to the CuTe DSL backend section of `README.md`. **1. Stale op list.** The "Enable CuTe DSL Backend" section claimed the backend provides `cutlass.cute` implementations of RMSNorm, cross entropy, and fused scaled cross entropy. It also ships **RoPE**, **SwiGLU**, and **fused linear cross entropy** — all exported from `src/liger_kernel/ops/cutedsl/ops/__init__.py`, all genuine `cutlass.cute` kernels (not fallbacks), and all covered by tests: - `src/liger_kernel/ops/cutedsl/ops/rope.py` → `test/transformers/test_cutedsl_rope.py` - `src/liger_kernel/ops/cutedsl/ops/swiglu.py` → `test/transformers/test_swiglu_cutedsl.py` - `src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py` and `fused_linear_cross_entropy_sm90.py` → `test/cutedsl/test_fused_linear_cross_entropy.py` Reformatted the list as bullets since it no longer fits comfortably inline. **2. Version mismatch.** The Optional Dependencies section documented `nvidia-cutlass-dsl >= 4.5.2`, but the `cutedsl` extra in `setup.py:41` pins `nvidia-cutlass-dsl>=4.6.0`. Aligned the README to setup.py. ## Testing Documentation-only change — no code paths touched. Co-authored-by: Zhipeng Wang <[email protected]> Co-authored-by: Claude Opus 5 (1M context) <[email protected]> |
||
|
|
62b691b721 |
perf(cutedsl): optimize Hopper FLSCE forward (#1370)
## Summary
- replaces the SM90 M-outer and M-fast FLSCE forwards with one fixed
cluster-M2 fragment kernel
- computes a logical `M128 x N320 x K64` tile as two N160 accumulators
and four N80 epilogue slices
- adds a measured runtime split-N lookup for the profiled 16K/32K
long-sequence shapes
- supports differentiable BF16 entropy through online weighted
statistics and a small split finalizer
- preserves `m_tiles_per_cluster` for API compatibility while removing
its slower schedule variants
- keeps the existing one-kernel N256 backward unchanged
- extracts shared SM90 forward helpers and removes 2,320 lines of
obsolete forward code
## Kernel design
- WG0: 24-register TMA producer
- WG1/WG2: 240-register WGMMA plus online-softmax epilogue
- cluster-M2 pairs process adjacent token tiles and multicast shared
weight panels
- safe one-stage AB release lag prevents TMA overwriting live WGMMA
operands
- two padded FP16 SMEM buffers stage four N80 slices; logits never
round-trip through HBM
- cluster pairs emit compact `(max, sum, target[, weighted])`
statistics; a small kernel finalizes NLL/LSE/entropy
- unprofiled shapes derive split count from active cluster capacity,
including uneven M-pair assignments
The H100 lookup for `H=4096`, `V=131072` selects:
| M | Entropy | `split_n` |
|---:|:---:|---:|
| 16384 | No | 4 |
| 16384 | Yes | 2 |
| 32768 | No | 2 |
| 32768 | Yes | 2 |
## H100 performance
BF16, `H=4096`, `V=131072`, cuTile 1.5.0. Effective TFLOPS count the
common projection work: forward `2*M*H*V`, backward `6*M*H*V`, and full
`8*M*H*V`.
### Forward
Medians from 60 interleaved samples per provider.
| M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback |
|---:|:---:|---:|---:|---:|
| 2048 | No | **3.12 ms / 706 TFLOPS** | 3.19 ms / 690 TFLOPS | 11.36 ms
/ 194 TFLOPS |
| 2048 | Yes | **3.13 ms / 703 TFLOPS** | 3.25 ms / 676 TFLOPS | 11.37
ms / 193 TFLOPS |
| 4096 | No | **6.05 ms / 727 TFLOPS** | 6.28 ms / 701 TFLOPS | 22.62 ms
/ 194 TFLOPS |
| 4096 | Yes | **6.11 ms / 720 TFLOPS** | 6.39 ms / 688 TFLOPS | 22.68
ms / 194 TFLOPS |
| 8192 | No | **12.25 ms / 718 TFLOPS** | 12.41 ms / 709 TFLOPS | 45.35
ms / 194 TFLOPS |
| 8192 | Yes | **12.30 ms / 715 TFLOPS** | 12.70 ms / 693 TFLOPS | 45.59
ms / 193 TFLOPS |
| 16384 | No | **23.84 ms / 738 TFLOPS** | 25.11 ms / 700 TFLOPS | 90.81
ms / 194 TFLOPS |
| 16384 | Yes | **24.09 ms / 730 TFLOPS** | 26.10 ms / 674 TFLOPS |
90.97 ms / 193 TFLOPS |
| 32768 | No | **49.69 ms / 708 TFLOPS** | 54.31 ms / 648 TFLOPS |
180.01 ms / 195 TFLOPS |
| 32768 | Yes | **49.85 ms / 706 TFLOPS** | 55.69 ms / 632 TFLOPS |
179.85 ms / 196 TFLOPS |
The tuned 16K/32K CuTe forward is 5.4-11.7% faster than cuTile.
### Backward
Medians from 30 interleaved samples per provider.
| M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback |
|---:|:---:|---:|---:|---:|
| 8192 | No | **37.93 ms / 696 TFLOPS** | 42.55 ms / 620 TFLOPS | 85.17
ms / 310 TFLOPS |
| 8192 | Yes | **37.89 ms / 696 TFLOPS** | 41.43 ms / 637 TFLOPS |
121.00 ms / 218 TFLOPS |
| 16384 | No | **79.42 ms / 665 TFLOPS** | 83.78 ms / 630 TFLOPS |
166.85 ms / 316 TFLOPS |
| 16384 | Yes | **78.20 ms / 675 TFLOPS** | 84.12 ms / 627 TFLOPS |
239.39 ms / 220 TFLOPS |
| 32768 | No | 163.34 ms / 646 TFLOPS | **163.21 ms / 647 TFLOPS** |
331.08 ms / 319 TFLOPS |
| 32768 | Yes | **160.64 ms / 657 TFLOPS** | 161.73 ms / 653 TFLOPS |
477.51 ms / 221 TFLOPS |
### Full forward + backward
| M | Entropy | CuTe SM90 | cuTile | Verl Torch fallback |
|---:|:---:|---:|---:|---:|
| 8192 | No | **50.39 ms / 698 TFLOPS** | 57.15 ms / 616 TFLOPS | 128.76
ms / 273 TFLOPS |
| 8192 | Yes | **50.22 ms / 701 TFLOPS** | 56.86 ms / 619 TFLOPS |
165.63 ms / 212 TFLOPS |
| 16384 | No | **104.77 ms / 672 TFLOPS** | 112.03 ms / 628 TFLOPS |
254.91 ms / 276 TFLOPS |
| 16384 | Yes | **103.65 ms / 679 TFLOPS** | 111.21 ms / 633 TFLOPS |
328.37 ms / 214 TFLOPS |
| 32768 | No | **211.30 ms / 666 TFLOPS** | 221.34 ms / 636 TFLOPS |
508.33 ms / 277 TFLOPS |
| 32768 | Yes | **210.60 ms / 668 TFLOPS** | 221.39 ms / 636 TFLOPS |
655.89 ms / 215 TFLOPS |
Against the removed generic CuTe forward, the fragment path measured:
| M | Generic | Fragment | Speedup |
|---:|---:|---:|---:|
| 2048 | 4.56 ms | 3.28 ms | **1.39x** |
| 4096 | 7.58 ms | 6.41 ms | **1.18x** |
## Validation
- H100: `47 passed, 19 skipped`
```bash
pytest -q \
test/cutedsl/test_fused_scaled_cross_entropy_forward_fragment_sm90.py \
test/cutedsl/test_fused_scaled_cross_entropy_sm90.py \
test/transformers/test_fused_linear_scaled_cross_entropy.py
```
- independent Torch oracle coverage for NLL, entropy, temperature,
ignore index, ragged M/H/V, K64/K4096, final-column targets,
hardware/uneven split assignments, entropy-only backward, and runtime
lookup selection
- `make checkstyle` passes
## Compatibility
The public input/output contract and original N256/wave-8 backward are
unchanged. Entropy and unprofiled shapes use the same fragment forward
rather than retaining slower fallback schedules.
---------
Co-authored-by: jkolehma <[email protected]>
Co-authored-by: Copilot <[email protected]>
|
||
|
|
b304c890f0 |
[cuTile] Add fused scaled cross entropy (#1359)
## Summary Adds a cuTile implementation of `LigerFusedLinearScaledCrossEntropyFunction`, matching the existing per-token scaled NLL and optional differentiable entropy contract. **Compared with the VERL-derived PPO fallback on the combined NLL-plus-entropy workload, cuTile is 3.1x faster on H100 and 3.7x faster on B200 while reducing peak memory by 55% and 51%, respectively.** | Hardware | Implementation | Forward | Backward | Full | Speedup vs VERL | Peak full memory | Memory saving vs VERL | |---|---|---:|---:|---:|---:|---:|---:| | H100 | **cuTile** | **6.24 ms** | **20.88 ms** | **26.15 ms** | **3.14x** | **2432 MiB** | **55%** | | H100 | VERL-derived fallback | 22.36 ms | 59.40 ms | 82.03 ms | 1.00x | 5384 MiB | — | | B200 | **cuTile** | **3.17 ms** | **9.18 ms** | **12.01 ms** | **3.68x** | **2640 MiB** | **51%** | | B200 | VERL-derived fallback | 12.34 ms | 31.63 ms | 44.16 ms | 1.00x | 5336 MiB | — | - registers the implementation behind `LIGER_KERNEL_IMPL=cutile` - supports FP32, FP16, and BF16 matching `input`/`weight` dtypes, scalar positive temperature, ragged dimensions, and `ignore_index` - bounds forward and backward logits workspaces to 256 MiB portably and automatically uses 512 MiB for large FP16/BF16 Blackwell workloads - balances token chunks when power-of-two rounding would add a launch or leave a less-than-half-sized final chunk - supports `LIGER_CUTILE_SCALED_CE_WORKSPACE_MB` for workload-specific workspace tuning - recomputes logits chunk-by-chunk in backward, transforms them into `dZ` in place, writes `dX` directly, and accumulates `dW` directly with `mm`/`addmm` outputs to avoid full `V x H` temporaries - tests the direct cuTile class against the independent VERL-derived fallback for combined, NLL-only, and entropy-only gradients - adds the cuTile provider to the scaled-cross-entropy benchmark and fixes that benchmark to reset both input and weight gradients - centralizes README documentation for the shared cuTile/CuTe/fallback public API No open PR currently implements the same cuTile scaled-cross-entropy backend. ## Testing Done - Hardware Type: NVIDIA H100 80GB HBM3 and NVIDIA B200 - [x] run `make test` to ensure correctness - equivalent command: `PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python3 -m pytest --disable-warnings --ignore=test/convergence --ignore=test/cutedsl test/` - result: `4096 passed, 929 skipped, 14 xfailed` - [x] run `make checkstyle` to ensure code style - result: repository-wide `ruff check .` and `ruff format --check .` passed - [ ] run `make test-convergence` to ensure convergence - not applicable: this change adds an opt-in operator implementation and does not patch a model - cuTile-targeted command: `LIGER_KERNEL_IMPL=cutile PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python3 -m pytest -q test/transformers/test_fused_linear_scaled_cross_entropy.py test/transformers/test_cutile_backend.py` - result: `29 passed` on both H100 and B200 H100 BF16 benchmark at `M=4096`, `H=4096`, `V=131072`: | Implementation | Forward | Backward | Full | Peak full memory | |---|---:|---:|---:|---:| | cuTile | 6.45 ms | 20.70 ms | 26.94 ms | 2.43 GiB | | CuTe SM90 | 7.26 ms | 19.81 ms | 27.74 ms | 2.43 GiB | B200 BF16 benchmark for the same shape, using automatic workspace selection: | Implementation | Forward | Backward | Full | Peak full memory | |---|---:|---:|---:|---:| | cuTile | 2.97 ms | 8.35 ms | 11.36 ms | 2.64 GiB | | Torch | 5.24 ms | 7.15 ms | 12.85 ms | 7.22 GiB | AI assistance was used. The submitter reviewed the changed lines and ran the validation above. --------- Co-authored-by: Copilot <[email protected]> |
||
|
|
a3f3654f93 |
Add fused linear scaled cross entropy frontend (#1354)
## Summary Adds a root-level `LigerFusedLinearScaledCrossEntropyFunction` frontend that: - dispatches exact NVIDIA SM90 inputs to the existing Hopper CuTe DSL implementation - uses a 512-token chunked PyTorch fallback adapted from Verl's fused linear PPO implementation on other devices and NVIDIA compute capabilities - preserves per-token FP32 NLL, optional differentiable entropy, temperature scaling, and `ignore_index` behavior across both paths - keeps the SM90 implementation isolated under `cutedsl` while exporting the frontend from `liger_kernel.ops` Also adds CPU fallback coverage and direct SM90-vs-fallback parity coverage for outputs and gradients across a fallback chunk boundary. ## Testing Done - Hardware Type: NVIDIA H100 80GB HBM3 (SM90) - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence - `python -m pytest -q test/transformers/test_fused_linear_scaled_cross_entropy.py test/cutedsl/test_fused_scaled_cross_entropy_sm90.py::test_sm90_frontend_matches_fused_linear_ppo_fallback` — 10 passed Co-authored-by: jkolehma <[email protected]> Co-authored-by: Copilot <[email protected]> |
||
|
|
073ae3283a |
Add Hopper CuTe DSL fused scaled cross entropy (#1334)
## Summary Adds `LigerFusedScaledCrossEntropySM90Function`, a Hopper-only BF16 CuTe DSL operator for Verl-style fused vocabulary projection, per-token NLL, and optional differentiable token entropy. - supports `temperature` with Verl-compatible `logits / temperature` semantics - returns per-token FP32 NLL and optional BF16 entropy; reductions stay in PyTorch - zeros NLL, entropy, and their gradient contributions on `ignore_index` rows - keeps the existing Triton `LigerFusedLinearCrossEntropyFunction` unchanged - uses a fixed 1024-token backward wave to bound the BF16 dZ workspace ## Kernel architecture ### Forward - 384 threads / 3 warp groups - WG0 + WG1 cooperatively issue `M128 x N256 x K64` WGMMA, split M64 per WG - warp 0 of WG0 also issues the 3-stage TMA load pipeline; `wait_group(1)` overlaps TMA refills with live WGMMA - WG2 owns one token row per thread and consumes a padded FP16 SMEM logits tile - NLL and entropy are computed in one online vocabulary pass using running max, exp sum, and exp-weighted logit sum; logits are never stored to HBM - cluster-16 DSMEM reduction combines vocabulary partitions - `m_tiles_per_cluster=1` selects the fastest M-outer schedule; larger values retain multiple online-softmax states for shape-specific tuning ### Backward - one persistent CUDA kernel with a device-side wave loop; no Python phase loop - at most one CTA per SM, launched as 66 cluster-2 pairs / 132 CTAs on H200 - one phase-serial ~224 KiB SMEM arena and one shared FP32 WGMMA accumulator - dZ: cooperative self-TMA `N256` recomputation with the M-fast static scheduler - dX: cooperative `N256/K64/s3`, with cluster-2 W multicast - dW: transposed TMA layout equivalent to `Step<_2,_1,_3>` / `order=(1,0,2)` - one monotone HBM atomic counter provides GPU-wide barriers after dZ and between reusable waves - wave zero plain-stores BF16 dW; later waves use BF16 TMA reduce-add - all TMA descriptors are prefetched at kernel entry ## Correctness - `22 passed` on an H200 SM90: aligned/ragged shapes, temperature, differentiable entropy, direct dependency-free Verl formula parity, ignored rows, large-logit overflow guard, sparse upstream gradients, multi-wave BF16 dW accumulation, exports and guards - `22 skipped` on a host without an SM90 GPU via a module-wide pytest gate - `make checkstyle` passes - Torch, Liger Triton, and CuTe benchmark providers smoke-tested ## H200 performance vs Verl BF16, temperature 1.0, entropy returned but not included in the scalar loss. Both full operations execute four logical GEMMs. `cold` is one launch after 3.5 s idle; `hot` is the median after five back-to-back warmups. Cells are `milliseconds / effective TFLOPS`. | M | H | V | CuTe cold | Verl cold | cold speedup | CuTe hot | Verl hot | hot speedup | |---:|---:|---:|---:|---:|---:|---:|---:|---:| | 512 | 4096 | 131072 | 4.78 / 460 | 6.97 / 315 | **1.46x** | 4.36 / 504 | 6.72 / 327 | **1.54x** | | 1024 | 4096 | 131072 | 7.72 / 570 | 13.27 / 331 | **1.72x** | 7.64 / 575 | 13.18 / 334 | **1.72x** | | 2048 | 4096 | 131072 | 13.21 / 666 | 25.77 / 341 | **1.95x** | 14.25 / 617 | 25.93 / 339 | **1.82x** | | 4096 | 1024 | 131072 | 11.26 / 391 | 29.11 / 151 | **2.59x** | 10.83 / 406 | 28.93 / 152 | **2.67x** | | 4096 | 2048 | 131072 | 15.03 / 585 | 36.39 / 242 | **2.42x** | 15.99 / 550 | 36.33 / 242 | **2.27x** | | 4096 | 4096 | 32768 | 6.82 / 645 | 13.57 / 324 | **1.99x** | 7.26 / 606 | 13.47 / 327 | **1.85x** | | 4096 | 4096 | 65536 | 12.73 / 691 | 25.96 / 339 | **2.04x** | 14.01 / 628 | 26.08 / 337 | **1.86x** | | 4096 | 4096 | 131072 | 24.87 / 707 | 50.91 / 346 | **2.05x** | 28.15 / 625 | 51.73 / 340 | **1.84x** | | 4096 | 4096 | 262144 | 49.86 / 706 | 100.13 / 351 | **2.01x** | 56.16 / 626 | 102.10 / 345 | **1.82x** | | 32000 | 4096 | 131072 | 218.71 / 628 | 398.56 / 345 | **1.82x** | 229.49 / 599 | 408.01 / 337 | **1.78x** | ## Peak incremental memory | M | H | V | CuTe MiB | Verl MiB | saving | |---:|---:|---:|---:|---:|---:| | 512 | 4096 | 131072 | 1156 | 2952 | **61%** | | 1024 | 4096 | 131072 | 1288 | 3984 | **68%** | | 2048 | 4096 | 131072 | 1296 | 3992 | **68%** | | 4096 | 1024 | 131072 | 520 | 1801 | **71%** | | 4096 | 2048 | 131072 | 784 | 2452 | **68%** | | 4096 | 4096 | 32768 | 352 | 1032 | **66%** | | 4096 | 4096 | 65536 | 672 | 2024 | **67%** | | 4096 | 4096 | 131072 | 1312 | 4008 | **67%** | | 4096 | 4096 | 262144 | 2592 | 7976 | **68%** | | 32000 | 4096 | 131072 | 1531 | 4226 | **64%** | --------- Co-authored-by: Copilot <[email protected]> |
||
|
|
72a4ed47a5 |
Add CUTLASS CuTe DSL RMSNorm kernel (#1299)
# Add CUTLASS CuTe DSL RMSNorm kernel This PR adds a genuine **NVIDIA CUTLASS CuTe DSL** (`import cutlass.cute`) RMSNorm implementation to the existing `cutedsl` backend introduced in #1279. It is opt-in only and selected with `LIGER_KERNEL_IMPL=cutedsl`; any op without a CuTe DSL kernel transparently falls back to the default Triton kernel, so selecting the backend is always safe. The backend provides forward and backward implementations compatible with the existing `LigerRMSNormFunction` interface (same public `rms_norm_forward` / `rms_norm_backward` / autograd `Function` signatures). ## Feature coverage - FP32 and BF16 (FP16 also covered by tests) - Llama, Gemma, and no-casting modes - Configurable weight offset - Affine and non-affine RMSNorm - In-place and out-of-place backward - Regular and non-power-of-two (irregular) hidden dimensions - Transformer module and functional API routing - Defensive DTensor detection for compatible Torch versions ## This PR also adds - RMSNorm registration on the existing CuTe DSL export surface (extends #1279; cross-entropy exports are preserved) - A Triton-vs-CuTe DSL benchmark runner for RMSNorm (`run_cutedsl_compare.py`) - Subprocess-isolated backend selection so Triton and CuTe DSL registration state can't leak between providers - Per-provider benchmark CSV routing (reuses the repo's existing `LIGER_BENCH_TARGET` / `LIGER_BENCH_PROVIDER_TAG` hooks) - A small Torch-compatibility fix to the upstream RMSNorm DTensor check (guards `torch.distributed.tensor.DTensor` on Torch builds that don't eagerly import that submodule) > **Companion PR (separate):** grouped-bar benchmark visualization, source-specific > plot filenames, and one-command generation of all metric/sweep plots live in a > separate visualizer PR to keep this kernel PR narrowly scoped and reviewable. ## Kernel design (genuine `cutlass.cute`) The kernels are written for correctness and clarity first (fp32 accumulation; one CTA per row for the reductions; the weight gradient reduced in parallel across row strips) and keep the hidden dimension fully dynamic, so a single compiled kernel serves every shape — regular or irregular. ### Forward - One CTA (thread block) per input row; 8 warps (256 threads) cooperate on the row. - Two-pass algorithm: 1. Sum of squares in FP32 → per-warp butterfly-shuffle reduction → cross-warp reduction via shared memory → `rstd = rsqrt(mean_square + eps)`, cached to an FP32 `RSTD` buffer for reuse in backward. 2. Normalize (`x * rstd`) and apply the optional affine weight (`offset + w`). - Per-column tail is bounds-guarded, so irregular hidden dimensions are supported without a separate kernel. - Casting mode and the affine flag are compile-time (`cutlass.Constexpr`) specializations. Llama round-trips the normalized value to the input dtype before the affine multiply; Gemma and no-casting stay in FP32 — matching Triton semantics. ### Backward - **dX kernel:** one CTA per row. Computes `dot(m, x)` with `m = dy * (w + offset)` via the same warp + cross-warp reduction, then `dx = rstd * (m - (1/N) * rstd^2 * dot(m, x) * x)`. Supports reusing `dY` storage for `dX` when `in_place=True` (each column is read then written by its single owning thread, fenced by the reduction barrier). - **dW kernel:** the weight-gradient reduction is parallelized across row strips with a 2-D grid (column-blocks × row-strips). Each strip (one per SM, capped at `n_rows`) reduces `dy * (x * rstd)` over just its slice of rows in FP32 and writes a partial `dW` row; the host sums the partials — the same strategy as the Triton backward's `sm_count` partials + `_dW.sum(0)`. `dW` is computed **before** `dX` so the in-place `dX` write can't clobber the `dY` the reduction reads. Skipped entirely for non-affine RMSNorm. ### Host - Torch tensors are marshaled via DLPack; kernels are compiled once with `cute.compile` and cached, keyed on **every** baked dtype (X, W, dY) plus casting mode and the affine flag, then launched on Torch's current CUDA stream. ## Correctness Ran the CuTe DSL routing and RMSNorm correctness suites on an NVIDIA GPU with `nvidia-cutlass-dsl` installed: ```bash LIGER_KERNEL_IMPL=cutedsl pytest test/transformers/test_cutedsl_rms_norm.py -v LIGER_KERNEL_IMPL=cutedsl pytest test/transformers/test_cutedsl_backend.py -v pytest test/transformers/test_cutedsl_cross_entropy.py -v # upstream regression guard make checkstyle ``` Result: **158 passed, 1 skipped.** The single skip is expected — it is the cross-entropy backend-routing test that only runs when `LIGER_KERNEL_IMPL=cutedsl` is set, which is intentionally not set for the cross-entropy regression-guard command above. `make checkstyle` (ruff check + format) passes clean. Coverage includes: - Explicit `LIGER_KERNEL_IMPL=cutedsl` routing validation (ops + transformer modules) - Forward parity against the default Triton RMSNorm reference - Input-gradient parity - Affine weight-gradient parity - FP32 and BF16 (FP16 in the parity matrix) - Llama, Gemma, and no-casting modes - Affine and non-affine configurations - In-place and out-of-place backward - Standard and irregular shapes (representative regular dims plus non-power-of-two hidden dimensions) - A same-process mixed weight-dtype case that guards the compiled-kernel cache key - Upstream CuTe DSL cross-entropy tests preserved as a backend export/import regression guard ## Benchmarks: CuTe DSL vs Triton <!-- The tables below are placeholders. Fill them with the genuine medians from YOUR H100/B200 runs — do NOT reuse any pre-existing CSV. On each GPU node, extract the full-pass table straight from the CSV this branch's runner produces: python - <<'PY' import pandas as pd f = "data/all_benchmark_data_cutedsl_rms_norm_h100.csv" # or _b200.csv df = pd.read_csv(f) d = df[(df.kernel_name=="rms_norm") & (df.metric_name=="speed") & (df.kernel_operation_mode=="full")] piv = d.pivot_table(index="x_value", columns="kernel_provider", values="y_value_50") piv["speedup"] = piv["liger_triton"] / piv["liger_cutedsl"] print(piv.round(4)) PY Repeat with metric_name=="memory" for the peak-memory comparison. --> Measurements use the median (`y_value_50`) from `triton.testing.do_bench`. Ratio = Triton / CuTe DSL; **values below 1.00x mean CuTe DSL is slower than Triton.** `full` = forward plus backward. **Benchmark environment** - GPU: NVIDIA H100 80GB HBM3 and NVIDIA B200 - Torch: 2.12.0.1+cu130 - CUDA runtime: 13.0 - Triton: 3.7.0 - nvidia-cutlass-dsl: 4.5.2 - Input dtype: BF16 The token-length sweep uses B=1, H=4096, and T = {1024, 2048, 4096, 8192}. The model-configuration sweep uses B=1, T=2048, and the hidden dimension from each model configuration. > **Honest performance note.** This is a correctness-first implementation (fp32 > accumulation, one CTA per row for the reductions, no vectorization/pipelining). The > weight-gradient (dW) reduction is parallelized across row strips (one strip per SM, > partials summed on the host — the same strategy as the Triton backward). The three-way > picture (H100 and B200 are the same shape): > > - **vs HuggingFace eager:** CuTe DSL is **faster in the forward pass in every case** > (11/11 on both GPUs), and faster on the **full pass in 9/11** cases on both GPUs and > the **backward in 9/11 (H100) / 11/11 (B200)** — a clear win over eager. The only > losses to HuggingFace are at the smallest size (T=1024), where the op is tiny and > launch-bound. > - **vs the existing Triton Liger kernel:** CuTe DSL is **not yet faster in any measured > case** (~0.5–0.8x), but after the parallel-dW backward it is now **stable across > sequence length** rather than collapsing — full-pass ratio holds around 0.54–0.57x > as T grows from 1k to 8k (it was 0.59x→0.09x before the optimization). > - **Peak memory:** identical to Triton to the MB, and ~4–5x below HuggingFace. > > Remaining headroom vs Triton (fusing dX/dW to load x/dy once, vectorized loads) is > deferred to a `liger-kernel-perf` follow-up. The `Ratio` columns below are > Triton/CuTeDSL, so values below 1.00x mean CuTe DSL is slower than Triton specifically. ### NVIDIA H100 80GB HBM3 Full pass (fwd+bwd) — token-length sweep (B=1, H=4096) | T | CuTe DSL (ms) | Triton (ms) | Ratio (Triton/CuTeDSL) | |---|---|---|---| | 1,024 | 0.4688 | 0.2554 | 0.54x | | 2,048 | 0.4641 | 0.2589 | 0.56x | | 4,096 | 0.4686 | 0.2549 | 0.54x | | 8,192 | 0.4762 | 0.2604 | 0.55x | Full pass (fwd+bwd) — model-config sweep (B=1, T=2048) | Model | CuTe DSL (ms) | Triton (ms) | Ratio (Triton/CuTeDSL) | |---|---|---|---| | Llama 2 7B | 0.3825 | 0.3155 | 0.82x | | Llama 3 8B | 0.3867 | 0.3122 | 0.81x | | Qwen 2.5 7B | 0.3842 | 0.3141 | 0.82x | | Qwen 2.5 14B | 0.3858 | 0.3160 | 0.82x | | Qwen 2.5 72B | 0.3828 | 0.3135 | 0.82x | | DeepSeek V2 Lite | 0.3791 | 0.3137 | 0.83x | | DeepSeek V3 | 0.3822 | 0.2577 | 0.67x | Geometric-mean ratio (Triton/CuTeDSL), H100: forward 0.47x; backward 0.49x; full 0.70x. The parallel-dW backward lifted the full-pass geomean from ~0.27x (token) to 0.70x and made it flat across sequence length. CuTe DSL beats HuggingFace eager in 11/11 forward, 9/11 backward, and 9/11 full-pass cases. Peak memory — full pass, token-length sweep (MB) | T | CuTe DSL | Triton | HuggingFace | |---|---|---|---| | 1,024 | 34.1 | 34.1 | 152.0 | | 2,048 | 66.1 | 66.1 | 304.0 | | 4,096 | 130.1 | 130.1 | 608.0 | | 8,192 | 258.1 | 258.1 | 1216.1 | CuTe DSL peak memory is identical to Triton (the dW strip partials match Triton's `_dW` buffer) and ~4–5x below HuggingFace eager across both sweeps. ### NVIDIA B200 Full pass (fwd+bwd) — token-length sweep (B=1, H=4096) | T | CuTe DSL (ms) | Triton (ms) | Ratio (Triton/CuTeDSL) | |---|---|---|---| | 1,024 | 0.2859 | 0.1598 | 0.56x | | 2,048 | 0.2846 | 0.1593 | 0.56x | | 4,096 | 0.2840 | 0.1574 | 0.55x | | 8,192 | 0.3298 | 0.1894 | 0.57x | Full pass (fwd+bwd) — model-config sweep (B=1, T=2048) | Model | CuTe DSL (ms) | Triton (ms) | Ratio (Triton/CuTeDSL) | |---|---|---|---| | Llama 2 7B | 0.2867 | 0.1635 | 0.57x | | Llama 3 8B | 0.2870 | 0.1615 | 0.56x | | Qwen 2.5 7B | 0.2864 | 0.1629 | 0.57x | | Qwen 2.5 14B | 0.2839 | 0.1621 | 0.57x | | Qwen 2.5 72B | 0.2857 | 0.1589 | 0.56x | | DeepSeek V2 Lite | 0.2857 | 0.1621 | 0.57x | | DeepSeek V3 | 0.2868 | 0.1613 | 0.56x | Geometric-mean ratio (Triton/CuTeDSL), B200: forward 0.53x; backward 0.41x; full 0.56x. Same shape as H100 and just as stable across sequence length after the parallel-dW backward. CuTe DSL beats HuggingFace eager in 11/11 forward, 11/11 backward, and 9/11 full-pass cases. Peak memory — full pass, token-length sweep (MB) | T | CuTe DSL | Triton | HuggingFace | |---|---|---|---| | 1,024 | 34.4 | 34.4 | 152.0 | | 2,048 | 66.4 | 66.4 | 304.0 | | 4,096 | 130.4 | 130.4 | 608.0 | | 8,192 | 258.4 | 258.4 | 1216.1 | As on H100, CuTe DSL peak memory is identical to Triton and ~4–5x below HuggingFace eager across both sweeps. ## Usage Install the CuTe DSL dependency and enable the backend: ```bash pip install "liger-kernel[cutedsl]" # nvidia-cutlass-dsl LIGER_KERNEL_IMPL=cutedsl python your_script.py ``` Run both Triton and CuTe DSL RMSNorm benchmarks into one CSV per GPU: ```bash cd benchmark/scripts python run_cutedsl_compare.py --kernel rms_norm --source h100 --overwrite python run_cutedsl_compare.py --kernel rms_norm --source h100 --sweep-mode model_config --overwrite # (repeat on the B200 node with --source b200) ``` Generate the plots (grouped bars — companion visualizer PR): ```bash cd benchmark python benchmarks_visualizer.py --kernel-name rms_norm --metric-name speed \ --sweep-mode token_length --data-file data/all_benchmark_data_cutedsl_rms_norm_h100.csv --overwrite ``` ## Known limitations - `row_mode` is accepted for interface compatibility but is not currently specialized by the CuTe DSL implementation. - The kernels are correctness-first (FP32 accumulation, one CTA per row for the reductions, strip-parallel dW). Remaining on-GPU tuning (fusing dX/dW to load x/dy once, vectorized/pipelined loads) is left to a follow-up using the `liger-kernel-perf` workflow. - No-casting mode accumulates in FP32 (vs the Triton kernel's in-dtype path); this is an intentional, tolerance-covered deviation. - **Performance does not yet match the existing Triton kernel.** On both H100 and B200, CuTe DSL is ~0.5–0.8x of Triton on the full pass — slower, but now **stable across sequence length** (it no longer collapses at large T) thanks to the strip-parallel dW backward. It beats HuggingFace eager in the forward pass everywhere and on the full pass in 9/11 cases on both GPUs; peak memory is identical to Triton. Closing the remaining gap to Triton (fused dX/dW, vectorized loads) is a `liger-kernel-perf` follow-up. ## Branch state - Branched directly from the latest `upstream/main`; no old prototype commits were cherry-picked. - No `import cuda.tile` anywhere under `src/liger_kernel/ops/cutedsl/`; the RMSNorm kernel uses `cutlass.cute`. - The existing CuTe DSL cross-entropy implementation and the upstream cuTile files are unchanged. - No tracked benchmark CSV is added or deleted; the genuine H100/B200 result CSVs are git-ignored and are not part of this PR. - Files changed vs `upstream/main`: `README.md`, `benchmark/scripts/run_cutedsl_compare.py`, `src/liger_kernel/ops/cutedsl/ops/__init__.py`, `src/liger_kernel/ops/cutedsl/ops/rms_norm.py`, `src/liger_kernel/ops/rms_norm.py`, `test/transformers/test_cutedsl_backend.py`, `test/transformers/test_cutedsl_rms_norm.py`. --------- Co-authored-by: Kenechukwu Enemuo <[email protected]> Co-authored-by: Copilot <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
|
7dc5417b3b |
[NPU] Bump Ascend NPU deps to torch/torch_npu 2.7.1 and triton-ascend 3.2.1 (#1262)
## Summary - Bump NPU default dependencies in `setup.py` from `torch==2.6.0` / `torch_npu==2.6.0` / `triton-ascend==3.2.0` to `torch==2.7.1` / `torch_npu==2.7.1` / `triton-ascend==3.2.1`. - Remove outdated TODO comments that blocked the 2.7.1 / 3.2.1 upgrade. - Add an **Ascend NPU** subsection at the end of the Installation section in `README.md`, documenting the source install command with `--extra-index-url` (required because `triton-ascend` is not on the default PyPI index). ## Motivation Align Liger-Kernel with the current Ascend NPU stack (torch 2.7.1 + triton-ascend 3.2.1) while keeping the existing `pip install -e ".[dev]"` workflow. The extra index URL must be passed on the same pip invocation so dependency resolution can find `triton-ascend`. On an Ascend NPU machine, run: ```bash pip install -e ".[dev]" --extra-index-url https://triton-ascend.osinfra.cn/pypi/simple ``` ## Testing Done All test cases have been fully verified in: - https://github.com/linkedin/Liger-Kernel/pull/1259 Hardware Type: All NPUs. - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence Co-authored-by: sunyi0505 <[email protected]> |
||
|
|
9e075c9917 |
[Feature] Support apply_liger_kernel_to_deepseek_v4 & unit tests (#1245)
## Summary Following the work in https://github.com/linkedin/Liger-Kernel/pull/1221, `Liger-Kernel` now supports `transformers v5.8.0`, which includes the DeepSeek v4 model—widely used by the community. This PR therefore requests `Liger-Kernel` to support this model, along with the corresponding unit tests. Note: - `transformers` gained full support for DeepSeek v4 in v5.8.1: https://github.com/huggingface/transformers/releases/tag/v5.8.1, so here we're validating against v5.9.0. We've tested that `Liger-Kernel` is also compatible. - `Liger-Kernel` requires `transformers` v4.52.0 or later, so for lower versions we should skip the test cases to avoid disrupting the rest of the test suite. ## Testing Done ### transformers=5.9.0 <img width="1384" height="581" alt="5-9-0" src="https://github.com/user-attachments/assets/b8704be4-2a77-4606-9b56-5ea286563165" /> ### transformers=4.52.0 <img width="1381" height="591" alt="4-52-0" src="https://github.com/user-attachments/assets/04594e89-a2e4-443a-9ab5-d082515aa60a" /> - Hardware Type: RTX 5090, A100-PCIE-40G - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
|
94236ea193 |
Improve ROCm source install extras (#1243)
## Summary Update ROCm source install instructions to install ROCm PyTorch wheels before installing Liger, so `pip install -e .` and `pip install -e ".[dev]"` work without replacing the ROCm torch stack. Align AMD CI with the same install order and ROCm wheel index. <!--- ## 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 |
||
|
|
8ef203a42d |
Add cutile jsd (#1228)
## Summary Add optional cuTile backend support for Liger Kernel. This change introduces `LIGER_KERNEL_BACKEND=cutile` backend selection, CUDA-only validation, cuTile JSD operator replacement, optional `cutile` package extras, benchmark labeling as `liger-cutile`, and documentation for enabling cuTile. Related to https://github.com/linkedin/Liger-Kernel/issues/1205 ## Details cuTile is only enabled explicitly through `LIGER_KERNEL_BACKEND=cutile`. Non-CUDA devices raise an error, and missing cuTile dependencies raise an import error with installation instructions. ## Testing Done - Ran JSD correctness tests with cuTile enabled: `LIGER_KERNEL_BACKEND=cutile pytest -s test/transformers/test_jsd.py -v` - Added backend selection test: `test/transformers/test_cutile_backend.py` - Ran targeted ruff checks on modified files. - Ran JSD benchmarks comparing `liger` vs `liger-cutile` on B300, will add B200 in the later commit - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence <img width="1239" height="354" alt="image" src="https://github.com/user-attachments/assets/3abeacdc-dd48-45e3-967f-09f04d254f12" /> --------- Co-authored-by: Vaibhav Jindal <[email protected]> Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]> |
||
|
|
decb1b73d7 |
[Gemma 4] Add multimodal support (apply_liger_kernel_to_gemma4 for Gemma4ForConditionalGeneration) (#1203)
## Summary Follow-up to #1196 (the text path) — adds `apply_liger_kernel_to_gemma4` for `Gemma4ForConditionalGeneration` (multimodal class; includes E2B / E4B / E4B-it which are loaded by `AutoModelForCausalLM` as the multimodal class even when only text is being trained). Closes the multimodal half of #1186. ### Why Gemma 4's text vocab is 262,144. Without FLCE the (B, T, V) bf16 logits tensor is ~17 GB at T=8192 (and ~34 GB once the loss path upcasts to fp32 for cross-entropy), which OOMs even 96 GB cards on `Gemma4ForConditionalGeneration` SFT — the OOM that originally motivated #1186. Routing loss through `LigerForCausalLMLoss` materializes only the loss scalar. ### Shape A single unified entry point that dispatches on class, per @Mecoli1219's preference in #1186: - `apply_liger_kernel_to_gemma4(model=Gemma4ForConditionalGeneration_instance)` — multimodal path. Class-level RMSNorm + GeGLU swaps via `apply_liger_kernel_to_gemma4_text`, FLCE forward via the new `multimodal_forward`, recurses into `model.model.language_model` for instance-level patches. - `apply_liger_kernel_to_gemma4(model=Gemma4ForCausalLM_instance)` — routes to `apply_liger_kernel_to_gemma4_text` for backwards compatibility, so the same entry point works for either shape. - Registry: adds `"gemma4": apply_liger_kernel_to_gemma4` alongside the existing `"gemma4_text"`. ### Drive-by fixes Two `isinstance(model, tuple_with_None_filter)` sites (one in our new dispatcher, one in the existing `apply_liger_kernel_to_gemma4_text`) raised `TypeError: isinstance() arg 2 must be a type, a tuple of types, or a union` when called under `with patch("transformers.models.gemma4.modeling_gemma4")`: `getattr(MagicMock_module, "Gemma4TextForCausalLM", None)` returns a MagicMock (not `None`) because `MagicMock` auto-creates attributes, so the `cls is not None` filter let it slip into the `isinstance` tuple. The text-path version was dormant — its existing test passes a `Gemma4ForCausalLM` which short-circuits the `isinstance` match before reaching the bad entry — but the multimodal recursive call into the text path passes a `Gemma4TextModel`, so no early match. Both sites now use `isinstance(cls, type)` as the filter. ### Out of scope (deferred) - **Vision / audio tower kernel swaps.** Gemma 4 loads both via `AutoModel.from_config(config.{vision,audio}_config)`, so the module classes are polymorphic. Out of scope here; FLCE on the LM head is what unblocks training OOM. - **`Gemma4MultimodalEmbedder` / projector norms.** Analogous to gemma3's `mm_soft_emb_norm` patching. Skipped for the same minimal-surface reason. - **PLE (Per-Layer Embeddings).** State passes through the inner forward unchanged; verified end-to-end on `google/gemma-4-E4B-it`. - **MoE experts (`Gemma4TextExperts`).** Guarded out via the same `enable_moe_block` check used by the text path in #1196. - **Multimodal mini convergence test (`mini_gemma4`).** Would need new `test/resources/fake_configs/Google/Gemma4/...` scaffolding for an image / audio processor — PR #1196 followed the same pattern (only added `mini_gemma4_text` to the non-multimodal convergence files). Happy to add it as a follow-up if you'd prefer it bundled here — let me know. ## Testing Done ### Hardware - **Hardware Type: RTX 5090 (Blackwell, sm_120, 32 GB), Vast.ai instance** - CUDA 13.0, NVIDIA driver 590.48.01 - Python 3.10.12, torch 2.11.0+cu130, transformers 5.7.0.dev0 (built from `huggingface/transformers@main` — gemma4 requires ≥ 5.5.0) - bf16 throughout for end-to-end numerics ### End-to-end numerical equivalence on real `google/gemma-4-E4B-it` Verified before authoring this PR with our internal `verify_patch_equivalence.py` (same shape as #1196's verification): | Axis | Threshold | Measured | | --- | --- | --- | | Fused-CE vs reference loss diff (seq=512) | `< 5e-3` | `0.0016` | | End-to-end top-1 token agreement (seq=256, patched vs stock SDPA) | `> 99 %` | **>99 %** | | End-to-end loss diff (patched vs stock SDPA) | `< 5e-3` | `0.0016` | | Logit-distribution KL (informational) | — | `~3.5e-2` | | SDPA EFFICIENT_ATTENTION vs MATH cos-sim (seq=4096) | `> 0.9999` | **>0.9999** | ### Liger-Kernel test gates - [x] `make checkstyle` — `All checks passed!`, 267 files already formatted - [x] `make test` — see log - [x] `make test-convergence` — see logs (per file) #### `make test` `3131 passed, 903 skipped, 12 xfailed, 3 failed` in 35:42. Our two new unit tests pass cleanly: ``` PASSED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text (0.70s) PASSED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_conditional_generation (0.22s) ``` All six `LigerGEGLUMLPForGemma4` edge-case tests added by #1196 also pass. The 3 unrelated failures are pre-existing on the parent branch (untouched by this PR) and reproduce on `main`: ``` FAILED test/transformers/test_grpo_loss.py::test_grpo_loss_with_bias_correction_kl[...] FAILED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_paligemma FAILED test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma3_conditional_generation ``` #### `make test-convergence` (per file) | File | Result | | --- | --- | | `fp32/test_mini_models.py` | 32 passed, 3 skipped | | `fp32/test_mini_models_multimodal.py` | 8 passed, 4 skipped, 1 xfailed, **1 failed** (`mini_qwen2_vl`) | | `fp32/test_mini_models_with_logits.py` | 29 passed, 3 skipped, 1 xpassed | | `bf16/test_mini_models.py` | 33 passed, **2 failed** (`mini_llama4`, `mini_gemma4_text`) | | `bf16/test_mini_models_multimodal.py` | 11 passed, 1 skipped, **2 failed** (`mini_qwen2_vl`, `mini_llama4`) | | `bf16/test_mini_models_with_logits.py` | 30 passed, 1 skipped, **2 failed** (`mini_llama4`, `mini_qwen3_moe`) | `mini_gemma4_text` passes in `bf16/test_mini_models_with_logits.py` and `fp32/test_mini_models.py` (PR #1196's text path). It fails only in `bf16/test_mini_models.py` with the same Blackwell bf16 logprob drift @eqy and reviewers flagged on #1196 (review comment `r4321013177`); this is in PR #1196's territory, not introduced by this multimodal patch. The other failures (`mini_llama4`, `mini_qwen2_vl`, `mini_qwen3_moe`) are also in models we don't touch and are pre-existing test instability on consumer Blackwell. cc @Mecoli1219 @lardinator @ruilin-gif --- 🤖 Drafted with [Claude Code](https://claude.com/claude-code) (Claude Opus 4.7), reviewed and posted by me. --------- Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]> |
||
|
|
30b8486a2b |
Add fused modulated RMSNorm (#1225)
## Summary Adds a fused Modulated RMSNorm kernel for the common pattern: ```python y = rms_norm(x, weight, eps) * (1 + scale) + shift ``` This is intended for FiLM/AdaRMSNorm-style modulation where RMSNorm output is scaled and optionally shifted by per-hidden-dimension modulation tensors. Related issue: #1224 ## Changes - Add `LigerModulatedRMSNormFunction` with fused forward/backward Triton kernels. - Add `LigerModulatedRMSNorm` module wrapper and `liger_modulated_rms_norm` functional API. - Support optional RMSNorm weight, optional shift, and existing RMSNorm casting modes: `llama`, `gemma`, `none`. - Support modulation shapes that flatten to `(hidden,)`, `(batch, hidden)`, or row-wise `(batch, seq, hidden)` when rows divide the hidden-state rows. - Add correctness tests across fp32/bf16, casting modes, modulation broadcast modes, optional shift, and optional affine weight. - Add benchmark script for speed and memory comparison. ## Benchmark Environment: NVIDIA GeForce RTX 3090, bf16, hidden size 4096, row-wise `scale` and `shift`. ### Speed, p50 ms | BT | Mode | Fused | LigerRMSNorm + torch mod | Speedup | |---:|---|---:|---:|---:| | 1024 | forward | 0.0430 | 0.1091 | 2.5x | | 1024 | backward | 0.0763 | 0.1403 | 1.8x | | 1024 | full | 0.1403 | 0.2642 | 1.9x | | 4096 | forward | 0.1546 | 0.3999 | 2.6x | | 4096 | backward | 0.2540 | 0.4895 | 1.9x | | 4096 | full | 0.4787 | 1.3793 | 2.9x | ### Memory, p50 MB | BT | Mode | Fused | LigerRMSNorm + torch mod | Savings | |---:|---|---:|---:|---:| | 1024 | forward | 32.02 | 56.02 | 42.8% | | 1024 | backward | 89.33 | 97.33 | 8.2% | | 1024 | full | 89.33 | 97.33 | 8.2% | | 4096 | forward | 128.03 | 224.03 | 42.9% | | 4096 | backward | 353.34 | 385.34 | 8.3% | | 4096 | full | 353.34 | 385.34 | 8.3% | |
||
|
|
3dd9e3a2b7 |
Extend LigerExperts patching to qwen3_vl_moe and glm4v_moe (#1192)
## Summary
- Add fused MoE expert (LigerExperts) monkey-patching to
**qwen3_vl_moe** and **glm4v_moe**, the two remaining MoE models that
were missing it
- Fix buggy glm4v_moe instance patching that was outside the `for` loop
with duplicate logic and setting `decoder_layer.mlp = None`
- Enable 3 previously-skipped qwen3_vl_moe convergence tests (bf16/fp32
with_logits, fp32 multimodal)
- All patching gated behind `IS_TRANSFORMERS_V5_OR_LATER` for backward
compatibility with transformers v4
### Models patched
| Model | Class-level patch | Instance-level patch |
|---|---|---|
| `qwen3_vl_moe` | `Qwen3VLMoeTextExperts = LigerExperts` |
Attribute-based detection of `experts` on MoE layers |
| `glm4v_moe` | `Glm4vMoeTextNaiveMoe = LigerExperts` | Attribute-based
detection of `experts` + `shared_experts` |
### Why not llama4?
`Llama4TextExperts.forward(hidden_states)` takes only 1 arg (routing
done externally in MoE block), incompatible with LigerExperts'
`forward(hidden_states,
top_k_index, top_k_weights)`.
### Compatibility
Verified across transformers v4.57.6, v5.0–v5.4, and v5.5.4. Class names
and forward signatures are stable across all v5 releases.
## Test plan
- [x] Instance patching tests (`test_monkey_patch.py`) for all 3
qwen3_vl_moe variants and glm4v_moe
- [x] Convergence tests (bf16 + fp32, FLCE + with_logits) for
qwen3_vl_moe and glm4v_moe
- [x] Enabled 3 previously-skipped qwen3_vl_moe convergence tests
- [x] Verified backward compatibility: on v4,
`IS_TRANSFORMERS_V5_OR_LATER` gates skip all LigerExperts code paths
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>
|
||
|
|
9e3f2618c3 |
[Gemma 4] Add apply_liger_kernel_to_gemma4_text (dense text, 31B-targeted) (#1196)
## Summary Adds Liger-Kernel support for Google's **Gemma 4** family (text, dense variants — primary target `google/gemma-4-31B`). Follows the exact integration pattern of the Gemma 3 port: module-level class swaps (`Gemma4RMSNorm`, `Gemma4TextMLP`), class-forward swap on `Gemma4ForCausalLM`, and `MODEL_TYPE_TO_APPLY_LIGER_FN["gemma4_text"]` registration so `AutoLigerKernelForCausalLM` routes transparently. ### What you get - `LigerRMSNormForGemma4` (ones-init, no `(1+w)` offset, fp32 compute, handles `with_scale=False` for `v_norm`) - `LigerGEGLUMLPForGemma4` (thin wrapper over `LigerGEGLUMLP` that absorbs `Gemma4TextMLP`'s `layer_idx` arg) - Fused-linear-CE `causal_forward` on both `Gemma4ForCausalLM` (stock HF) and `Gemma4TextForCausalLM` (some users extract this subclass to dodge HF issue #45200's mm_token_type_ids check) - Registration + exports from `liger_kernel.transformers` ### What you don't get (explicit non-goals) - Multimodal (`Gemma4ForConditionalGeneration`, `Gemma4VisionModel`, `Gemma4AudioModel`) - MoE (`Gemma4TextExperts`/`Gemma4TextRouter`, used by 26B-A4B) - PLE (Per-Layer Embeddings — used by E2B/E4B) - Double-wide MLP on KV-shared layers (31B doesn't use either) - **Rope kernel swap**: HF Gemma 4's `apply_rotary_pos_emb(x, cos, sin)` is single-tensor; `liger_rotary_pos_emb(q, k, cos, sin)` is dual-tensor. Signatures are incompatible. `rope=True` is a no-op with a warning. The large memory win (LCE) is unaffected. ## Headline: Peak HBM at seq_len=8192 Measured on **AMD MI250X (LUMI)** with a 6-layer Gemma-4-shaped mini model at full 31B vocab (262,144), `bf16`, batch=1, seq=8192, forward+backward: | Peak HBM | Value | |---|---| | HF baseline | **41.47 GB** | | Liger (patched) | **10.89 GB** | | **Saved** | **30.58 GB (73.7% reduction)** | Driver: fused-linear-CE (`skip_logits=True`) eliminates the 262,144 × 8192 × bf16 logits tensor (~4 GB materialized) **plus** its gradient, activation buffers, and softcap intermediates — combined ~30 GB freed. Loss parity on the same configuration: `loss_HF = 12.6841440`, `loss_liger = 12.6840858` (abs diff ~6e-5 — numerically identical). ## Numerical correctness ### fp32 (kernel correctness) Whole-model forward, 6-layer Gemma-4-shaped config, vocab 32000, shape `(2, 256, 32000)`: | statistic | value | |---|---| | max \|logits_liger − logits_hf\| | **2.55e-03** | | mean abs diff | 5.11e-05 | | p99 abs diff | 3.84e-04 | ### bf16 (expected dtype noise) | statistic | value | |---|---| | max abs diff | 7.97e-01 | | mean abs diff | 3.05e-02 | | p99 abs diff | 1.92e-01 | Interpretation: kernels are numerically correct (fp32 max 2.55e-3); bf16 drift is 6-layer dtype-inherent precision noise, well inside industry-standard ranges. ## Discovered during LUMI verification (fixed in-PR) Three plan-level assumptions that only showed up under real model construction — all fixed here: 1. **`v_norm` uses `with_scale=False`** — no weight parameter, so the naive subclass broke class-level swap. Fix: `LigerRMSNormForGemma4` accepts `with_scale` kwarg and delegates forward to a plain-torch path when weight is absent. 2. **`Gemma4TextMLP(config, layer_idx)` signature** (vs `LigerGEGLUMLP(config)`) — crashed at layer construction. Fix: `LigerGEGLUMLPForGemma4` wrapper absorbs the extra arg. 3. **`apply_rotary_pos_emb(x, cos, sin)` single-tensor signature** (vs `liger_rotary_pos_emb(q, k, cos, sin)`) — incompatible. Fix: `rope=True` becomes a no-op + warning for Gemma 4. Plus: `Gemma4RMSNorm` inherits `Gemma3nRMSNorm` (not `Gemma3RMSNorm`) — reusing `LigerRMSNormForGemma3` would silently diverge training via its `(1+w)` offset. New subclass is mandatory, not cosmetic. ## Testing Done - Hardware Type: AMD MI250X (LUMI HPC) - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence ### CPU-runnable (included in `make test`) \`\`\` $ python -m pytest test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text -v test/transformers/test_monkey_patch.py::test_apply_liger_kernel_to_instance_for_gemma4_text PASSED [100%] ======================== 1 passed, 1 warning in 21.27s ========================= \`\`\` ### checkstyle \`\`\` $ ruff check --output-format=concise . All checks passed! $ ruff format --check . 264 files already formatted \`\`\` ### LUMI (MI250X, ROCm 6.2.4, PyTorch 2.7.1, transformers 5.5.4) - [x] `pytest test/convergence/bf16/test_mini_models.py -k mini_gemma4_text` — MI250X ROCm - [x] `pytest test/convergence/bf16/test_mini_models_with_logits.py -k mini_gemma4_text` — MI250X ROCm - [x] Gemma 3 regression (3 tests) — all pass, no Gemma 3 behaviour change - [x] Revert helper functional test — identical outputs (`max diff 0.00e+00`) after revert - [x] HBM benchmark at seq_len=8192 — **30.58 GB saved**, **73.7% reduction** Tolerance selection for `mini_gemma4_text`: - `loss_atol=5e-2` (vs gemma3's `1e-2`) — 6-layer mini (minimum for sliding+global cycle) has ~0.05 bf16 drift vs gemma3's 4-layer which fits 1e-2 - `logprobs_atol=5e-1` — accumulated bf16 noise flips 3 of ~20,480 top-k logprob ranks on near-tie tokens --- 🤖 Drafted with [Claude Code](https://claude.com/claude-code) (Claude Opus 4.7). No auto-generation skill was used end-to-end — scaffolding and LUMI-verified fixes were developed and reviewed interactively. --------- Co-authored-by: amrtini <[email protected]> Co-authored-by: Claude Opus 4.7 <[email protected]> |
||
|
|
1c013e286c |
[Feature] Use Liger's Relu_Squared kernel for Nemotron models (#1176)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Use relu_squared in nemotron. This PR is generated using the liger-autopatch skill and tests the changes in #1177 . <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> Class patching and instance patching of relu_squared function. ## 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: H100 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
|
8e6232a08f |
Add Liger Kernel support for Ministral models (#1166)
This PR was generated automatically by Claude Code with minimal human intervention. Claude uses the skill written in #1167 to generate this PR and test it out. Ministral (model_type: "ministral") is a separate model type in transformers v5+ with its own MinistralForCausalLM class. This adds full Liger Kernel optimization support including RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, and FusedLinearCrossEntropy. ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> - Adds full Liger Kernel optimization for Ministral models (model_type: "ministral") in transformers v5+ - Patches: RoPE, RMSNorm, SwiGLU, CrossEntropyLoss, FusedLinearCrossEntropy - All 5 tests (instance patching + 4 convergence) validated on H100 GPU <!--- ## 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> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
|
eac9ce6ede |
Add Liger Kernel support for nemotron models (#1165)
This PR is generated by Claude code using the skill in #1167. ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> - Adds monkey patching support for NVIDIA Nemotron models with apply_liger_kernel_to_nemotron - Supports FusedLinearCrossEntropy (default) and CrossEntropyLoss optimizations - RoPE, MLP, and LayerNorm patching are intentionally excluded due to Nemotron's unique architecture <!--- ## 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. ---> - Instance patching test passes - All 4 convergence tests pass on H100 GPU (bf16/fp32 x FLCE/with_logits) - Lint passes <!-- 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> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 (1M context) <[email protected]> |
||
|
|
a01a0020b4 |
Add support for Qwen3.5 dense models (#1123)
## Summary Adds Liger Kernel support for the Qwen3.5 dense model family (0.8B / 2B / 4B / 9B / 27B). Closes #1119 Qwen3.5 uses the same hybrid GDN + full-attention architecture as Qwen3 Next but is a distinct model type (`qwen3_5`) in transformers v5.2+. This PR adds a dedicated `apply_liger_kernel_to_qwen3_5` entry point so auto-patching works correctly. Patched kernels: - **RMSNorm** — Gemma-style offset RMSNorm (`LigerRMSNormForQwen3Next`, `offset=1.0, casting_mode="gemma"`) - **SwiGLU** — Fused SwiGLU MLP (`LigerQwen3MoeSwiGLUMLP`, matching Qwen3.5's `__init__` signature) - **Fused Linear Cross Entropy** — avoids materializing the full logits tensor (vocab_size=248,320) - **RoPE** — disabled (raises `NotImplementedError`); only 1 in 4 layers uses RoPE via full attention ## Details **New file:** - `src/liger_kernel/transformers/model/qwen3_5.py` — `lce_forward` for `Qwen3_5ForCausalLM` **Modified files:** - `monkey_patch.py` — `apply_liger_kernel_to_qwen3_5` with class-level and instance-level patching; registered as `qwen3_5` in `MODEL_TYPE_TO_APPLY_LIGER_FN` - `__init__.py` — export - `test/utils.py` — `revert_liger_kernel_to_qwen3_5` - `test_monkey_patch.py` — instance patching test - `bf16/test_mini_models.py` — convergence test - `bf16/test_mini_models_with_logits.py` — convergence with logits - `fp32/test_mini_models.py` — convergence (skipped; `ChunkGatedDeltaRuleFunction` doesn't support float32) - `fp32/test_mini_models_with_logits.py` — convergence with logits (skipped, same reason) - `README.md` — added Qwen3.5 to supported models table **Validation on real models (Qwen3.5-0.8B through 27B):** | Model | Logits cosine sim | Training loss rel diff | |-------|:-----------------:|:----------------------:| | 0.8B | > 0.999 | 0.32% | | 2B | > 0.999 | 0.20% | | 4B | > 0.999 | 0.28% | | 9B | > 0.999 | 0.39% | | 27B | > 0.999 | 0.41% | ## Testing Done - Hardware Type: NVIDIA H200 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence ``` test_apply_liger_kernel_to_instance_for_qwen3_5 PASSED test_mini_model[mini_qwen3_5 bf16] PASSED (6.47s) test_mini_model[mini_qwen3_5 bf16 with_logits] PASSED (7.23s) test_mini_model[mini_qwen3_5 fp32] SKIPPED (ChunkGatedDeltaRuleFunction no fp32) test_mini_model[mini_qwen3_5 fp32 with_logits] SKIPPED (same) ``` |
||
|
|
134188d997 |
[NPU] Add Ascend NPU CI status in README (#1131)
Summary ----------- Due to the self-hosted runner permission issue, the test workflow of Ascend NPU cannot be pushed upstream for the time being (https://github.com/linkedin/Liger-Kernel/issues/1022). Now it is placed in the fork repository for running, and echo the test results to the upstream README. Source workflow file ---------- https://github.com/xuedinge233/Liger-Kernel/blob/main/.github/workflows/ascend_npu_ci.yml |
||
|
|
5841280d38 |
Add mHC kernel documentation to README and API reference (#1132)
## Summary - Add `mHC (Hyper-Connections)` entry to the Model Kernels table in `README.md` and `docs/Low-Level-APIs.md` - Add detailed description section in `docs/Low-Level-APIs.md` with architecture overview, usage example, and functional API reference ## Reference Issue - Follows up on #1065 which added the mHC fused kernels but did not update documentation |
||
|
|
64d6ce592c |
[Model] Pixtral Support (#253)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> This PR aims to support pixtral <!--- ## 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. ---> tested model + tested monkey patch <!-- 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: 3090 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude Opus 4.6 <[email protected]> |
||
|
|
6c2565b49b |
Update discord channel link and announcement for meetup (#984)
Update discord channel link and announcement for Liger Kernel China meetup ## Summary Update discord channel link and announcement for Liger Kernel China meetup ## Testing Done - Hardware Type: <BLANK> - [ ] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Shao Tang <[email protected]> Co-authored-by: Tcc0403 <[email protected]> |
||
|
|
38e7953b4e |
Update README.md (#970)
## 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> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
|
720cc68486 |
[feat]: Add support for gpt-oss (#949)
## Summary This work adds support for gpt oss models ## Testing Done: Inference Benchmark Execution Total inference runs: 18 executions 12 successful benchmark measurements 6 pre-run warmup iterations Configurations tested: 3 scenarios with 2 runs each All scenarios passed with consistent, measurable improvements - Hardware Type: RTX A6000 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Claude <[email protected]> |
||
|
|
8a9339806f |
Olmo3 model support [ready for review] (#946)
## Summary Add support for Olmo 3 models: https://huggingface.co/allenai/Olmo-3-1125-32B ## Details Olmo 3 is similar to Olmo 3 in all of the ways that matter to Liger-kernel. It is a separate model type in `transformers` (in order to support SWA) so instead of just aliasing Olmo 3 to Olmo 2 I add a new model type here. ## Testing Done Hardware Type: NVIDIA RTX 4090 - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence --------- Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
|
98d9e197c4 |
add hunyuanv1 dense and moe model (#940)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> This pull request introduces support for the Hunyuanv1 dense and moe model within the Liger-Kernel framework. [HunyuanV1 Model PR](https://github.com/huggingface/transformers/pull/39606) <!--- ## 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. ---> <img width="3482" height="595" alt="1763112849386" src="https://github.com/user-attachments/assets/eaab7b15-7737-4285-9f23-1d01cc09ee91" /> And a simple test: <img width="2748" height="1051" alt="image" src="https://github.com/user-attachments/assets/fa6155f3-9456-4054-a874-41e1c25e4b47" /> <!-- 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: H20 - [ x ] run `make test` to ensure correctness - [ x ] run `make checkstyle` to ensure code style - [ x ] run `make test-convergence` to ensure convergence Co-authored-by: aureli <[email protected]> Co-authored-by: Steven Shimizu <[email protected]> |
||
|
|
9584b94119 |
Add InternVL3 support (#878)
## Summary This PR added InternVL3 models(e.g. https://huggingface.co/OpenGVLab/InternVL3-1B-hf) to liger-kernel. From InternVL2.5-MPO(e.g. https://huggingface.co/OpenGVLab/InternVL2_5-8B-MPO-hf) to InternVL3.5 they share same model type `internvl` but I focused on supporting InternVL3. Other versions may work but not sure. This model is present after transformers v4.52.1: https://github.com/huggingface/transformers/releases/tag/v4.52.1 ## Testing Done - Hardware Type: H100 - [x] run `make test` to ensure correctness: all passed ``` ====================================================================================================== 2072 passed, 215 skipped, 44 warnings in 918.28s (0:15:18) ======================================================================================================= ``` - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence: glm fail, not sure why ``` ======================================================================================================================== short test summary info ======================================================================================================================== FAILED test/convergence/bf16/test_mini_models.py::test_mini_model[mini_glm4v-32-1e-05-dtype14-0.01-0.01-0.1-0.01-0.01-0.01] - AssertionError: [Loss]Number of mismatched elements: 3 FAILED test/convergence/bf16/test_mini_models.py::test_mini_model[mini_glm4v_moe-32-1e-05-dtype15-0.01-0.2-0.1-0.01-0.01-0.01] - AssertionError: [Loss]Number of mismatched elements: 2 ========================================================================================================== 2 failed, 19 passed, 1 warning in 117.34s (0:01:57) ========================================================================================================== ``` ## Important note for testing Existing test failed for multimodal model test, so I fixed some implementation, not sure I've done right. Please have a look and give me review. https://github.com/linkedin/Liger-Kernel/pull/878/files#diff-76dbff2421c7cc7c9833d3bfaa5e1c55da5ad7ce96087d356aefb7d340acaf5bR954-R962 --------- Co-authored-by: Shao Tang <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
|
49fdc2dfd4 |
Remove daily test cron job (#890)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Daily tests have become costly in terms of Modal credits since we now run them on H100s. At the same time, they don’t add much value because we already run tests on every push to the main branch. The only real benefit is catching regressions when upstream dependencies like transformers change, but those issues will surface soon enough once new changes are pushed to the main branch. <!--- ## 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> - [x] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
|
d1441a0de8 |
Fixed AMD CI issue #793 (#887)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> This PR fixes #793 The AMD CI for `Liger-Kernel` was failing due to PyTorch ROCm not being installed correctly. The CI was installing a CUDA build instead of ROCm because PyPI packages were taking precedence over the extraindexurl nightly ROCm index URL. This caused the error: ``` and getattr(torch, infer_device()).get_device_properties(0).total_memory < 36e9, E AttributeError: module 'torch.cpu' has no attribute 'get_device_properties' ``` Since CUDA PyTorch couldn't detect the GPU on the AMD Device. This change updates the README and setup so that AMD users explicitly install the ROCm nightly PyTorch after the dev requirements, ensuring the correct ROCm build is used and torch can detect GPU. --- ## Details 1. **Intel CI issues remain unresolved.** 2. AMD users now need to run one additional step to (as updated in readme to downalod rocm torch, I tried using index-url instead of extra-index-url that helped to downalod correct version but still installed cuda version because of precedence so its better to manually downalod as specified in readme): ``` pip install -e .[dev] pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.3/ ``` This ensures the ROCm build overrides any CUDA installations from PyPI. --- ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> Ran: `python -m pytest test/transformers/test_tvd.py -v ` (was earlier giving AttributeError: module 'torch.cpu' has no attribute 'get_device_properties' , now all testcase passed) <img width="1440" height="778" alt="testcasepassed" src="https://github.com/user-attachments/assets/ec1f9d81-2afa-45d2-9b29-ccbacccf257f" /> <img width="1440" height="774" alt="alltestcaserun" src="https://github.com/user-attachments/assets/5a427ebf-f8e4-4d38-bbf0-0bd9012504bc" /> All test cases passed on AMD MI300X. <!-- 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: AMD MI300X - [x] run `make test` to ensure correctness - [ ] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence |
||
|
|
65c0ad123e |
docs(README): fix intel ci link (#842)
## 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 |
||
|
|
1b9a844a67 |
docs(README): fix gpumode discord badge (#835)
## 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 Co-authored-by: Vaibhav Jindal <[email protected]> |
||
|
|
bc67f94595 |
Update README.md (#823)
Updated the citation to ICML2025 official BibTex record. ## 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> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
|
5e8f12b315 |
Update README.md (#811)
Updated the paper url to https://openreview.net/forum?id=36SjAIT42G, and added contact names for external collaboration. ## 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> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
|
68c70bf8e3 |
Fix minor typo (#809)
## Summary This PR fixes a minor typo in README.md Co-authored-by: Shao Tang <[email protected]> |
||
|
|
9548e70318 |
Update README.md (#808)
Adding CoreAI PoC ([email protected]) to support external collaboration, and updated the paper citation to ICML 2025 CodeML Workshop. ## 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> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence |
||
|
|
8e231a3538 |
Liger support for Llama4 (#740)
## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> Adding liger support to llama4 (text and multimodal). Liger rope is not supported yet for llama4 as it has a different implementation. The tolerance are set higher than usual to pass the tests. <!--- ## 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: NVIDIA H100 80GB HBM3 - [ ] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [ ] run `make test-convergence` to ensure convergence --------- Co-authored-by: Shao Tang <[email protected]> Co-authored-by: Manan Shah <[email protected]> Co-authored-by: Vaibhav Jindal <[email protected]> |