liger-kerneltritonllm-traininggpu-optimizationpytorchtutorial

Liger Kernel: Fused Triton Kernels for LLM Training

September 23, 2026

Liger Kernel Architecture: Fused Triton Kernels for LLM Training, Memory Reduction, and Throughput Acceleration

The Problem: The Logits Memory Wall in LLM Training

When fine-tuning, aligning, or pre-training modern large language models in production data systems like AiFetchly, machine learning engineers routinely encounter unexpected GPU Out-of-Memory (OOM) exceptions.

Surprisingly, the primary consumer of GPU High-Bandwidth Memory (HBM) during backpropagation is often not the model’s static weights or optimizer states, but activation memory generated by standard un-fused PyTorch operators.

Among all activation bottlenecks in contemporary foundation models, the single largest memory killer is the Unembedding Linear Layer combined with Cross-Entropy Loss.

The Standard PyTorch Logits Memory Bottleneck:
┌────────────────────────────────────────────────────────────────────────────────────────┐
│ 1. Hidden State (H)          │ [Batch × SequenceLength, HiddenDim] (e.g. [8192, 4096]) │
├──────────────────────────────┼─────────────────────────────────────────────────────────┤
│ 2. Unembedding Projection    │ Matrix multiplication: Logits = H × W_head^T            │
├──────────────────────────────┼─────────────────────────────────────────────────────────┤
│ 3. Logits Tensor Materialized│ Shape: [8192, 128256] in FP32                           │
│                              │ Consumes ~4.20 Gigabytes of VRAM per micro-batch!       │
├──────────────────────────────┼─────────────────────────────────────────────────────────┤
│ 4. Stored for Backpropagation│ Retained across entire forward pass to compute dL/dLogits│
└──────────────────────────────┴─────────────────────────────────────────────────────────┘

1. The 128k Vocabulary Explosion

In modern frontier and open architectures—such as Meta LLaMA 3, Qwen 2.5, and DeepSeek-V3—the tokenizer vocabulary size ($V$) expanded from earlier baselines of 32,000 tokens to 128,000 or 129,280 tokens.

While larger vocabularies dramatically improve compression efficiency and multilingual tokenization, they multiply the memory footprint of the final loss layer by 4x:

$$\text{Logits Memory (Bytes)} = B \times S \times V \times \text{sizeof(float32)}$$

Consider fine-tuning a LLaMA-3 8B model with a modest batch size $\times$ sequence length of $B \times S = 8,192$ tokens:

Logits Tensor Size: $8,192 \times 128,256 \times 4 \text{ bytes} \approx \mathbf{4.20 \text{ GB}}$.
Backward Activation Buffer: Standard PyTorch Autograd retains the entire 4.2 GB FP32 logits tensor in global VRAM to compute the gradients $\frac{\partial \mathcal{L}}{\partial \text{logits}}$ during backpropagation.
If scaling to a sequence length of 32,768 tokens ($B \times S = 32,768$), the un-fused logits tensor alone requires 16.8 GB of VRAM!

This single transient tensor forces teams to reduce micro-batch sizes, incur expensive activation checkpointing overheads, or deploy complex tensor parallelism across multiple GPUs.

2. High-Bandwidth Memory (HBM) Roundtrips in Elementwise Layers

Beyond the unembedding layer, standard PyTorch implementations of transformer building blocks—such as RMSNorm, LayerNorm, SwiGLU, CrossEntropy, and Rotary Position Embeddings (RoPE)—are written as modular, separate operator calls:

# Standard PyTorch un-fused SwiGLU: launches 3 separate CUDA kernels!
gate = self.gate_proj(x)       # Global HBM Write
up = self.up_proj(x)           # Global HBM Write
activated = F.silu(gate) * up  # Global HBM Read + Math + Global HBM Write

Every un-fused operator reads inputs from global GPU HBM, performs simple arithmetic, and writes the output back to HBM. Because these elementwise operations have near-zero arithmetic intensity, Streaming Multiprocessors (SMs) spend up to 70% of their execution time waiting for memory bus transactions.

The Architecture: How Liger Kernel Optimizes Training

