Source-linked AI summary
Liger Kernel: Efficient Triton Kernels for LLM Training
Pin-Lun Hsu, Yun Dai, Vignesh Kothapalli, Qingquan Song, Shao Tang, Siyu Zhu, Steven Shimizu, Shivam Sahni, Haowen Ning, Yanning Chen
TL;DR
LLM training faces efficiency bottlenecks and complex optimization demands. Liger-Kernel provides an open-source Triton-kernel library using fusion and related kernel techniques; benchmarks report reduced memory use and improved throughput in two Medusa training modes.
Problem
LLM training efficiency is challenged by infrastructure bottlenecks and the complexity of optimization.
Source: §1 · p. 1Method
Liger-Kernel is an open-source library of Triton kernels for LLM training that streamlines tensor operations using kernel fusion.
Source: §1 · p. 1Results
In both benchmarked Medusa training modes, Liger kernels showed reduced memory usage and improved throughput.
Source: §4.2 · p. 13Takeaways & Limitations
The library is designed for flexible training use, and its techniques may also be adapted to optimize model inference.
Source: §5 · p. 15The report focuses on performance benchmarking; generating effective LLM heads to accelerate LLaMA3-8B inference is outside its scope.
Source: §4.2 · p. 14
Abstract
from arXiv · showhide
Training Large Language Models (LLMs) efficiently at scale presents a formidable challenge, driven by their ever-increasing computational demands and the need for enhanced performance. In this work, we introduce Liger-Kernel, an open-sourced set of Triton kernels developed specifically for LLM training. With kernel optimization techniques like kernel operation fusing and input chunking, our kernels achieve on average a 20% increase in training throughput and a 60% reduction in GPU memory usage for popular LLMs compared to HuggingFace implementations. In addition, Liger-Kernel is designed with modularity, accessibility, and adaptability in mind, catering to both casual and expert users. Comprehensive benchmarks and integration tests are built in to ensure compatibility, performance, correctness, and convergence across diverse computing environments and model architectures. The source code is available under a permissive license at: github.com/linkedin/Liger-Kernel.
1 Introduction
LLM training efficiency is constrained by infrastructure bottlenecks and the complexity of optimization. Liger-Kernel addresses this with an open-source Triton-kernel library designed to streamline and customize training.
1 Introduction
- LLM training efficiency depends on infrastructure stability and faces memory-management and tensor-operation latency-bandwidth bottlenecks.The passage also describes kernel-level fusion as an optimization opportunity for reducing memory copying and improving parallel efficiency.Source: §1 · p. 1
- Liger-Kernel is an open-source library of efficient Triton kernels for LLM training.It uses kernel fusion and offers integration with distributed frameworks including PyTorch FSDP and DeepSpeed ZeRO.Source: §1 · p. 1
- The library offers a flexible interface for novice users and modular components and adaptive layer configurations for advanced users.It requires PyTorch and Triton and is designed to work across computing environments and hardware platforms.Source: §1 · p. 1
2 Preliminaries
The preliminaries explain the overheads of eager execution and review compilation and operation-fusion approaches to GPU optimization. They position Triton kernels within this landscape and note related implementations that inform Liger-Kernel.
2 Preliminaries
- Eager PyTorch execution adds dispatch, function-call, and kernel-launch overhead, while materializing intermediate activations also consumes GPU memory.Custom Triton operation fusion is one approach practitioners use to replace native PyTorch execution.Source: §2 · pp. 1–2
2.1 Model Compiler
- Model compilers translate high-level model descriptions into optimized low-level code for specialized hardware.The passage names torch.compile, TVM, XLA, and nvFuser and describes their differing backends and optimization targets.Source: §2.1 · p. 2
2.2 An Algorithmic Perspective of Operation Fusion
- GPU operation fusion combines standalone operations to reduce per-operation time and memory overhead from step-by-step execution.FlashAttention uses GPU memory hierarchies and smaller blocks to reduce attention memory complexity from quadratic to linear; FlashAttention-2 further addresses register spilling and parallelism.Source: §2.2 · p. 2
2.3 Custom Operation Fusion with Triton
- Triton provides a Python-like language and compiler for high-performance GPU kernels, and Liger-Kernel draws on related Triton projects as references.The cited examples include Transformer building blocks, Triton implementations for fine-tuning and inference, and block-wise fused linear projection with cross-entropy.Source: §2.3 · pp. 2–3
3.1 API Design
Liger-Kernel's API aims to fit existing code with minimal disruption while supporting different levels of customization. Users can apply automatic or model-specific patches, or compose individual kernels into custom models.
3.1 API Design
- The API is designed to be minimally disruptive to existing code while allowing varying levels of customization.The passage presents several ways to apply Liger kernels, ranging from automatic use to custom composition.Source: §3.1 · p. 3
- AutoLigerKernelForCausalLM automatically patches supported model types without requiring model-specific patching API imports.The example loads a model through the AutoLigerKernelForCausalLM class.Source: §3.1 · p. 3
- Model-specific patching APIs offer finer control and can be used beyond causal language models, including sequence classification.The example applies a LLaMA patch before loading a sequence-classification model.Source: §3.1 · p. 3
- Advanced users can compose individual Liger kernels to build custom models, including using LigerLayerNorm and LigerCrossEntropyLoss in a LigerTransformer.These flexible options support integration into varied workflows for LLM training and deployment.Source: §3.1 · pp. 3–4
3.2 Kernels
The section presents fused Triton kernels for normalization, rotary embeddings, gated activations, and cross-entropy. For fused linear cross-entropy, it chunks hidden-state processing to address the memory cost of materializing logits.
RMSNorm and LayerNorm
- RMSNorm fuses normalization and scaling, while LayerNorm centers the input and applies learned scale and bias.The RMSNorm scale parameter and LayerNorm scale and bias are shared across input vectors, so their gradients are summed.Source: §3.2 · p. 4
RoPE
- RoPE fuses query and key rotation into one kernel and uses the HuggingFace rotation-matrix form rather than Su et al. (2023)'s form.The implementation adopts an efficient computation described in Su et al. (2023) for the sparse rotation matrix.Source: §3.2 · p. 5
SwiGLU and GeGLU
- SwiGLU and GeGLU fuse element-wise gated projections; SwiGLU uses β=1 SiLU, while GeGLU uses a tanh-approximate GELU.For both, the input is transformed by two learned affine projections before the element-wise gated computation.Source: §3.2 · pp. 5–6
CrossEntropy (CE)
- CrossEntropy moves gradient computation into the forward function, replaces logits in place, and computes gradients with online softmax.The implementation also uses a safe log operation to avoid numerical instabilities.Source: §3.2 · p. 6
FusedLinearCrossEntropy (FLCE)
- Fused linear cross-entropy processes hidden states in chunks before cross-entropy, addressing the memory cost of materializing large logit tensors.The linear projection is applied sequentially to chunks, whose logits are passed to the cross-entropy kernel; chunk size balances memory allocation and processing speed.Source: §3.2 · pp. 6–7
3.3 Testing Best Practices
The testing approach emphasizes exactness and evaluates correctness, speed and memory use, convergence behavior, and tensor contiguity. It includes irregular input shapes and checks results under conditions intended to reflect practical training.
3.3 Testing Best Practices
- Testing is treated as central to kernel development, with exactness framed as essential to precision and reliability.The stated testing approach draws on both research and practical experience.Source: §3.3 · p. 7
3.3.1 Correctness
- Correctness tests compare Triton kernels with a PyTorch baseline across regular and irregular shapes and multiple data types.The stated tolerances are atol=10^-7 and rtol=10^-5 for fp32, and atol=10^-3 and rtol=10^-2 for bf16; large-dimension indexing uses int64 to avoid int32 overflow.Source: §3.3.1 · p. 8
3.3.2 Performance
- Performance tests assess both speed and memory use using training-like input dimensions, including batch size 4 and hidden dimension 2048.Sequence length varies across the tested shapes.Source: §3.3.2 · p. 8
3.3.3 Convergence Test
- Convergence tests mimic practical tensor contiguity, shapes, and data types at smaller scale, then verify logits, weights, and loss at training end.The tests check that these values are exact.Source: §3.3.3 · p. 8
3.3.4 Contiguity
- Non-contiguous tensors can cause illegal memory access or incorrect outputs; a non-contiguous attention derivative in RoPE production training caused significant loss divergence.The authors recommend ensuring tensors are contiguous before passing them to a kernel.Source: §3.3.4 · p. 8
3.4 Integrations
Liger integrates with several popular training frameworks, typically requiring only a flag to patch model code with Liger kernels.
3.4 Integrations
- Liger integrates with Hugging Face Trainer and SFTTrainer, Axolotl, and LLaMA-Factory, typically with a single flag and minimal code changes.The integrations are presented as demonstrating API flexibility and ease of use.Source: §3.4 · p. 8
- The example uses TRL’s SFTTrainer with use_liger=True to load a model through AutoLigerKernelForCausalLM.The code excerpt names meta-llama/Meta-Llama-3-8B as the model.Source: §3.4 · p. 8
- Correctness tests may need tolerances relaxed by one or two orders of magnitude, and convergence tests are used when tolerance must be loose.Source: §3.4 · p. 8
4 Numerical Experiments
Across kernel-level and end-to-end benchmarks, Liger kernels improve speed and/or memory use; the experiments also include Medusa training, while the report focuses on performance benchmarking.
4.1 Kernel Benchmark
- At vocab size 163840, CrossEntropy runs approximately 3× faster and uses approximately 5× less memory; GeGLU and SwiGLU cut peak memory roughly 1.6× at sequence length 16384.GeGLU and SwiGLU maintain baseline speed; their backward pass recomputes SiLU(·) and GELU(·) outputs.Source: §4.1 · p. 9
- At hidden dimension 16384, RMSNorm is approximately 7× faster with roughly 3× lower peak memory, while LayerNorm execution time falls approximately 30% with minimal memory overhead.RMSNorm fuses normalization and scaling and caches root mean square values; LayerNorm caches the inverse root mean square for backward use.Source: §4.1 · pp. 9–10
- At hidden size 16384, RoPE achieves approximately 8× speedup with approximately 3× lower memory consumption.The kernel represents the rotation matrix as a flattened 1D tensor and leverages repeated blocks.Source: §4.1 · p. 10
4.2 Usecase Benchmark
- LLaMA 3-8B: throughput +42.8%, GPU memory −54.8%; Qwen2: throughput +25.5%, GPU memory −56.8%; Gemma: throughput +11.9%, GPU memory −51.8%; Mistral: throughput +27%, GPU memory −21%; Phi3: throughput +17%, GPU memory −13%.The respective batch sizes were 64, 48, 48, 128, and 128.Source: §4.2 · p. 10, Figure 4 · p. 11, Figure 5 · p. 11
- Both Medusa training stages show reduced memory use and improved throughput with Liger; without it, experiments are highly prone to out-of-memory issues.Stage 1 trains only added Medusa heads while freezing the backbone; Stage 2 tunes the backbone and heads together. The end-to-end benchmarks use four A100 GPUs to fine-tune five LLMs on Alpaca.Source: §4.2 · pp. 10, 13
- The report focuses on performance benchmarking and does not evaluate generating effective LLM heads to accelerate LLaMA 3-8B inference.The report notes that this work requires training-data selection, hyperparameter tuning, and warmup techniques for convergence.Source: §4.2 · p. 14
5 Conclusions
Liger-Kernel aims to make efficient LLM training broadly usable through accessible APIs, performance-focused kernels, and operational and community commitments. The authors also propose adapting its techniques to model inference.
5 Conclusions
- Liger-Kernel combines a user-friendly API and popular-framework integration with optimized Triton kernels for efficient training.The project aspires to become a preferred open-source library for efficient, scalable LLM training.Source: §5 · p. 14
- Ease-of-use priorities include intuitive APIs, broad model support, and wide hardware compatibility.Source: §5 · p. 14
- The performance focus pairs computational efficiency with exactness.Source: §5 · p. 14
- Operational priorities include stable continuous integration, rigorous testing, and an active community, alongside ecosystem engagement through events and industry collaborations.Source: §5 · p. 14
- The authors state that the same techniques used for training can be adapted to optimize model inference.Source: §5 · p. 15
6 Contributors and Acknowledgements
The contributor acknowledgements describe work spanning kernel implementation, interfaces, testing, framework integration, model support, and team leadership.
6.1 Core Contributors
- Pin-Lun Hsu led the project and architected and implemented kernels, the public interface, and test suite.Source: §6.1 · p. 15
- Yun Dai contributed RoPE, GeGLU, and Fused Linear CrossEntropy work; Vignesh Kothapalli implemented Fused Linear CrossEntropy and designed its scaling and sharding formula.Source: §6.1 · p. 15
- Qingquan Song implemented SwiGLU and led convergence testing and PyTorch Lightning integration; Shao Tang implemented Layer Norm variants and derived gradient formulas.Source: §6.1 · p. 15
- Siyu Zhu adapted Fused Linear CrossEntropy for Medusa and benchmarked it, while leading Hugging Face integration; Steven Shimizu also improved that integration and tests.Source: §6.1 · p. 15
- Shivam Sahni expanded model support; Haowen Ning led LLM training infrastructure, and Yanning Chen managed the team.Source: §6.1 · p. 15