Liger Kernel — Efficient Triton Kernels for LLM Training

Byron 의 진단으로 시작하는 강의. flash attention 이 attention 의 HBM 트래픽을 잡은 뒤, 학습 루프의 새 병목은 elementwise op 들의 launch overhead와 cross_entropy 의 logits memory spike 다. Liger Kernel 은 그 두 자리를 정확히 노린다.

강의가 답하려는 두 줄 —

1. 학습 루프에서 시간이 어디로 새는가 — Hugging Face Llama 의 forward + backward 를 profiler 로 까보면 무엇이 dominant 인가.
2. 그 자리를 어떻게 Triton 으로 다시 짜는가 — 그리고 사용자가 코드 한 줄도 바꾸지 않고 어떻게 채택하게 하는가 (monkey-patch).

강의의 인지적 frame 은 분명하다 — “flash attention 이 attention 의 fusion 을 풀었으니, 우리는 attention 바깥의 fusion 을 푼다”. cross_entropy, RMSNorm, RoPE, SwiGLU. 학습 루프의 hot path 의 60% 가 이 4 개 안에 있다는 게 Byron 의 측정 결과.

강의 끝에 손에 잡혀야 할 것은 — 네 커널의 fusion 자리(cross_entropy / RMSNorm / RoPE / SwiGLU), 각 커널의 메모리 절감 메커니즘, 그리고 monkey-patch 한 줄로 채택하는 패턴이다.

Byron 이 강의 첫 데모로 띄운 것 — Hugging Face 의 vanilla Llama-3 8B 학습 step 의 PyTorch profiler trace. 두 가지가 시각적으로 잡힌다 — elementwise 커널의 dense 한 launch와 cross_entropy 직후의 메모리 곡선 spike.

Byron 이 “이게 Liger 의 가장 영향력 있는 단일 커널이다” 라고 못 박은 자리. 학습 시 cross_entropy 의 logits 텐서가 backward 까지 메모리에 살아있어야 한다는 점이 핵심 문제 — 큰 vocab 모델에서 메모리 곡선의 spike 가 여기서 생긴다.

  GB
  ▲                       ▲ logits
80┤                      / \   spike
70┤                     /   \
60┤   ─── activations ─/     \─── backward
50┤  ─
40┤ ─
30┤─
  └────────────────────────────────────►
   forward                       backward

   ▲ logits spike 가 30 GB 추가 → fused 로 0 추가

# Liger 의 fused cross_entropy 사용 — 통상 PyTorch 대체
from liger_kernel.transformers import LigerCrossEntropyLoss

loss_fn = LigerCrossEntropyLoss(reduction="mean")
loss    = loss_fn(logits.view(-1, vocab),
                  labels.view(-1))      # logits 를 인자로 받지만 ...
# 더 강력한 변형 — LM head 의 weight 를 받아 logits 자체를 안 만든다
from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
loss = LigerFusedLinearCrossEntropyLoss()(
    hidden_states,                  # (B*S, hidden)
    lm_head_weight,                 # (vocab, hidden)
    labels)                         # logits = h @ W^T 가 inside-kernel 에서 chunk 별로

RMSNorm 은 LLaMA / Mistral / Qwen 의 표준 normalization. 식 자체는 짧지만, eager 모드에서는 4–7 개의 elementwise 커널 시퀀스 로 분해된다. Liger 의 RMSNorm 커널이 이걸 한 커널 안에 묶는다.

# Liger RMSNorm 커널의 핵심 — Triton
@triton.jit
def rms_norm_forward(X, W, Y, RSTD, eps,
                     N: tl.constexpr, BLOCK: tl.constexpr):
    pid    = tl.program_id(0)
    X_row  = X + pid * N
    Y_row  = Y + pid * N

    cols   = tl.arange(0, BLOCK)
    mask   = cols < N
    x      = tl.load(X_row + cols, mask=mask, other=0.).to(tl.float32)
    w      = tl.load(W + cols, mask=mask, other=0.).to(tl.float32)

    # 한 커널 안에서 mean(x²) → rsqrt → multiply → cast 까지
    var    = tl.sum(x * x, axis=0) / N
    rstd   = 1.0 / tl.sqrt(var + eps)
    y      = (x * rstd) * w

    tl.store(Y_row + cols, y.to(tl.bfloat16), mask=mask)
    tl.store(RSTD + pid, rstd)        # backward 에서 재사용

RoPE(Rotary Position Embedding)는 LLaMA 이후 표준이 된 위치 인코딩. 식 자체는 sin/cos 와의 elementwise multiply 두 줄 이지만 — 매 layer 의 Q, K 양쪽에 적용되니까 모델 전체에서 호출 횟수가 layer × 2 = 64+ 회. 그 누적이 launch overhead 의 큰 부분.

# RoPE 의 수학 — half-rotation 패턴
# q_rot = q * cos + rotate_half(q) * sin
# 여기서 rotate_half 는 [x1,x2] -> [-x2,x1] 의 swap