To eliminate this memory tax and restore hardware efficiency, open-source engineers at LinkedIn developed and released Liger Kernel: Efficient Triton Kernels for LLM Training (arXiv:2410.10989).

Liger Kernel is a comprehensive collection of drop-in, hardware-aligned GPU kernels implemented in OpenAI Triton. Its defining breakthrough is the Fused Linear Cross-Entropy Kernel, which eliminates the intermediate logits tensor entirely.

Liger Fused Linear Cross-Entropy Architecture: Eliminating Full Logits Materialization via Chunked Online LogSumExp in OpenAI Triton

1. Fused Linear Cross-Entropy with Online LogSumExp

In traditional PyTorch, the unembedding layer ($Z = H W^T$) and the loss function ($\mathcal{L} = \text{CrossEntropy}(Z, Y)$) are decoupled into two distinct operations.

Liger Kernel fuses the linear projection and the cross-entropy computation into a single Triton program. It processes tokens in streaming tiles or chunks along the sequence dimension:

Chunked Hidden State Streaming: Instead of projecting all $N = B \times S$ hidden states at once, the kernel streams a compact block of hidden states ($H_{\text{chunk}} \in \mathbb{R}^{\text{chunk_size} \times D}$) directly into high-speed GPU on-chip Shared Memory (SRAM).
In-SRAM Tile Matrix Multiplication: The tile matrix product $H_{\text{chunk}} W^T$ is computed inside SRAM.
Online LogSumExp Accumulation: Similar to how FlashAttention computes online softmax without materializing full attention matrices, Liger calculates the running maximum and denominator of the softmax on the fly: $$\text{LSE}(z) = m + \log\left(\sum_{j} \exp(z_j - m)\right), \quad \text{where } m = \max_j(z_j)$$
Immediate Target Extraction: The kernel extracts the target token logit $z_{y}$ directly within the kernel loop.
Zero Logits Materialization: The full $[N, V]$ logits matrix is never written to global HBM. Only the final scalar loss value and running normalizers are stored!

2. In-Kernel Gradient Computation During the Backward Pass

During backpropagation, standard PyTorch requires loading the saved $[N, V]$ logits tensor from HBM to compute:

$$\frac{\partial \mathcal{L}}{\partial z_i} = P(z_i) - \mathbb{I}(i = y) = \frac{\exp(z_i - \text{LSE})}{\sum_j \exp(z_j - \text{LSE})} - \mathbb{I}(i = y)$$

$$\frac{\partial \mathcal{L}}{\partial H} = \left(\frac{\partial \mathcal{L}}{\partial Z}\right) W, \quad \frac{\partial \mathcal{L}}{\partial W} = \left(\frac{\partial \mathcal{L}}{\partial Z}\right)^T H$$

Because Liger saved only the scalar $\text{LSE}$ per token rather than the full logits matrix, its custom Triton backward kernel:

Recomputes the necessary tile logits $H_{\text{chunk}} W^T$ inside SRAM.
Derives the gradient vector $\frac{\partial \mathcal{L}}{\partial z}$ directly in cache.
Multiplies by $W$ to immediately accumulate $\frac{\partial \mathcal{L}}{\partial H}$ and streams gradients back to earlier layers.

By trading a tiny amount of redundant in-SRAM computation for massive memory bandwidth reduction, Liger Kernel achieves up to 80% to 90% reduction in peak loss memory and accelerates the backward pass.

3. The Full Liger Fused Operator Suite

Beyond the Linear Cross-Entropy operator, Liger Kernel replaces all memory-bound transformer layers with optimized Triton implementations:

Operator	Standard PyTorch Behavior	Liger Kernel Triton Optimization
Linear + CrossEntropy	Materializes full $[B \times S, V]$ FP32 logits tensor (~4.2 GB)	Fused in-SRAM: Zero global logits tensor allocated (<35 MB buffer)
RMSNorm / LayerNorm	Separate mean, variance, scale, and cast kernel launches	Fused single-kernel pass: Combines squaring, summation, and scaling
SwiGLU / GeGLU	3 separate memory read/write passes across global HBM	Fused activation: Computes gating and elementwise product in SRAM
RoPE (Rotary Embedding)	Separate slicing, frequency generation, and rotation kernels	Fused coordinate rotation: Applies trigonometric rotation in registers
Kullback-Leibler (KL) Divergence	Materializes dual probability tensors for policy and reference	Fused streaming KL: Ideal for RLHF, PPO, and GRPO alignment pipelines

