liger-kernel 0.4.1__tar.gz → 0.4.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (60) hide show
  1. {liger_kernel-0.4.1/src/liger_kernel.egg-info → liger_kernel-0.4.2}/PKG-INFO +14 -80
  2. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/README.md +13 -79
  3. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/pyproject.toml +1 -1
  4. liger_kernel-0.4.2/src/liger_kernel/chunked_loss/dpo_loss.py +57 -0
  5. liger_kernel-0.4.2/src/liger_kernel/chunked_loss/fused_linear_preference.py +206 -0
  6. liger_kernel-0.4.2/src/liger_kernel/chunked_loss/orpo_loss.py +63 -0
  7. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/fused_linear_cross_entropy.py +1 -0
  8. liger_kernel-0.4.2/src/liger_kernel/transformers/model/__init__.py +0 -0
  9. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/qwen2_vl.py +43 -17
  10. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/monkey_patch.py +5 -2
  11. {liger_kernel-0.4.1 → liger_kernel-0.4.2/src/liger_kernel.egg-info}/PKG-INFO +14 -80
  12. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel.egg-info/SOURCES.txt +4 -0
  13. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/LICENSE +0 -0
  14. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/NOTICE +0 -0
  15. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/setup.cfg +0 -0
  16. {liger_kernel-0.4.1/src/liger_kernel/ops → liger_kernel-0.4.2/src/liger_kernel/chunked_loss}/__init__.py +0 -0
  17. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/env_report.py +0 -0
  18. {liger_kernel-0.4.1/src/liger_kernel/transformers/model → liger_kernel-0.4.2/src/liger_kernel/ops}/__init__.py +0 -0
  19. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/cross_entropy.py +0 -0
  20. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/experimental/embedding.py +0 -0
  21. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/experimental/mm_int8int2.py +0 -0
  22. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/fused_linear_jsd.py +0 -0
  23. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/geglu.py +0 -0
  24. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/group_norm.py +0 -0
  25. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/jsd.py +0 -0
  26. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/kl_div.py +0 -0
  27. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/layer_norm.py +0 -0
  28. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/rms_norm.py +0 -0
  29. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/rope.py +0 -0
  30. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/swiglu.py +0 -0
  31. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/ops/utils.py +0 -0
  32. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/__init__.py +0 -0
  33. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/auto_model.py +0 -0
  34. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/cross_entropy.py +0 -0
  35. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/experimental/embedding.py +0 -0
  36. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/functional.py +0 -0
  37. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/fused_linear_cross_entropy.py +0 -0
  38. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/fused_linear_jsd.py +0 -0
  39. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/geglu.py +0 -0
  40. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/group_norm.py +0 -0
  41. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/jsd.py +0 -0
  42. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/kl_div.py +0 -0
  43. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/layer_norm.py +0 -0
  44. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/gemma.py +0 -0
  45. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/gemma2.py +0 -0
  46. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/llama.py +0 -0
  47. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/mistral.py +0 -0
  48. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/mixtral.py +0 -0
  49. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/mllama.py +0 -0
  50. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/phi3.py +0 -0
  51. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/model/qwen2.py +0 -0
  52. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/rms_norm.py +0 -0
  53. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/rope.py +0 -0
  54. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/swiglu.py +0 -0
  55. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/transformers/trainer_integration.py +0 -0
  56. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/triton/__init__.py +0 -0
  57. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel/triton/monkey_patch.py +0 -0
  58. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel.egg-info/dependency_links.txt +0 -0
  59. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel.egg-info/requires.txt +0 -0
  60. {liger_kernel-0.4.1 → liger_kernel-0.4.2}/src/liger_kernel.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: liger_kernel
3
- Version: 0.4.1
3
+ Version: 0.4.2
4
4
  Summary: Efficient Triton kernels for LLM Training
5
5
  License: BSD 2-CLAUSE LICENSE
6
6
  Copyright 2024 LinkedIn Corporation
@@ -99,7 +99,8 @@ Requires-Dist: seaborn; extra == "dev"
99
99
 
100
100
  <details>
101
101
  <summary>Latest News 🔥</summary>
