DOCSEARCHbyAlgolia
Submit Docs

Liger Kernel

repository·main·Indexed 27 days ago

https://github.com/linkedin/liger-kernel

A collection of efficient Triton kernels designed to increase throughput and reduce memory usage during LLM training. It provides exact, Hugging Face compatible implementations of common layers and loss functions, including RMSNorm, RoPE, SwiGLU, and CrossEntropy. Liger Kernel supports patching for various model architectures (e.g., Llama, Mistral, Gemma, Qwen) and integrates with PyTorch FSDP, DeepSpeed, and multiple trainer frameworks.

Tokens
17.4K
Snippets
43
Records
71
Agent score
92%

What's inside liger-kernel

  1. Overview of Liger Kernel

    main

    Liger Kernel is a collection of Triton kernels specifically designed to optimize LLM training. It provides Hugging Face compatible implementations of layers like RMSNorm, RoPE, SwiGLU, and CrossEntropy.

    Key benefits include:

    • Increased Throughput: Up to 20% increase in multi-GPU training throughput.
    • Reduced Memory Usage: Up to 60% reduction in memory usage for standard training, and up to 80% for post-training tasks (alignment/distillation).
    • Compatibility: Works out of the box with Flash Attention, PyTorch FSDP, and Microsoft DeepSpeed.
    • Exact Computation: No approximations are used; kernels are mathematically exact and verified against standard implementations.

  2. Use Liger FlexChunkLoss for Alignment and Distillation

    main
    Liger FlexChunkLoss provides optimized implementations for post-training loss functions such as alignment (DPO, ORPO, CPO, KTO) and distillation. It delivers up to 80% memory savings and a 10% throughput boost by using chunking and fused kernel optimizations. It achieves this by fusing the final linear layer with loss computation and calculating backward gradients during the forward pass to reduce intermediate activation storage.

  3. Liger Kernel Key Features

    main

    Liger Kernel offers several advantages for LLM developers:

    • Ease of use: Patch Hugging Face models with one line of code or compose models using Liger Kernel modules.
    • Efficiency: Uses kernel fusion, in-place replacement, and chunking for layers like RMSNorm, RoPE, SwiGLU, and CrossEntropy.
    • Lightweight: Minimal dependencies, requiring only Torch and Triton.
    • Multi-GPU support: Compatible with PyTorch FSDP, DeepSpeed, and DDP.
    • Trainer Framework Integration: Supports Axolotl, LLaMa-Factory, SFTTrainer, Hugging Face Trainer, SWIFT, and oumi.

  4. Install Liger Kernel for ROCm

    main

    To use Liger Kernel on ROCm, you must first install ROCm PyTorch from the official PyTorch index, then install Liger Kernel from source.

    1. Install ROCm PyTorch:

    pip3 install torch torchvision --index-url https://download.pytorch.org/whl/rocm7.2

    2. Install Liger Kernel from source:

    git clone https://github.com/linkedin/Liger-Kernel.git
cd Liger-Kernel
pip install -e .

  5. Configure custom Cross-Entropy for Megatron-LM

    main
    For training setups that require explicit kernel configuration (such as custom ignore_index or label_smoothing), do not use the high-level apply_liger_kernel_to_megatron patch. Instead, instantiate LigerMegatronCrossEntropy directly and wire it into your model.

  6. Apply model-specific patching APIs

    main

    For more control, you can use model-specific patching APIs to swap Hugging Face models with optimized Liger kernels. You can either call the patching function without arguments to apply all default optimizations, or specify exactly which kernels to enable/disable.

    Example for Llama models:

    • apply_liger_kernel_to_llama(): Automatically monkey-patches the model with default Liger kernels.
    • apply_liger_kernel_to_llama(rope=True, swiglu=True, ...): Allows granular control over which kernels (like rope, swiglu, cross_entropy, fused_linear_cross_entropy, or rms_norm) are applied.
    import transformers
from liger_kernel.transformers import apply_liger_kernel_to_llama

# 1a. Adding this line automatically monkey-patches the model with the optimized Liger kernels
apply_liger_kernel_to_llama()

# 1b. You could alternatively specify exactly which kernels are applied
apply_liger_kernel_to_llama(
  rope=True,
  swiglu=True,
  cross_entropy=True,
  fused_linear_cross_entropy=False,
  rms_norm=False
)

# 2. Instantiate patched model
model = transformers.AutoModelForCausalLM("path/to/llama/model")

  8. Set up the Liger-Kernel development environment

    main

    To develop on Liger-Kernel, clone the repository, install the package in editable mode with development dependencies, and set up pre-commit hooks using prek.

    If you encounter a no matches found: .[dev] error during installation, use the quoted syntax.

  9. Run tests and check code style

    main

    Use the provided Makefile to run the full test suite or pytest for specific tests.

    Makefile commands:

    • make test: Runs correctness tests.
    • make checkstyle: Ensures code follows the project style.
    • make test-convergence: Runs convergence tests.

    Pytest command: To run a specific test function in a file: python -m pytest test_sample.py::test_function_name

    make test
make checkstyle
make test-convergence

# Run a single test function
python -m pytest test_sample.py::test_function_name

  10. Run a Non-model dimension sweep (D1)

    main

    A D1 sweep varies non-model dimensions (like sequence length) while keeping the model configuration fixed.

    Implementation Steps:

    1. Parse args and resolve model via get_benchmark_model_config(args.model).
    2. Compute sweep config using compute_seq_len_sweep_config() (for seq_len) or use BT directly.
    3. Build x_values (e.g., powers of 2) and extra_benchmark_configs with fixed dimensions.
    4. Execute via run_benchmarks().

    CLI Usage:

    # Default model (llama_3_8b)
python benchmark_geglu.py

# Specific model
python benchmark_geglu.py --model llama_2_7b

# Overwrite existing CSV entries
python benchmark_geglu.py --model llama_3_8b --overwrite
    python benchmark_geglu.py --model llama_3_8b --overwrite

  11. Add a new benchmark script

    main

    To add a new benchmark for a kernel, create a Python script in the benchmark/scripts/ directory using the naming convention benchmark_<kernel_name>.py (e.g., benchmark_geglu.py).

    To ensure accuracy, import reference (non-Liger) kernels from the existing test suite (e.g., test/transformers/test_<kernel>.py) to use as baselines. This prevents code duplication and keeps benchmarks in sync with tests.

  12. Run kernel benchmarks via CLI

    main

    Benchmarks are executed by running scripts located in benchmark/scripts/. Results are saved to benchmark/data/all_benchmark_data.csv.

    To run a model configuration sweep:

    python scripts/benchmark_kto_loss.py --sweep-mode model_config [--model llama_3_8b]

    To run a token length sweep:

    python scripts/benchmark_kto_loss.py [--sweep-mode token_length] [--bt 2048]
    cd benchmark
python scripts/benchmark_kto_loss.py --sweep-mode model_config [--model llama_3_8b]
python scripts/benchmark_kto_loss.py [--sweep-mode token_length] [--bt 2048]