Implementation: Simulating Fused Linear Cross-Entropy in Python

To understand how Liger Kernel eliminates logits materialization, let us implement a reference simulation in Python comparing the standard un-fused PyTorch workflow against Liger’s chunked online streaming algorithm:

import numpy as np
from typing import Tuple, Dict, Any

class LinearCrossEntropyBenchmark:
    """
    Demonstrates the mathematical mechanics of standard un-fused Cross-Entropy
    versus Liger Kernel's chunked streaming formulation.
    """
    def __init__(self, batch_seq: int = 128, hidden_dim: int = 512, vocab_size: int = 4096):
        self.N = batch_seq    # Batch size * Sequence length
        self.D = hidden_dim   # Hidden dimension
        self.V = vocab_size   # Vocabulary size

    def run_standard_unfused(
        self, H: np.ndarray, W: np.ndarray, targets: np.ndarray
    ) -> Tuple[float, np.ndarray, int]:
        """
        Simulates standard PyTorch un-fused execution:
        1. Materializes full [N, V] logits tensor in global memory.
        2. Computes cross-entropy loss over full matrix.
        3. Backpropagates gradients through logits buffer.
        """
        # Step 1: Materialize full [N, V] logits in memory
        logits = np.dot(H, W.T)  # Shape: [N, V]
        peak_logits_bytes = logits.nbytes

        # Step 2: Numerically stable softmax
        max_logits = np.max(logits, axis=-1, keepdims=True)
        exp_logits = np.exp(logits - max_logits)
        sum_exp = np.sum(exp_logits, axis=-1, keepdims=True)
        log_probs = (logits - max_logits) - np.log(sum_exp)

        # Cross-entropy loss: -log(P(target))
        loss = -np.mean(log_probs[np.arange(self.N), targets])

        # Step 3: Gradient with respect to logits: (P - 1) at target
        probs = exp_logits / sum_exp
        grad_logits = probs.copy()
        grad_logits[np.arange(self.N), targets] -= 1.0
        grad_logits /= self.N

        # Gradient with respect to hidden states H
        grad_H = np.dot(grad_logits, W)

        return float(loss), grad_H, peak_logits_bytes

    def run_liger_fused_streaming(
        self, H: np.ndarray, W: np.ndarray, targets: np.ndarray, chunk_size: int = 16
    ) -> Tuple[float, np.ndarray, int]:
        """
        Simulates Liger Kernel's Fused Triton streaming execution:
        Processes hidden states in small chunks, computing online LogSumExp
        and gradients without ever allocating the full [N, V] tensor.
        """
        total_loss = 0.0
        grad_H = np.zeros_like(H)
        max_tile_bytes = 0

        # Stream sequence in chunks (analogous to Triton threadblock tiles)
        for i in range(0, self.N, chunk_size):
            chunk_H = H[i : i + chunk_size]
            chunk_targets = targets[i : i + chunk_size]
            chunk_len = chunk_H.shape[0]

            # In Triton, this matrix tile is kept in high-speed SRAM / registers
            chunk_logits = np.dot(chunk_H, W.T)
            max_tile_bytes = max(max_tile_bytes, chunk_logits.nbytes)

            # Online LogSumExp computation
            max_val = np.max(chunk_logits, axis=-1, keepdims=True)
            exp_logits = np.exp(chunk_logits - max_val)
            sum_exp = np.sum(exp_logits, axis=-1, keepdims=True)
            log_sum_exp = np.log(sum_exp) + max_val

            # Accumulate scalar loss for target tokens
            target_logits = chunk_logits[np.arange(chunk_len), chunk_targets]
            chunk_loss = np.sum(log_sum_exp.squeeze(-1) - target_logits)
            total_loss += chunk_loss

            # Fused gradient computation inside kernel
            probs = exp_logits / sum_exp
            probs[np.arange(chunk_len), chunk_targets] -= 1.0
            probs /= self.N

            grad_H[i : i + chunk_size] = np.dot(probs, W)

        avg_loss = total_loss / self.N
        return float(avg_loss), grad_H, max_tile_bytes

    def evaluate(self) -> Dict[str, Any]:
        np.random.seed(42)
        H = np.random.randn(self.N, self.D).astype(np.float32)
        W = np.random.randn(self.V, self.D).astype(np.float32)
        targets = np.random.randint(0, self.V, size=(self.N,))

        loss_std, grad_std, mem_std = self.run_standard_unfused(H, W, targets)
        loss_fused, grad_fused, mem_fused = self.run_liger_fused_streaming(
            H, W, targets, chunk_size=16
        )

        return {
            "loss_std": loss_std,
            "loss_fused": loss_fused,
            "loss_diff": abs(loss_std - loss_fused),
            "max_grad_diff": float(np.max(np.abs(grad_std - grad_fused))),
            "mem_std_bytes": mem_std,
            "mem_fused_bytes": mem_fused,
            "mem_savings_pct": (1.0 - (mem_fused / mem_std)) * 100.0,
        }