102
-
102
+
103
+ - [2024/11/6] We release [v0.4.0](https://github.com/linkedin/Liger-Kernel/releases/tag/v0.4.0): Full AMD support, Tech Report, Modal CI, Llama-3.2-Vision!
103
104
  - [2024/10/21] We have released the tech report of Liger Kernel on Arxiv: https://arxiv.org/pdf/2410.10989
104
105
  - [2024/9/6] We release v0.2.1 ([X post](https://x.com/liger_kernel/status/1832168197002510649)). 2500+ Stars, 10+ New Contributors, 50+ PRs, 50k Downloads in two weeks!
105
106
  - [2024/8/31] CUDA MODE talk, [Liger-Kernel: Real-world Triton kernel for LLM Training](https://youtu.be/gWble4FreV4?si=dxPeIchhkJ36Mbns), [Slides](https://github.com/cuda-mode/lectures?tab=readme-ov-file#lecture-28-liger-kernel)
@@ -127,18 +128,12 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
127
128
 
128
129
  ## Examples
129
130
 
130
- ### Basic
131
-
132
- | **Example** | **Description** | **Lightning Studio** |
133
- |------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
134
- | [**Hugging Face Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/huggingface) | Train LLaMA 3-8B ~20% faster with over 40% memory reduction on Alpaca dataset using 4 A100s with FSDP | TBA |
135
- | [**Lightning Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/lightning) | Increase 15% throughput and reduce memory usage by 40% with LLaMA3-8B on MMLU dataset using 8 A100s with DeepSpeed ZeRO3 | TBA |
136
131
 
137
- ### Advanced
138
-
139
- | **Example** | **Description** | **Lightning Studio** |
140
- |------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
141
- | [**Medusa Multi-head LLM (Retraining Phase)**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa) | Reduce memory usage by 80% with 5 LM heads and improve throughput by 40% using 8 A100s with FSDP | TBA |
132
+ | **Use Case** | **Description** |
133
+ |------------------------------------------------|---------------------------------------------------------------------------------------------------|
134
+ | [**Hugging Face Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/huggingface) | Train LLaMA 3-8B ~20% faster with over 40% memory reduction on Alpaca dataset using 4 A100s with FSDP |
135
+ | [**Lightning Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/lightning) | Increase 15% throughput and reduce memory usage by 40% with LLaMA3-8B on MMLU dataset using 8 A100s with DeepSpeed ZeRO3 |
136
+ | [**Medusa Multi-head LLM (Retraining Phase)**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa) | Reduce memory usage by 80% with 5 LM heads and improve throughput by 40% using 8 A100s with FSDP | |
142
137
 
143
138
  ## Key Features
144
139
 
@@ -149,13 +144,6 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
149
144
  - **Multi-GPU supported:** Compatible with multi-GPU setups (PyTorch FSDP, DeepSpeed, DDP, etc.).
150
145
  - **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift)
151
146
 
152
- ## Target Audiences
153
-
154
- - **Researchers**: Looking to compose models using efficient and reliable kernels for frontier experiments.
155
- - **ML Practitioners**: Focused on maximizing GPU training efficiency with optimal, high-performance kernels.
156
- - **Curious Novices**: Eager to learn how to write reliable Triton kernels to enhance training efficiency.
157
-
158
-
159
147
  ## Installation
160
148
 
161
149
  ### Dependencies
@@ -261,23 +249,6 @@ loss = loss_fn(model.weight, input, target)
261
249
  loss.backward()
262
250
  ```
263
251
 
264
-
265
- ## Structure
266
-
267
- ### Source Code
268
-
269
- - `ops/`: Core Triton operations.
270
- - `transformers/`: PyTorch `nn.Module` implementations built on Triton operations, compliant with the `transformers` API.
271
-
272
- ### Tests
273
-
274
- - `transformers/`: Correctness tests for the Triton-based layers.
275
- - `convergence/`: Patches Hugging Face models with all kernels, runs multiple iterations, and compares weights, logits, and loss layer-by-layer.
276
-
277
- ### Benchmark
278
-
279
- - `benchmark/`: Execution time and memory benchmarks compared to Hugging Face layers.
280
-
281
252
  ## APIs
282
253
 
283
254
  ### AutoModel
@@ -346,54 +317,17 @@ $$\text{GeGLU}(x)=\text{GELU}(xW+b)\otimes(xV+c)$$
346
317
  - **Embedding**: [Embedding](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) is implemented by fusing embedding lookup and output operations. It achieves a peak speedup of ~1.5x in the forward pass and an overall speedup of ~1.1x.
347
318
  - **Matmul int2xint8**: is implemented by using the cache tiled matrix multiplication and by fusing the matmul with the unpacking process which achieves a considerable speed up and performs on par with @torch.compile
348
319
  <!-- TODO: be more specific about batch size -->
349
- > **Note:**
350
- > Reported speedups and memory reductions are with respect to the LLaMA 3-8B Hugging Face layer implementations. All models use 4K hidden size and 4K sequence length and are evaluated based on memory usage and wall time for the forward+backward pass on a single NVIDIA A100 80G GPU using small batch sizes. Liger kernels exhibit more efficient scaling to larger batch sizes, detailed further in the [Benchmark](./benchmark) folder.
351
-
352
- ## Contributing
353
-
354
- [CONTRIBUTING GUIDE](https://github.com/linkedin/Liger-Kernel/blob/main/CONTRIBUTING.md)
355
-
356
- ## Acknowledgement
357
-
358
-
359
- ### Design
360
-
361
- - [@claire_yishan](https://twitter.com/claire_yishan) for the LOGO design
362
- - [Wave Snippets](https://www.wavesnippets.com/) for generating the animated code snippets
363
-
364
- ### Code
365
-
366
- We referenced or used the following projects:
367
-
368
-
369
-
370
- | # | Project | Description | Location | License |
371
- |---|----------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------|
372
- | 1 | [Unsloth](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/unsloth/kernels/utils.py#L43) | `calculate_settings` to determine block size and warp; We reuse it for Norm and MLP | [Liger Kernel Utils](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/utils.py#L23) | [Apache](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/LICENSE) |
373
- | 2 | [Unsloth](https://github.com/unslothai/unsloth/blob/976d11a10d54383aeb7a692c69e01151a20bfd72/unsloth/kernels/rms_layernorm.py#L48) | We modified and added dW calculation on top of Unsloth implementation | [Liger Kernel RMS Norm](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/rms_norm.py#L50) | [Apache](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/LICENSE) |
374
- | 3 | [Triton tutorial](https://triton-lang.org/main/index.html) | We modified on top of triton tutorials | [Liger Kernel RMS Norm](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/rms_norm.py#L50) | [MIT](https://github.com/triton-lang/triton/blob/main/LICENSE) |
375
- | 4 | [tiny shakespeare dataset](https://huggingface.co/datasets/karpathy/tiny_shakespeare) | We use tiny shakespeare dataset to conduct convergence test on mini model | [Liger Kernel Convergence](https://github.com/linkedin/Liger-Kernel/tree/main/test/convergence) | N/A |
376
- | 5 | [Efficient Cross Entropy](https://github.com/mgmalek/efficient_cross_entropy) | We use the idea of gradient-in-forward and chunking | [Liger Kernel Linear Cross Entropy](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py) | [MIT](https://github.com/mgmalek/efficient_cross_entropy/blob/main/LICENSE) |
377
- | 6 | [Flash attn](https://github.com/Dao-AILab/flash-attention) | We take many optimization ideas from the work, such as tiling and recomputation | | [BSD](https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE) |
378
- | 7 | [AutoAWQ](https://github.com/casper-hansen/AutoAWQ) | We reference the design of automodel | [Liger Kernel Auto Model](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/auto_model.py) | [MIT](https://github.com/casper-hansen/AutoAWQ/blob/main/LICENSE) |
379
- | 8 | [llm.c](https://github.com/karpathy/llm.c) | We reference the design of end-to-end testing | [Liger Kernel Convergence Tests](https://github.com/linkedin/Liger-Kernel/tree/main/test/convergence) | [MIT](https://github.com/karpathy/llm.c/blob/master/LICENSE) |
380
-
381
- Many thanks to the contributors to these projects for their invaluable work that helped make Liger possible.
382
-
383
- ## License
384
320
 
385
- This project is licensed under the [BSD 2-CLAUSE](https://github.com/linkedin/Liger-Kernel/blob/main/LICENSE) License (see `LICENSE` for details).
386
- It also includes components from projects licensed under:
321
+ ## Contributing, Acknowledgements, and License
387
322
 
388
- - Apache License 2.0 (see `LICENSE-APACHE-2.0` for details).
389
- - MIT License (see `LICENSE-MIT-AutoAWQ` for details).
390
- - MIT License (see `LICENSE-MIT-Efficient Cross Entropy` for details).
391
- - MIT License (see `LICENSE-MIT-llmc` for details).
392
- - MIT License (see `LICENSE-MIT-triton` for details).
323
+ - [Contributing Guidelines](https://github.com/linkedin/Liger-Kernel/blob/main/docs/CONTRIBUTING.md)
324
+ - [Acknowledgements](https://github.com/linkedin/Liger-Kernel/blob/main/docs/Acknowledgement.md)
325
+ - [License Information](https://github.com/linkedin/Liger-Kernel/blob/main/docs/License.md)
393
326
 
394
327
  ## Contact
395
328
 
396
- - For public discussion, join [our discord channel](https://discord.gg/vNBDpjhb)
329
+ - For issues, create a Github ticket in this repository
330
+ - For open discussion, join [our discord channel](https://discord.gg/gpumode)
397
331
  - For formal collaboration, send an email to byhsu@linkedin.com
398
332
 
399
333
  ## Cite this work
@@ -52,7 +52,8 @@
52
52
 
53
53
  <details>
54
54
  <summary>Latest News 🔥</summary>
55
-
55
+
56
+ - [2024/11/6] We release [v0.4.0](https://github.com/linkedin/Liger-Kernel/releases/tag/v0.4.0): Full AMD support, Tech Report, Modal CI, Llama-3.2-Vision!
56
57
  - [2024/10/21] We have released the tech report of Liger Kernel on Arxiv: https://arxiv.org/pdf/2410.10989
57
58
  - [2024/9/6] We release v0.2.1 ([X post](https://x.com/liger_kernel/status/1832168197002510649)). 2500+ Stars, 10+ New Contributors, 50+ PRs, 50k Downloads in two weeks!
58
59
  - [2024/8/31] CUDA MODE talk, [Liger-Kernel: Real-world Triton kernel for LLM Training](https://youtu.be/gWble4FreV4?si=dxPeIchhkJ36Mbns), [Slides](https://github.com/cuda-mode/lectures?tab=readme-ov-file#lecture-28-liger-kernel)
@@ -80,18 +81,12 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
80
81
 
81
82
  ## Examples
82
83
 
83
- ### Basic
84
-
85
- | **Example** | **Description** | **Lightning Studio** |
86
- |------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
87
- | [**Hugging Face Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/huggingface) | Train LLaMA 3-8B ~20% faster with over 40% memory reduction on Alpaca dataset using 4 A100s with FSDP | TBA |
88
- | [**Lightning Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/lightning) | Increase 15% throughput and reduce memory usage by 40% with LLaMA3-8B on MMLU dataset using 8 A100s with DeepSpeed ZeRO3 | TBA |
89
84
 
90
- ### Advanced
91
-
92
- | **Example** | **Description** | **Lightning Studio** |
93
- |------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
94
- | [**Medusa Multi-head LLM (Retraining Phase)**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa) | Reduce memory usage by 80% with 5 LM heads and improve throughput by 40% using 8 A100s with FSDP | TBA |
85
+ | **Use Case** | **Description** |
86
+ |------------------------------------------------|---------------------------------------------------------------------------------------------------|
87
+ | [**Hugging Face Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/huggingface) | Train LLaMA 3-8B ~20% faster with over 40% memory reduction on Alpaca dataset using 4 A100s with FSDP |
88
+ | [**Lightning Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/lightning) | Increase 15% throughput and reduce memory usage by 40% with LLaMA3-8B on MMLU dataset using 8 A100s with DeepSpeed ZeRO3 |
89
+ | [**Medusa Multi-head LLM (Retraining Phase)**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa) | Reduce memory usage by 80% with 5 LM heads and improve throughput by 40% using 8 A100s with FSDP | |
95
90
 
96
91
  ## Key Features
97
92
 
@@ -102,13 +97,6 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
102
97
  - **Multi-GPU supported:** Compatible with multi-GPU setups (PyTorch FSDP, DeepSpeed, DDP, etc.).
103
98
  - **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift)
104
99
 
105
- ## Target Audiences
106
-
107
- - **Researchers**: Looking to compose models using efficient and reliable kernels for frontier experiments.
108
- - **ML Practitioners**: Focused on maximizing GPU training efficiency with optimal, high-performance kernels.
109
- - **Curious Novices**: Eager to learn how to write reliable Triton kernels to enhance training efficiency.
110
-
111
-
112
100
  ## Installation
113
101
 
114
102
  ### Dependencies
@@ -214,23 +202,6 @@ loss = loss_fn(model.weight, input, target)
214
202
  loss.backward()
215
203
  ```
216
204
 
217
-
218
- ## Structure
219
-
220
- ### Source Code
221
-
222
- - `ops/`: Core Triton operations.
223
- - `transformers/`: PyTorch `nn.Module` implementations built on Triton operations, compliant with the `transformers` API.
224
-
225
- ### Tests
226
-
227
- - `transformers/`: Correctness tests for the Triton-based layers.
228
- - `convergence/`: Patches Hugging Face models with all kernels, runs multiple iterations, and compares weights, logits, and loss layer-by-layer.
229
-
230
- ### Benchmark
231
-
232
- - `benchmark/`: Execution time and memory benchmarks compared to Hugging Face layers.
233
-
234
205
  ## APIs
235
206
 
236
207
  ### AutoModel
@@ -299,54 +270,17 @@ $$\text{GeGLU}(x)=\text{GELU}(xW+b)\otimes(xV+c)$$
299
270
  - **Embedding**: [Embedding](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) is implemented by fusing embedding lookup and output operations. It achieves a peak speedup of ~1.5x in the forward pass and an overall speedup of ~1.1x.
300
271
  - **Matmul int2xint8**: is implemented by using the cache tiled matrix multiplication and by fusing the matmul with the unpacking process which achieves a considerable speed up and performs on par with @torch.compile
301
272
  <!-- TODO: be more specific about batch size -->
302
- > **Note:**
303
- > Reported speedups and memory reductions are with respect to the LLaMA 3-8B Hugging Face layer implementations. All models use 4K hidden size and 4K sequence length and are evaluated based on memory usage and wall time for the forward+backward pass on a single NVIDIA A100 80G GPU using small batch sizes. Liger kernels exhibit more efficient scaling to larger batch sizes, detailed further in the [Benchmark](./benchmark) folder.
304
-
305
- ## Contributing
306
-
307
- [CONTRIBUTING GUIDE](https://github.com/linkedin/Liger-Kernel/blob/main/CONTRIBUTING.md)
308
-
309
- ## Acknowledgement
310
-
311
-
312
- ### Design
313
-
314
- - [@claire_yishan](https://twitter.com/claire_yishan) for the LOGO design
315
- - [Wave Snippets](https://www.wavesnippets.com/) for generating the animated code snippets
316
-
317
- ### Code
318
-
319
- We referenced or used the following projects:
320
-
321
-
322
-
323
- | # | Project | Description | Location | License |
324
- |---|----------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------|
325
- | 1 | [Unsloth](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/unsloth/kernels/utils.py#L43) | `calculate_settings` to determine block size and warp; We reuse it for Norm and MLP | [Liger Kernel Utils](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/utils.py#L23) | [Apache](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/LICENSE) |
326
- | 2 | [Unsloth](https://github.com/unslothai/unsloth/blob/976d11a10d54383aeb7a692c69e01151a20bfd72/unsloth/kernels/rms_layernorm.py#L48) | We modified and added dW calculation on top of Unsloth implementation | [Liger Kernel RMS Norm](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/rms_norm.py#L50) | [Apache](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/LICENSE) |
327
- | 3 | [Triton tutorial](https://triton-lang.org/main/index.html) | We modified on top of triton tutorials | [Liger Kernel RMS Norm](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/rms_norm.py#L50) | [MIT](https://github.com/triton-lang/triton/blob/main/LICENSE) |
328
- | 4 | [tiny shakespeare dataset](https://huggingface.co/datasets/karpathy/tiny_shakespeare) | We use tiny shakespeare dataset to conduct convergence test on mini model | [Liger Kernel Convergence](https://github.com/linkedin/Liger-Kernel/tree/main/test/convergence) | N/A |
329
- | 5 | [Efficient Cross Entropy](https://github.com/mgmalek/efficient_cross_entropy) | We use the idea of gradient-in-forward and chunking | [Liger Kernel Linear Cross Entropy](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py) | [MIT](https://github.com/mgmalek/efficient_cross_entropy/blob/main/LICENSE) |
330
- | 6 | [Flash attn](https://github.com/Dao-AILab/flash-attention) | We take many optimization ideas from the work, such as tiling and recomputation | | [BSD](https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE) |
331
- | 7 | [AutoAWQ](https://github.com/casper-hansen/AutoAWQ) | We reference the design of automodel | [Liger Kernel Auto Model](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/auto_model.py) | [MIT](https://github.com/casper-hansen/AutoAWQ/blob/main/LICENSE) |
332
- | 8 | [llm.c](https://github.com/karpathy/llm.c) | We reference the design of end-to-end testing | [Liger Kernel Convergence Tests](https://github.com/linkedin/Liger-Kernel/tree/main/test/convergence) | [MIT](https://github.com/karpathy/llm.c/blob/master/LICENSE) |
333
-
334
- Many thanks to the contributors to these projects for their invaluable work that helped make Liger possible.
335
-
336
- ## License
337
273
 
338
- This project is licensed under the [BSD 2-CLAUSE](https://github.com/linkedin/Liger-Kernel/blob/main/LICENSE) License (see `LICENSE` for details).
339
- It also includes components from projects licensed under:
274
+ ## Contributing, Acknowledgements, and License
340
275
 
341
- - Apache License 2.0 (see `LICENSE-APACHE-2.0` for details).
342
- - MIT License (see `LICENSE-MIT-AutoAWQ` for details).
343
- - MIT License (see `LICENSE-MIT-Efficient Cross Entropy` for details).
344
- - MIT License (see `LICENSE-MIT-llmc` for details).
345
- - MIT License (see `LICENSE-MIT-triton` for details).
276
+ - [Contributing Guidelines](https://github.com/linkedin/Liger-Kernel/blob/main/docs/CONTRIBUTING.md)
277
+ - [Acknowledgements](https://github.com/linkedin/Liger-Kernel/blob/main/docs/Acknowledgement.md)
278
+ - [License Information](https://github.com/linkedin/Liger-Kernel/blob/main/docs/License.md)
346
279
 
347
280
  ## Contact
348
281
 
349
- - For public discussion, join [our discord channel](https://discord.gg/vNBDpjhb)
282
+ - For issues, create a Github ticket in this repository
283
+ - For open discussion, join [our discord channel](https://discord.gg/gpumode)
350
284
  - For formal collaboration, send an email to byhsu@linkedin.com
351
285
 
352
286
  ## Cite this work
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "liger_kernel"
7
- version = "0.4.1"
7
+ version = "0.4.2"
8
8
  description = "Efficient Triton kernels for LLM Training"
9
9
  urls = { "Homepage" = "https://github.com/linkedin/Liger-Kernel" }
10
10
  readme = { file = "README.md", content-type = "text/markdown" }
@@ -0,0 +1,57 @@
1
+ import torch.nn.functional as F
2
+
3
+ from liger_kernel.chunked_loss.fused_linear_preference import (
4
+ LigerFusedLinearPreferenceBase,
5
+ )
6
+
7
+
8
+ class LigerFusedLinearDPOFunction(LigerFusedLinearPreferenceBase):
9
+
10
+ @staticmethod
11
+ def preference_loss_fn(chosen_logps, rejected_logps, beta=0.1):
12
+ """
13
+ Compute DPO loss (Direct Preference Optimization).
14
+ Args:
15
+ chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
16
+ rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
17
+ beta (float): Weight for the direct preference loss.
18
+ """
19
+ logits_diff = beta * (chosen_logps - rejected_logps)
20
+ losses = -F.logsigmoid(logits_diff)
21
+ return losses.sum()
22
+
23
+ @staticmethod
24
+ def forward(
25
+ ctx,
26
+ _input,
27
+ weight,
28
+ target,
29
+ bias=None,
30
+ ignore_index=-100,
31
+ beta=0.1,
32
+ compute_nll_loss=True,
33
+ compiled=True,
34
+ ):
35
+ """
36
+ Fused linear layer with DPO (Direct Preference Optimization) loss.
37
+ Handles both the forward and backward pass of the final linear layer with DPO loss.
38
+ """
39
+ return LigerFusedLinearPreferenceBase.forward(
40
+ ctx=ctx,
41
+ _input=_input,
42
+ weight=weight,
43
+ target=target,
44
+ bias=bias,
45
+ loss_fn=LigerFusedLinearDPOFunction.preference_loss_fn,
46
+ compute_nll_loss=compute_nll_loss,
47
+ ignore_index=ignore_index,
48
+ beta=beta,
49
+ compiled=compiled,
50
+ )
51
+
52
+ @staticmethod
53
+ def backward(ctx, grad_output):
54
+ # Get gradients for _input, weight, bias, and target from the base class
55
+ grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
56
+ # Return these gradients, followed by None for the remaining inputs
57
+ return *grads, None, None, None, None
@@ -0,0 +1,206 @@
1
+ from abc import abstractmethod
2
+ from functools import partial
3
+
4
+ import torch
5
+ from torch.nn import functional as F
6
+
7
+
8
+ class LigerFusedLinearPreferenceBase(torch.autograd.Function):
9
+
10
+ @abstractmethod
11
+ def preference_loss_fn(chosen_logps, rejected_logps, beta=0.1):
12
+ """
13
+ Compute preference loss.
14
+ Args:
15
+ chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
16
+ rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
17
+ beta (float): Weight for the odds ratio loss.
18
+ """
19
+ raise NotImplementedError("Preference loss function must be implemented.")
20
+
21
+ @staticmethod
22
+ def forward(
23
+ ctx,
24
+ _input,
25
+ weight,
26
+ target,
27
+ bias=None,
28
+ loss_fn=None,
29
+ chunk_size=1,
30
+ compute_nll_loss=True,
31
+ ignore_index=-100,
32
+ beta=0.1,
33
+ compiled=True,
34
+ ):
35
+ """
36
+ Base class for fused linear layer with preference loss.
37
+ Expects _input to be stacked with chosen and rejected inputs on the batch dimension.
38
+
39
+ Args:
40
+ _input (torch.Tensor): Input tensor. Shape: (batch_size, seq_len, hidden_size).
41
+ weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size).
42
+ target (torch.Tensor): Target tensor. Shape: (batch_size, seq_len).
43
+ bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
44
+ loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
45
+ chunk_size (int): Size of a chunk (# of batches of stacked chosen and rejected inputs).
46
+ compute_nll_loss (bool): Whether to compute NLL loss.
47
+ ignore_index (int): Index to ignore for loss computation.
48
+ beta (float): Weight for the odds ratio loss.
49
+ compiled (bool): Whether to use torch compile for chunk accumulation.
50
+ """
51
+ # TODO: Tune CHUNK_SIZE to fully utilize the GPU
52
+ CHUNK_SIZE = chunk_size
53
+
54
+ grad_weight = torch.zeros_like(weight)
55
+ grad_chosen_inputs = []
56
+ grad_rejected_inputs = []
57
+ grad_bias = torch.zeros_like(bias) if bias is not None else None
58
+ loss_acc = torch.zeros((), device=_input.device)
59
+
60
+ chunks = max(1, _input.shape[0] // (2 * CHUNK_SIZE))
61
+ loss_func_to_call = partial(
62
+ LigerFusedLinearPreferenceBase._compute_loss,
63
+ preference_loss_fn=loss_fn,
64
+ ignore_index=ignore_index,
65
+ beta=beta,
66
+ compute_nll_loss=compute_nll_loss,
67
+ full_target=target,
68
+ )
69
+
70
+ def accumulate_chunk(input_chunk, target_chunk):
71
+ if bias is not None:
72
+ (chunk_grad_input, chunk_grad_weight, chunk_grad_bias), (
73
+ chunk_loss,
74
+ (chunk_or_loss, chunk_chosen_logps, chunk_rejected_logps),
75
+ ) = torch.func.grad_and_value(
76
+ loss_func_to_call, argnums=(0, 1, 3), has_aux=True
77
+ )(
78
+ input_chunk, weight, target_chunk, bias
79
+ )
80
+ grad_bias.add_(chunk_grad_bias)
81
+ else:
82
+ (chunk_grad_input, chunk_grad_weight), (
83
+ chunk_loss,
84
+ (chunk_or_loss, chunk_chosen_logps, chunk_rejected_logps),
85
+ ) = torch.func.grad_and_value(
86
+ loss_func_to_call, argnums=(0, 1), has_aux=True
87
+ )(
88
+ input_chunk, weight, target_chunk
89
+ )
90
+ grad_weight.add_(chunk_grad_weight)
91
+ loss_acc.add_(chunk_loss)
92
+ return chunk_grad_input
93
+
94
+ len_chosen = target.shape[0] // 2
95
+ _chosen_input_chunks = torch.chunk(_input[:len_chosen], chunks=chunks, dim=0)
96
+ _chosen_target_chunks = torch.chunk(target[:len_chosen], chunks=chunks, dim=0)
97
+ _rejected_input_chunks = torch.chunk(_input[len_chosen:], chunks=chunks, dim=0)
98
+ _rejected_target_chunks = torch.chunk(target[len_chosen:], chunks=chunks, dim=0)
99
+
100
+ for (
101
+ chosen_input_chunk,
102
+ rejected_input_chunk,
103
+ chosen_target_chunk,
104
+ rejected_target_chunk,
105
+ ) in zip(
106
+ _chosen_input_chunks,
107
+ _rejected_input_chunks,
108
+ _chosen_target_chunks,
109
+ _rejected_target_chunks,
110
+ ):
111
+ input_chunk = torch.cat([chosen_input_chunk, rejected_input_chunk], dim=0)
112
+ target_chunk = torch.cat(
113
+ [chosen_target_chunk, rejected_target_chunk], dim=0
114
+ )
115
+
116
+ if compiled:
117
+ accumulate_chunk = torch.compile(accumulate_chunk)
118
+ grad_input = accumulate_chunk(input_chunk, target_chunk)
119
+
120
+ grad_chosen_inputs.append(grad_input[: chosen_target_chunk.shape[0]])
121
+ grad_rejected_inputs.append(grad_input[chosen_target_chunk.shape[0] :])
122
+
123
+ # combine grad_chosen_inputs and grad_rejected_inputs
124
+ grad_inputs = grad_chosen_inputs + grad_rejected_inputs
125
+
126
+ ctx.save_for_backward(
127
+ torch.cat(grad_inputs, dim=0),
128
+ grad_weight,
129
+ grad_bias,
130
+ )
131
+ return loss_acc
132
+
133
+ @staticmethod
134
+ def backward(ctx, grad_output):
135
+ grad_input, grad_weight, grad_bias = ctx.saved_tensors
136
+ if torch.ne(grad_output, torch.tensor(1.0, device=grad_output.device)):
137
+ grad_input = grad_input * grad_output
138
+ grad_weight = grad_weight * grad_output
139
+ grad_bias = grad_bias * grad_output if grad_bias is not None else None
140
+
141
+ return grad_input, grad_weight, None, grad_bias, None, None, None
142
+
143
+ @staticmethod
144
+ def _compute_loss(
145
+ input_chunk,
146
+ weight,
147
+ target_chunk,
148
+ bias=None,
149
+ preference_loss_fn=None,
150
+ full_target=None,
151
+ ignore_index=-100,
152
+ beta=0.1,
153
+ compute_nll_loss=True,
154
+ **loss_kwargs,
155
+ ):
156
+ """
157
+ Compute the total loss for a chunk of input and target, while using an alignment/preference loss function.
158
+ Args:
159
+ preference_loss_fn (callable): Loss function to compute the loss on a chunk of input/target.
160
+ input_chunk (torch.Tensor): Chunk of input tensor. Shape: (2 * chunk_size, sequence_length, hidden_size).
161
+ weight (torch.Tensor): Weight tensor. Shape: (vocab_size, hidden_size).
162
+ target_chunk (torch.Tensor): Chunk of target tensor. Shape: (2 * chunk_size, sequence_length).
163
+ bias (torch.Tensor, optional): Bias tensor. Shape: (vocab_size,).
164
+ full_target (torch.Tensor): Full target tensor. Shape: (batch_size, sequence_length).
165
+ ignore_index (int): Index to ignore for loss computation.
166
+ beta (float): Weight for the odds ratio loss.
167
+ loss_kwargs (dict): Additional arguments for the loss function.
168
+ """
169
+ len_chosen_chunk = target_chunk.shape[0] // 2
170
+
171
+ logits_chunk = input_chunk @ weight.t() # chunk_size x V
172
+ if bias is not None:
173
+ logits_chunk = logits_chunk + bias
174
+ log_probs_chunk = F.log_softmax(logits_chunk.float(), dim=-1)
175
+
176
+ chosen_nll_loss = 0.0
177
+ if compute_nll_loss:
178
+ chosen_nll_loss = F.nll_loss(
179
+ log_probs_chunk[:len_chosen_chunk].view(-1, log_probs_chunk.shape[-1]),
180
+ target_chunk[:len_chosen_chunk].view(-1),
181
+ reduction="sum",
182
+ ignore_index=ignore_index,
183
+ )
184
+ chosen_nll_loss = (
185
+ chosen_nll_loss
186
+ / (full_target[: full_target.shape[0] // 2] != ignore_index).sum()
187
+ )
188
+
189
+ loss_mask = target_chunk != ignore_index
190
+ label_chunk = torch.where(loss_mask, target_chunk, 0)
191
+
192
+ per_token_logps = log_probs_chunk.gather(-1, label_chunk.unsqueeze(-1)).squeeze(
193
+ -1
194
+ )
195
+ average_log_prob = (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1)
196
+
197
+ chosen_logps = average_log_prob[:len_chosen_chunk]
198
+ rejected_logps = average_log_prob[len_chosen_chunk:]
199
+
200
+ alignment_loss = preference_loss_fn(
201
+ chosen_logps, rejected_logps, beta=beta, **loss_kwargs
202
+ )
203
+ alignment_loss = alignment_loss / (full_target.shape[0] // 2)
204
+
205
+ loss = chosen_nll_loss - alignment_loss
206
+ return loss, (alignment_loss, chosen_logps, rejected_logps)
@@ -0,0 +1,63 @@
1
+ import torch
2
+ import torch.nn.functional as F
3
+
4
+ from liger_kernel.chunked_loss.fused_linear_preference import (
5
+ LigerFusedLinearPreferenceBase,
6
+ )
7
+
8
+
9
+ class LigerFusedLinearORPOFunction(LigerFusedLinearPreferenceBase):
10
+
11
+ @staticmethod
12
+ def preference_loss_fn(chosen_logps, rejected_logps, beta=0.1):
13
+ """
14
+ Compute odds-ratio loss.
15
+ Args:
16
+ chosen_logps (torch.Tensor): Avg log probabilities of chosen tokens. Shape: (batch_size,).
17
+ rejected_logps (torch.Tensor): Avg log probabilities of rejected tokens. Shape: (batch_size,).
18
+ beta (float): Weight for the odds ratio loss.
19
+ """
20
+ log_odds = (chosen_logps - rejected_logps) - (
21
+ torch.log1p(-torch.exp(chosen_logps))
22
+ - torch.log1p(-torch.exp(rejected_logps))
23
+ )
24
+ ratio = F.logsigmoid(log_odds)
25
+ return beta * ratio.sum()
26
+
27
+ @staticmethod
28
+ def forward(
29
+ ctx,
30
+ _input,
31
+ weight,
32
+ target,
33
+ bias=None,
34
+ ignore_index=-100,
35
+ beta=0.1,
36
+ compute_nll_loss=True,
37
+ compiled=True,
38
+ ):
39
+ """
40
+ Fused linear layer with ORPO (Odds-Ratio Preference Optimization) loss.
41
+ Handles both the forward and backward pass of the final linear layer with ORPO loss.
42
+ Inspired from LigerFusedLinearCrossEntropyFunction (https://arxiv.org/abs/2410.10989) which fuses final linear layer and CE loss.
43
+ """
44
+
45
+ return LigerFusedLinearPreferenceBase.forward(
46
+ ctx=ctx,
47
+ _input=_input,
48
+ weight=weight,
49
+ target=target,
50
+ bias=bias,
51
+ loss_fn=LigerFusedLinearORPOFunction.preference_loss_fn,
52
+ compute_nll_loss=compute_nll_loss,
53
+ ignore_index=ignore_index,
54
+ beta=beta,
55
+ compiled=compiled,
56
+ )
57
+
58
+ @staticmethod
59
+ def backward(ctx, grad_output):
60
+ # Get gradients for _input, weight, bias, and target from the base class
61
+ grads = LigerFusedLinearPreferenceBase.backward(ctx, grad_output)[:4]
62
+ # Return these gradients, followed by None for the remaining inputs
63
+ return *grads, None, None, None, None
@@ -229,6 +229,7 @@ class LigerFusedLinearCrossEntropyFunction(torch.autograd.Function):
229
229
  label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing.
230
230
  reduction: reduction to apply
231
231
  """
232
+
232
233
  loss, grad_input, grad_weight, grad_bias = fused_linear_cross_entropy_forward(
233
234
  _input,
234
235
  weight,
@@ -1,7 +1,9 @@
1
1
  from typing import List, Optional, Tuple, Union
2
2
 
3
3
  import torch
4
+ from packaging import version
4
5
  from torch.nn import CrossEntropyLoss
6
+ from transformers import __version__ as transformers_version
5
7
  from transformers.models.qwen2_vl.modeling_qwen2_vl import (
6
8
  _CONFIG_FOR_DOC,
7
9
  QWEN2_VL_INPUTS_DOCSTRING,
@@ -80,8 +82,6 @@ def lce_forward(
80
82
  >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
81
83
  "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
82
84
  ```"""
83
- # FIXME: The code is outdated and not compatible with transformer >= 4.46.1
84
-
85
85
  output_attentions = (
86
86
  output_attentions
87
87
  if output_attentions is not None
@@ -100,27 +100,53 @@ def lce_forward(
100
100
  inputs_embeds = self.model.embed_tokens(input_ids)
101
101
  if pixel_values is not None:
102
102
  pixel_values = pixel_values.type(self.visual.get_dtype())
103
- image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).to(
104
- inputs_embeds.device
103
+ image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
104
+ n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
105
+ n_image_features = image_embeds.shape[0]
106
+ if n_image_tokens != n_image_features:
107
+ raise ValueError(
108
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
109
+ )
110
+ image_mask = (
111
+ (input_ids == self.config.image_token_id)
112
+ .unsqueeze(-1)
113
+ .expand_as(inputs_embeds)
114
+ .to(inputs_embeds.device)
105
115
  )
106
- image_mask = input_ids == self.config.image_token_id
107
- if self.training:
108
- inputs_embeds = inputs_embeds.clone()
109
- inputs_embeds[image_mask] = image_embeds
116
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
117
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
118
+
110
119
  if pixel_values_videos is not None:
111
120
  pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
112
- video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).to(
113
- inputs_embeds.device
121
+ video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
122
+ n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
123
+ n_video_features = video_embeds.shape[0]
124
+ if n_video_tokens != n_video_features:
125
+ raise ValueError(
126
+ f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
127
+ )
128
+ video_mask = (
129
+ (input_ids == self.config.video_token_id)
130
+ .unsqueeze(-1)
131
+ .expand_as(inputs_embeds)
132
+ .to(inputs_embeds.device)
114
133
  )
115
- video_mask = input_ids == self.config.video_token_id
116
- inputs_embeds[video_mask] = video_embeds
134
+ video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
135
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
136
+
117
137
  if attention_mask is not None:
118
138
  attention_mask = attention_mask.to(inputs_embeds.device)
119
- # The code is copied from https://github.com/huggingface/transformers/pull/33487
120
- if position_ids is None and input_ids is not None:
121
- position_ids, _ = self.get_rope_index(
122
- input_ids, image_grid_thw, video_grid_thw, attention_mask
123
- )
139
+
140
+ if version.parse(transformers_version) > version.parse("4.46.2"):
141
+ # NOTE: this bug fix for qwen2-vl is not applied until transformers 4.47.0
142
+ # https://github.com/huggingface/transformers/issues/33401
143
+ # While correct, this breaks equivalence with past versions of Qwen2-VL from
144
+ # transformers and leads to failed tests or users noticing differences in results.
145
+ # TODO: remove above conditional when liger drops support for transformers<4.47.0
146
+ if position_ids is None and input_ids is not None:
147
+ position_ids, _ = self.get_rope_index(
148
+ input_ids, image_grid_thw, video_grid_thw, attention_mask
149
+ )
124
150
 
125
151
  outputs = self.model(
126
152
  input_ids=None,
@@ -56,12 +56,15 @@ def _bind_method_to_module(module, method_name: str, new_method: Callable):
56
56
  module.__dict__[method_name] = new_method.__get__(module, module.__class__)
57
57
 
58
58
 
59
- def _patch_rms_norm_module(module, offset=0.0, eps=1e-6, casting_mode="llama"):
59
+ def _patch_rms_norm_module(
60
+ module, offset=0.0, eps=1e-6, casting_mode="llama", in_place=True
61
+ ):
60
62
  module.offset = offset
61
63
  module.casting_mode = casting_mode
62
64
  module.variance_epsilon = (
63
65
  getattr(module, "variance_epsilon", None) or getattr(module, "eps", None) or eps
64
66
  )
67
+ module.in_place = in_place
65
68
  _bind_method_to_module(module, "forward", LigerRMSNorm.forward)
66
69
  _bind_method_to_module(module, "extra_repr", LigerRMSNorm.extra_repr)
67
70
 
@@ -510,7 +513,7 @@ def apply_liger_kernel_to_gemma2(
510
513
  LigerRMSNorm, offset=1.0, casting_mode="gemma", init_fn="zeros", in_place=False
511
514
  )
512
515
  _patch_rms_norm_module_for_gemma2 = partial(
513
- _patch_rms_norm_module, offset=1.0, casting_mode="gemma"
516
+ _patch_rms_norm_module, offset=1.0, casting_mode="gemma", in_place=False
514
517
  )
515
518
 
516
519
  if rope:
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: liger_kernel
3
- Version: 0.4.1
3
+ Version: 0.4.2
4
4
  Summary: Efficient Triton kernels for LLM Training
5
5
  License: BSD 2-CLAUSE LICENSE
6
6
  Copyright 2024 LinkedIn Corporation
@@ -99,7 +99,8 @@ Requires-Dist: seaborn; extra == "dev"
99
99
 
100
100
  <details>
101
101
  <summary>Latest News 🔥</summary>
102
-
102
+
103
+ - [2024/11/6] We release [v0.4.0](https://github.com/linkedin/Liger-Kernel/releases/tag/v0.4.0): Full AMD support, Tech Report, Modal CI, Llama-3.2-Vision!
103
104
  - [2024/10/21] We have released the tech report of Liger Kernel on Arxiv: https://arxiv.org/pdf/2410.10989
104
105
  - [2024/9/6] We release v0.2.1 ([X post](https://x.com/liger_kernel/status/1832168197002510649)). 2500+ Stars, 10+ New Contributors, 50+ PRs, 50k Downloads in two weeks!
105
106
  - [2024/8/31] CUDA MODE talk, [Liger-Kernel: Real-world Triton kernel for LLM Training](https://youtu.be/gWble4FreV4?si=dxPeIchhkJ36Mbns), [Slides](https://github.com/cuda-mode/lectures?tab=readme-ov-file#lecture-28-liger-kernel)
@@ -127,18 +128,12 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
127
128
 
128
129
  ## Examples
129
130
 
130
- ### Basic
131
-
132
- | **Example** | **Description** | **Lightning Studio** |
133
- |------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
134
- | [**Hugging Face Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/huggingface) | Train LLaMA 3-8B ~20% faster with over 40% memory reduction on Alpaca dataset using 4 A100s with FSDP | TBA |
135
- | [**Lightning Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/lightning) | Increase 15% throughput and reduce memory usage by 40% with LLaMA3-8B on MMLU dataset using 8 A100s with DeepSpeed ZeRO3 | TBA |
136
131
 
137
- ### Advanced
138
-
139
- | **Example** | **Description** | **Lightning Studio** |
140
- |------------------------------------------------|---------------------------------------------------------------------------------------------------|----------------------|
141
- | [**Medusa Multi-head LLM (Retraining Phase)**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa) | Reduce memory usage by 80% with 5 LM heads and improve throughput by 40% using 8 A100s with FSDP | TBA |
132
+ | **Use Case** | **Description** |
133
+ |------------------------------------------------|---------------------------------------------------------------------------------------------------|
134
+ | [**Hugging Face Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/huggingface) | Train LLaMA 3-8B ~20% faster with over 40% memory reduction on Alpaca dataset using 4 A100s with FSDP |
135
+ | [**Lightning Trainer**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/lightning) | Increase 15% throughput and reduce memory usage by 40% with LLaMA3-8B on MMLU dataset using 8 A100s with DeepSpeed ZeRO3 |
136
+ | [**Medusa Multi-head LLM (Retraining Phase)**](https://github.com/linkedin/Liger-Kernel/tree/main/examples/medusa) | Reduce memory usage by 80% with 5 LM heads and improve throughput by 40% using 8 A100s with FSDP | |
142
137
 
143
138
  ## Key Features
144
139
 
@@ -149,13 +144,6 @@ With one line of code, Liger Kernel can increase throughput by more than 20% and
149
144
  - **Multi-GPU supported:** Compatible with multi-GPU setups (PyTorch FSDP, DeepSpeed, DDP, etc.).
150
145
  - **Trainer Framework Integration**: [Axolotl](https://github.com/axolotl-ai-cloud/axolotl), [LLaMa-Factory](https://github.com/hiyouga/LLaMA-Factory), [SFTTrainer](https://github.com/huggingface/trl/releases/tag/v0.10.1), [Hugging Face Trainer](https://github.com/huggingface/transformers/pull/32860), [SWIFT](https://github.com/modelscope/ms-swift)
151
146
 
152
- ## Target Audiences
153
-
154
- - **Researchers**: Looking to compose models using efficient and reliable kernels for frontier experiments.
155
- - **ML Practitioners**: Focused on maximizing GPU training efficiency with optimal, high-performance kernels.
156
- - **Curious Novices**: Eager to learn how to write reliable Triton kernels to enhance training efficiency.
157
-
158
-
159
147
  ## Installation
160
148
 
161
149
  ### Dependencies
@@ -261,23 +249,6 @@ loss = loss_fn(model.weight, input, target)
261
249
  loss.backward()
262
250
  ```
263
251
 
264
-
265
- ## Structure
266
-
267
- ### Source Code
268
-
269
- - `ops/`: Core Triton operations.
270
- - `transformers/`: PyTorch `nn.Module` implementations built on Triton operations, compliant with the `transformers` API.
271
-
272
- ### Tests
273
-
274
- - `transformers/`: Correctness tests for the Triton-based layers.
275
- - `convergence/`: Patches Hugging Face models with all kernels, runs multiple iterations, and compares weights, logits, and loss layer-by-layer.
276
-
277
- ### Benchmark
278
-
279
- - `benchmark/`: Execution time and memory benchmarks compared to Hugging Face layers.
280
-
281
252
  ## APIs
282
253
 
283
254
  ### AutoModel
@@ -346,54 +317,17 @@ $$\text{GeGLU}(x)=\text{GELU}(xW+b)\otimes(xV+c)$$
346
317
  - **Embedding**: [Embedding](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) is implemented by fusing embedding lookup and output operations. It achieves a peak speedup of ~1.5x in the forward pass and an overall speedup of ~1.1x.
347
318
  - **Matmul int2xint8**: is implemented by using the cache tiled matrix multiplication and by fusing the matmul with the unpacking process which achieves a considerable speed up and performs on par with @torch.compile
348
319
  <!-- TODO: be more specific about batch size -->
349
- > **Note:**
350
- > Reported speedups and memory reductions are with respect to the LLaMA 3-8B Hugging Face layer implementations. All models use 4K hidden size and 4K sequence length and are evaluated based on memory usage and wall time for the forward+backward pass on a single NVIDIA A100 80G GPU using small batch sizes. Liger kernels exhibit more efficient scaling to larger batch sizes, detailed further in the [Benchmark](./benchmark) folder.
351
-
352
- ## Contributing
353
-
354
- [CONTRIBUTING GUIDE](https://github.com/linkedin/Liger-Kernel/blob/main/CONTRIBUTING.md)
355
-
356
- ## Acknowledgement
357
-
358
-
359
- ### Design
360
-
361
- - [@claire_yishan](https://twitter.com/claire_yishan) for the LOGO design
362
- - [Wave Snippets](https://www.wavesnippets.com/) for generating the animated code snippets
363
-
364
- ### Code
365
-
366
- We referenced or used the following projects:
367
-
368
-
369
-
370
- | # | Project | Description | Location | License |
371
- |---|----------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------------|
372
- | 1 | [Unsloth](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/unsloth/kernels/utils.py#L43) | `calculate_settings` to determine block size and warp; We reuse it for Norm and MLP | [Liger Kernel Utils](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/utils.py#L23) | [Apache](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/LICENSE) |
373
- | 2 | [Unsloth](https://github.com/unslothai/unsloth/blob/976d11a10d54383aeb7a692c69e01151a20bfd72/unsloth/kernels/rms_layernorm.py#L48) | We modified and added dW calculation on top of Unsloth implementation | [Liger Kernel RMS Norm](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/rms_norm.py#L50) | [Apache](https://github.com/unslothai/unsloth/blob/fd753fed99ed5f10ef8a9b7139588d9de9ddecfb/LICENSE) |
374
- | 3 | [Triton tutorial](https://triton-lang.org/main/index.html) | We modified on top of triton tutorials | [Liger Kernel RMS Norm](https://github.com/linkedin/Liger-Kernel/blob/e249eee723978bf8610ff1ea2297d048a2417e20/src/liger_kernel/ops/rms_norm.py#L50) | [MIT](https://github.com/triton-lang/triton/blob/main/LICENSE) |
375
- | 4 | [tiny shakespeare dataset](https://huggingface.co/datasets/karpathy/tiny_shakespeare) | We use tiny shakespeare dataset to conduct convergence test on mini model | [Liger Kernel Convergence](https://github.com/linkedin/Liger-Kernel/tree/main/test/convergence) | N/A |
376
- | 5 | [Efficient Cross Entropy](https://github.com/mgmalek/efficient_cross_entropy) | We use the idea of gradient-in-forward and chunking | [Liger Kernel Linear Cross Entropy](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/ops/fused_linear_cross_entropy.py) | [MIT](https://github.com/mgmalek/efficient_cross_entropy/blob/main/LICENSE) |
377
- | 6 | [Flash attn](https://github.com/Dao-AILab/flash-attention) | We take many optimization ideas from the work, such as tiling and recomputation | | [BSD](https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE) |
378
- | 7 | [AutoAWQ](https://github.com/casper-hansen/AutoAWQ) | We reference the design of automodel | [Liger Kernel Auto Model](https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/auto_model.py) | [MIT](https://github.com/casper-hansen/AutoAWQ/blob/main/LICENSE) |
379
- | 8 | [llm.c](https://github.com/karpathy/llm.c) | We reference the design of end-to-end testing | [Liger Kernel Convergence Tests](https://github.com/linkedin/Liger-Kernel/tree/main/test/convergence) | [MIT](https://github.com/karpathy/llm.c/blob/master/LICENSE) |
380
-
381
- Many thanks to the contributors to these projects for their invaluable work that helped make Liger possible.
382
-
383
- ## License
384
320
 
385
- This project is licensed under the [BSD 2-CLAUSE](https://github.com/linkedin/Liger-Kernel/blob/main/LICENSE) License (see `LICENSE` for details).
386
- It also includes components from projects licensed under:
321
+ ## Contributing, Acknowledgements, and License
387
322
 
388
- - Apache License 2.0 (see `LICENSE-APACHE-2.0` for details).
389
- - MIT License (see `LICENSE-MIT-AutoAWQ` for details).
390
- - MIT License (see `LICENSE-MIT-Efficient Cross Entropy` for details).
391
- - MIT License (see `LICENSE-MIT-llmc` for details).
392
- - MIT License (see `LICENSE-MIT-triton` for details).
323
+ - [Contributing Guidelines](https://github.com/linkedin/Liger-Kernel/blob/main/docs/CONTRIBUTING.md)
324
+ - [Acknowledgements](https://github.com/linkedin/Liger-Kernel/blob/main/docs/Acknowledgement.md)
325
+ - [License Information](https://github.com/linkedin/Liger-Kernel/blob/main/docs/License.md)
393
326
 
394
327
  ## Contact
395
328
 
396
- - For public discussion, join [our discord channel](https://discord.gg/vNBDpjhb)
329
+ - For issues, create a Github ticket in this repository
330
+ - For open discussion, join [our discord channel](https://discord.gg/gpumode)
397
331
  - For formal collaboration, send an email to byhsu@linkedin.com
398
332
 
399
333
  ## Cite this work
@@ -8,6 +8,10 @@ src/liger_kernel.egg-info/SOURCES.txt
8
8
  src/liger_kernel.egg-info/dependency_links.txt
9
9
  src/liger_kernel.egg-info/requires.txt
10
10
  src/liger_kernel.egg-info/top_level.txt
11
+ src/liger_kernel/chunked_loss/__init__.py
12
+ src/liger_kernel/chunked_loss/dpo_loss.py
13
+ src/liger_kernel/chunked_loss/fused_linear_preference.py
14
+ src/liger_kernel/chunked_loss/orpo_loss.py
11
15
  src/liger_kernel/ops/__init__.py
12
16
  src/liger_kernel/ops/cross_entropy.py
13
17
  src/liger_kernel/ops/fused_linear_cross_entropy.py
File without changes
File without changes
File without changes