@triton.jit
def rope_forward(Q, K, COS, SIN,
                  Q_OUT, K_OUT, ...):
    pid_b = tl.program_id(0)        # batch · head · seq
    # Q · K 두 텐서를 같은 program 에서 처리
    q   = tl.load(Q + ...)
    k   = tl.load(K + ...)
    cos = tl.load(COS + ...)
    sin = tl.load(SIN + ...)

    # half-rotation: 한 row 안에서 self-shuffle
    q1, q2 = split(q)
    q_rot  = concat(q1*cos - q2*sin,
                    q2*cos + q1*sin)

    tl.store(Q_OUT + ..., q_rot)
    # K 도 같은 자리에서 — Q/K 두 launch 가 한 launch 로

LLaMA 의 MLP 는 down_proj(silu(gate_proj(x)) * up_proj(x)). 이 식의 끝 — silu(gate) * up — 이 elementwise activation. eager 에서는 silu 한 번 + multiply 한 번 + 그 사이 임시 텐서 하나가 박힌다. Liger SwiGLU 커널이 그걸 한 커널로.

강의에서 Byron 이 가장 시간을 많이 쓴 자리 중 하나 — Liger 의 모든 커널은 학습 convergence 가 baseline 과 일치한다는 검증. 정확도 차이가 일정 수준을 넘으면 학습이 망가진다.

  loss
   ▲
2.4┤▓
   │▓▒
2.2┤  ▓▒
   │   ▓▒
2.0┤    ▓▒  ▓ vanilla
   │      ▓▒
1.8┤        ▓▒  ▒ Liger fused
   │          ▓▒
1.6┤────────────────────────────────►
   0    250M    500M    750M   1B tok

   두 곡선이 visible 하게 겹친다 — diff < 1%

강의에서 Byron 이 가장 자랑한 디자인 결정 — 사용자가 모델 코드를 한 줄도 바꾸지 않는다. 한 import + 한 함수 호출로, 모델 안 RMSNorm/RoPE/SwiGLU/cross_entropy 가 자동으로 Liger 커널로 교체된다.

from liger_kernel.transformers import apply_liger_kernel_to_llama

# Hugging Face 모델 로드 그대로
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")

# ← 한 줄. 모델 안 RMSNorm/RoPE/SwiGLU 가 monkey-patch 됨
apply_liger_kernel_to_llama(
    rope=True,
    rms_norm=True,
    swiglu=True,
    fused_linear_cross_entropy=True,
)

# 그 후로는 Hugging Face Trainer 그대로
trainer = Trainer(model=model, args=...)
trainer.train()                    # ← Liger 커널이 자동으로 끼어듬

Liger 의 효과는 모델/세팅에 따라 다르다. 강의에서 Byron 이 보여준 표를 큰 그림으로 정리.

손에 새기기 — 실습 시퀀스

1. baseline 학습 step profiling — Hugging Face Llama 의 한 step 을 PyTorch profiler 로 잡는다. RMSNorm/RoPE/SwiGLU/cross_entropy 가 차지하는 % 를 확인.
2. Liger 적용 후 같은 profile — apply_liger_kernel_to_llama 한 줄 추가. 같은 step 의 trace 를 비교 — 커널 수가 얼마나 줄었는지.
3. RMSNorm 직접 짜보기 — Liger 의 RMSNorm 코드를 보지 않고 first principle 로. row-per-program, BLOCK = pow2(N), upcast fp32. 그 후 Liger 와 diff.
4. cross_entropy fusion 이해 — fused linear CE 가 logits 를 안 만드는 trick 을 코드로 본다. chunk 단위 hidden × W^T → softmax → NLL 의 흐름.
5. numerics gate 재현 — bf16 input · fp32 accum · bf16 output. baseline 과 output diff 가 1e-3 이하인지 확인.
6. monkey-patch 직접 짜보기 — 자기가 가진 모델(예: Mistral) 의 RMSNorm class 를 동적으로 교체하는 짧은 함수. 모델 코드 그대로.
7. 학습 1B token convergence — Liger on/off 두 가지로 1B token 학습. loss curve 를 plot 해서 두 곡선이 visible 하게 겹치는지 확인.
8. vocab 큰 모델에서 효과 측정 — Qwen2 같은 152K vocab 모델에서 fused linear CE 의 메모리 절감을 직접 측정.

• torch.compile + Liger 의 호환성 — torch.compile 의 inductor backend 가 같은 자리에서 fusion 하려고 한다. 두 시스템이 어떻게 합쳐지는가? graph break 없이 monkey-patch 가 작동하는지.
• flash attention 과의 fusion — Liger 는 attention 바깥. attention 자체와 RMSNorm/RoPE 를 더 fusion 할 여지가 있는가? fa3 의 tile-based 디자인과 Liger 패턴의 합집합.
• FP8 시대의 Liger — bf16 + fp32 accum 이 현재 패턴. H100/H200 의 FP8 mma 와 어떻게 통합되는가? Transformer Engine 과의 조율.
• Inference 에도 적용 가능한가 — Liger 는 학습 위주. inference 에서도 같은 패턴이 win 인가? cross_entropy 가 없으니 효과가 다를 듯.
• 큰 vocab 의 한계 — vocab 이 1M 이상으로 가면 fused linear CE 도 chunk 단위가 너무 커진다. 어떤 점에서 split 이 필요한가.
• 비-Llama 모델 지원 — RWKV, Mamba 같은 다른 아키텍처에 같은 패턴이 적용되는가? 이미 patch 함수가 모델별 분리됐는데, 새 모델 추가의 cost 는 어디.