def print_liger_benchmark():
    print("=" * 70)
    print("LIGER FUSED LINEAR CROSS-ENTROPY BENCHMARK SIMULATION")
    print("=" * 70)

    bench = LinearCrossEntropyBenchmark(batch_seq=128, hidden_dim=512, vocab_size=4096)
    res = bench.evaluate()

    print(f"Tokens (Batch x SeqLen):    128")
    print(f"Hidden Dimension:           512")
    print(f"Vocabulary Size:            4,096")
    print(f"Standard Logits Memory:     {res['mem_std_bytes'] / 1024:.2f} KB (Full matrix allocated)")
    print(f"Liger Streaming Memory:     {res['mem_fused_bytes'] / 1024:.2f} KB (Tile SRAM buffer)")
    print(f"Peak Buffer Reduction:      {res['mem_savings_pct']:.1f}%\n")

    print("NUMERICAL EQUIVALENCE CHECK:")
    print(f"  Standard Loss:            {res['loss_std']:.6f}")
    print(f"  Liger Fused Loss:         {res['loss_fused']:.6f}")
    print(f"  Absolute Loss Error:      {res['loss_diff']:.2e}")
    print(f"  Max Gradient Error:       {res['max_grad_diff']:.2e}\n")

    # Production LLaMA-3 Scale Projection
    prod_tokens = 8192
    prod_vocab = 128256
    prod_std_gb = (prod_tokens * prod_vocab * 4) / (1024**3)
    prod_liger_mb = (64 * prod_vocab * 4) / (1024**2)  # 64-token chunk

    print("PRODUCTION SCALE PROJECTION (LLaMA-3 / Qwen-2.5 128k Vocab, B*S=8192):")
    print(f"  PyTorch Full Logits Tensor: {prod_std_gb:.2f} GB VRAM per micro-batch!")
    print(f"  Liger Fused Tile Buffer:    {prod_liger_mb:.2f} MB VRAM")
    print(f"  Total VRAM Freed per Step:  {prod_std_gb - (prod_liger_mb / 1024):.2f} GB")
    print("=" * 70)

    assert res["loss_diff"] < 1e-4, "Loss values must match numerically"
    assert res["max_grad_diff"] < 1e-4, "Gradients must match numerically"
    print("  [SUCCESS] Liger Kernel mathematical equivalence verified!")


if __name__ == "__main__":
    print_liger_benchmark()

Integrating Liger Kernel in Hugging Face and TRL

One of Liger Kernel’s primary design goals is seamless drop-in compatibility. You do not need to rewrite your training loops or re-implement custom forward passes.

1. One-Line Monkey Patching for Hugging Face Transformers

To patch existing Hugging Face model architectures with Liger’s fused Triton kernels, install the package and apply the model-specific monkey patch before instantiating your model:

