Elamur Research⌕
☰
home / papers / 2024 / Liger Kernel: Efficient Triton Kernels for LLM Training

Liger Kernel: Efficient Triton Kernels for LLM Training

Paper · 2024 · Pin-Lun Hsu, Yun Dai, Vignesh Kothapalli et al.

Summary

Training large language models is costly. Running PyTorch step by step in eager mode adds overhead from dispatching and kernel launches, and storing every intermediate activation for the backward pass uses a lot of GPU memory. A LinkedIn team presents Liger-Kernel, an open-source library of GPU kernels written in Triton and aimed specifically at LLM training. The kernels merge several operations into one so that data moves less often between slow HBM and fast on-chip SRAM. They also split inputs into chunks so that large intermediate tensors, such as the full logits matrix, never have to be held in memory. Compared with HuggingFace reference implementations on popular LLMs, the authors report about 20% higher training throughput on average and roughly 60% lower GPU memory use. The library depends only on PyTorch and Triton and works with common distributed training setups. Built-in benchmarks and tests check correctness and convergence. Because the gains require changing only a few lines of code, efficient training becomes accessible to practitioners who lack deep GPU-programming expertise.

Key points

  • The library offers three levels of integration:
    • AutoLigerKernelForCausalLM, which patches supported models automatically;
    • model-specific patching functions such as apply_liger_kernel_to_llama;
    • standalone modules like LigerLayerNorm and LigerCrossEntropyLoss for building custom models.
  • Implemented kernels include RMSNorm, LayerNorm, RoPE, SwiGLU, GeGLU and a Fused Linear CrossEntropy. The last one computes the loss block by block and never materializes the logits.
  • It is compatible with PyTorch FSDP, DeepSpeed ZeRO and ZeRO++. Continuous integration runs on NVIDIA, AMD and Intel GPUs.
  • In Medusa multi-token-prediction training of LLaMA 3-8B (8× A100 80 GB, batch size 4, bfloat16, AdamW), the fused CE kernel reduced memory and raised throughput. This matters with LLaMA-3's 128k vocabulary, where runs without Liger frequently ran out of memory.
  • The code is released under a permissive license at github.com/linkedin/Liger-Kernel.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🏢LinkedInSunnyvale · United States3 works · 91 citations

 Fields, topics and concepts

📚 Shelves 6The field of science this work belongs to (OpenAlex fields).Computer Science · 18,614 worksArtificial Intelligence · 17,965 worksArts and Humanities · 17,450 worksLanguage and Linguistics · 17,296 worksComputer Vision and Pattern Recognition · 3,486 worksHardware and Architecture · 167 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,894 works

Cite

Cite as Pin-Lun Hsu, Yun Dai, Vignesh Kothapalli, Qingquan Song, Shao Tang, Siyu Zhu, Steven Shimizu, Shivam Sahni, Haowen Ning and Yanning Chen. 2024. Liger Kernel: Efficient Triton Kernels for LLM Training. arXiv:2410.10989.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2410.10989 · arXiv 2410.10989

@misc{hsu2024liger,
  title = {{Liger Kernel: Efficient Triton Kernels for LLM Training}},
  author = {Pin-Lun Hsu, Yun Dai, Vignesh Kothapalli, Qingquan Song, Shao Tang, Siyu Zhu, Steven Shimizu, Shivam Sahni, Haowen Ning and Yanning Chen},
  year = {2024},
  doi = {10.48550/arxiv.2410.10989},
  eprint = {2410.10989},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2410.10989}
}

References

30 references, 17 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Martı́n Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. {TensorFlow}: a system for {Large-Scale} machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16), page… — TensorFlow: A system for large-scale machine learning
  2. Marah Abdin, Sam Ade Jacobs, Ammar Ahmad Awan, Jyoti Aneja, Ahmed Awadallah, Hany Awadalla, Nguyen Bach, Amit Bahree, Arash Bakhtiari, Harkirat Behl, et al. Phi-3 technical report: A highly capable language model locally on your phone. arXiv preprint arXiv:2404.14219, 2024.
  3. Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, et al. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International…
  4. Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. stat, 1050:21, 2016.
  5. Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. In Proceedings of the 34th International Conference on Neural Information Processing Systems, pages…
  6. Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple llm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401.10774, 2024.
  7. Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. TVM: An automated End-to-End optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pages 578–5…
  8. Yun Dai, Tejas Dharamsi, Byron Hsu, Tao Song, and Hamed Firooz. Enhancing stability for large models training in constrained bandwidth networks. arXiv preprint arXiv:2407.01614, 2024. — [arXiv:2407.01614](https://arxiv.org/abs/2407.01614) find
  9. Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691, 2023.
  10. Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv:2205.14135, 2022. find
  11. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024.
  12. Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning, 4(9), 2018. find
  13. Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016. find
  14. Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685, 2021.
  15. Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023. find
  16. Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tintore, Susan Zhang, Patrick Labatut, Daniel Haziza, Luca Wehrstedt, Jeremy Reizenstein, and Grigory Sizov. xFormers: A modular and hackable transformer modelling library. http… find
  17. Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
  18. Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 3505–3506, 2020. find
  19. Amit Sabne. XLA : Compiling machine learning for peak performance, 2020. find
  20. Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202, 2020. find
  21. J Su, Y Lu, S Pan, A Murtadha, B Wen, and Y Liu Roformer. Enhanced transformer with rotary position embedding., 2021. DOI: https://doi.org/10.1016/j. neucom, 2023.
  22. Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805, 2023.
  23. Philippe Tillet, Hsiang-Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19, 2019.
  24. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
  25. A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
  26. Guanhua Wang, Heyang Qin, Sam Ade Jacobs, Connor Holmes, Samyam Rajbhandari, Olatunji Ruwase, Feng Yan, Lei Yang, and Yuxiong He. Zero++: Extremely efficient collective communication for giant model training. arXiv preprint arXiv:2306.10209, 2023. find
  27. Jason Wei, Yi Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yogatama, Maarten Bosma, Denny Zhou, Donald Metzler, et al. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022.
  28. W Hwu Wen-Mei, David B Kirk, and Izzat El Hajj. Programming Massively Parallel Processors: A Hands-on Approach. Morgan Kaufmann, 2022. find
  29. Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems, 32, 2019. find
  30. Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, et al. Pytorch FSDP: Experiences on scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16(12):3848–3860, 2023. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/hsu-2024-liger-kernel-efficient-triton-kernels/