pip install liger-kernel

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from liger_kernel.transformers import apply_liger_kernel_to_llama

# Apply Liger Kernel patch BEFORE loading model weights
apply_liger_kernel_to_llama(
    rope=True,
    rms_norm=True,
    swiglu=True,
    cross_entropy=True,
    fused_linear_cross_entropy=True
)

# Model now natively uses fused Triton kernels!
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

Liger provides automated patching functions across all major model families:

apply_liger_kernel_to_llama()
apply_liger_kernel_to_qwen2()
apply_liger_kernel_to_mistral()
apply_liger_kernel_to_gemma2()
apply_liger_kernel_to_phi3()

2. Native Integration with TRL SFTTrainer

If you use Hugging Face TRL (Transformer Reinforcement Learning) for Supervised Fine-Tuning (SFT) or Direct Preference Optimization (DPO), pass use_liger_kernel=True directly inside TrainingArguments or SFTConfig:

from trl import SFTTrainer, SFTConfig
from datasets import load_dataset

training_args = SFTConfig(
    output_dir="./llama3-liger-finetuned",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    max_seq_length=8192,
    bf16=True,
    logging_steps=10,
    use_liger_kernel=True,  # Enables all fused Triton kernels
)

trainer = SFTTrainer(
    model="meta-llama/Meta-Llama-3-8B",
    train_dataset=load_dataset("imdb", split="train"),
    args=training_args,
)

trainer.train()

3. Distributed Training with FSDP and DeepSpeed ZeRO

Liger Kernel is fully compatible with multi-GPU distributed training frameworks:

PyTorch FSDP (Fully Sharded Data Parallel): Compatible with standard FULL_SHARD and HYBRID_SHARD strategies.
DeepSpeed ZeRO-2 and ZeRO-3: Functions out of the box because the forward and backward kernels adhere strictly to standard PyTorch autograd contracts.
Gradient Checkpointing: Because fused linear cross-entropy eliminates the largest activation tensor, many training configurations can disable activation checkpointing entirely, achieving a free 30% speedup by avoiding re-computation.

Production Performance Benchmarks

In empirical evaluations conducted across NVIDIA H100 and A100 GPU clusters:

Liger Kernel Performance Summary on LLaMA-3 8B (Sequence Length 8192):
┌─────────────────────────┬───────────────────┬───────────────────┬──────────────┐
│ Benchmark Metric        │ PyTorch Baseline  │ With Liger Kernel │ Net Dividend │
├─────────────────────────┼───────────────────┼───────────────────┼──────────────┤
│ Peak Activation VRAM    │ 28.4 GB           │ 11.6 GB           │ -59.1% VRAM  │
│ Training Throughput     │ 2,140 tokens/sec  │ 2,680 tokens/sec  │ +25.2% speed │
│ Max Safe Context Window │ 8k tokens         │ 32k tokens        │ 4x expansion │
│ Max Micro-Batch Size    │ 2                 │ 6                 │ 3x capacity  │
└─────────────────────────┴───────────────────┴───────────────────┴──────────────┘

59% Reduction in Peak Activation Memory: By eliminating the FP32 logits buffer and fusing elementwise operations, total training VRAM drops by more than half.
25% Throughput Speedup: Eliminating redundant HBM read/write transactions allows Tensor Cores to operate closer to peak saturation.
4x Context Length Expansion: Workloads that previously ran out of memory at 8k tokens scale smoothly to 32k tokens on identical hardware.

Conclusion

Liger Kernel addresses one of the most glaring architectural inefficiencies in deep learning infrastructure: the massive memory penalty of un-fused elementwise operations and intermediate logits buffers.

By leveraging OpenAI Triton to fuse Linear projections directly with Cross-Entropy Loss, RMSNorm, and SwiGLU, Liger Kernel slashes training memory by up to 60% while accelerating overall throughput by 20% to 30%. With seamless drop-in integrations across Hugging Face Transformers, TRL, FSDP, and DeepSpeed, Liger Kernel represents an essential component of the modern LLM training toolkit.