LinkedIn에서 개발한 효율적인 Triton 커널 모음으로, 대규모 언어 모델 훈련에 최적화되어 훈련 속도를 20% 향상시키고 메모리 사용량을 60% 줄입니다.

Liger-Kernel 프로젝트 상세 소개

프로젝트 개요

Liger-Kernel은 LinkedIn에서 개발한 Triton 커널 모음으로, 대규모 언어 모델(LLM) 훈련을 위해 특별히 설계되었습니다. 이 프로젝트는 다중 GPU 훈련 처리량을 20% 효과적으로 향상시키고 메모리 사용량을 60% 줄일 수 있습니다. 프로젝트 이름 "Liger"는 "LinkedIn GPU Efficient Runtime"의 약자로, 효율적인 GPU 런타임이라는 핵심 이념을 반영합니다.

핵심 특징

성능 이점

• 훈련 속도 향상: 커널 융합, 제자리 대체 및 분할 기술을 통해 다중 GPU 훈련 처리량 20% 향상
• 메모리 효율성: 메모리 사용량 60% 감소, 더 긴 컨텍스트 길이, 더 큰 배치 크기 및 방대한 어휘 지원
• 후훈련 최적화: 정렬 및 증류 작업을 위해 후훈련 커널은 최대 80%의 메모리 절약 가능

기술 구현

• 정확한 계산: 근사 계산 없음, 순방향 및 역방향 전파 모두 엄격한 단위 테스트를 거침
• 경량 의존성: Torch 및 Triton만 필요하며 추가 라이브러리 의존성 없음
• 강력한 호환성: 즉시 사용 가능하며 Flash Attention, PyTorch FSDP 및 Microsoft DeepSpeed와 호환

지원 모델 및 작업

지원 모델 아키텍처

프로젝트는 다음과 같은 다양한 주요 대규모 언어 모델 아키텍처를 지원합니다.

• LLaMA 시리즈: LLaMA 2, LLaMA 3, LLaMA 3.2-Vision
• Mistral 시리즈: Mistral, Mixtral
• Gemma 시리즈: Gemma1, Gemma2, Gemma3
• Qwen 시리즈: Qwen2, Qwen2.5, Qwen2-VL, Qwen3 등
• 기타 모델: Phi3, Granite, OLMo2, GLM-4 등

핵심 커널 작업

프로젝트는 다양한 최적화된 커널 작업을 구현합니다.

기본 커널

• LigerRMSNorm: RMS 정규화
• LigerLayerNorm: 레이어 정규화
• liger_rotary_pos_emb: 회전 위치 임베딩(RoPE)
• LigerSwiGLUMLP: SwiGLU 활성화 함수
• LigerGEGLUMLP: GeGLU 활성화 함수
• LigerCrossEntropyLoss: 교차 엔트로피 손실
• LigerFusedLinearCrossEntropyLoss: 융합 선형 교차 엔트로피 손실

후훈련 커널

다양한 정렬 및 선호도 최적화 손실 함수 지원:

• LigerFusedLinearDPOLoss: DPO 손실
• LigerFusedLinearORPOLoss: ORPO 손실
• LigerFusedLinearCPOLoss: CPO 손실
• LigerFusedLinearSimPOLoss: SimPO 손실
• LigerFusedLinearKTOLoss: KTO 손실

사용 방법

1. 자동 통합 방식

from liger_kernel.transformers import AutoLigerKernelForCausalLM


model = AutoLigerKernelForCausalLM.from_pretrained("path/to/some/model")

2. 수동 패치 방식

import transformers
from liger_kernel.transformers import apply_liger_kernel_to_llama


apply_liger_kernel_to_llama()


apply_liger_kernel_to_llama(
    rope=True,
    swiglu=True,
    cross_entropy=True,
    fused_linear_cross_entropy=False,
    rms_norm=False
)


model = transformers.AutoModelForCausalLM("path/to/llama/model")

3. 저수준 API 방식

from liger_kernel.transformers import LigerFusedLinearCrossEntropyLoss
import torch.nn as nn
import torch

model = nn.Linear(128, 256).cuda()
loss_fn = LigerFusedLinearCrossEntropyLoss()

input = torch.randn(4, 128, requires_grad=True, device="cuda")
target = torch.randint(256, (4, ), device="cuda")
loss = loss_fn(model.weight, input, target)
loss.backward()

4. 후훈련 손실 사용 예시

from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss

orpo_loss = LigerFusedLinearORPOLoss()
y = orpo_loss(lm_head.weight, x, target)

설치 방법

안정 버전 설치

pip install liger-kernel

개발 버전 설치

pip install liger-kernel-nightly

소스 코드에서 설치

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

개발 환경 설치

pip install -e ".[dev]"

시스템 요구 사항

NVIDIA GPU 환경

• torch >= 2.1.2
• triton >= 2.3.0

AMD GPU 환경

• torch >= 2.5.0
• triton >= 3.0.0

기타 의존성

• transformers >= 4.x: transformers 모델 패치 API를 사용하는 경우

성능 벤치마크 테스트

벤치마크 테스트 조건:

• 모델: LLaMA 3-8B
• 배치 크기: 8
• 데이터 유형: bf16
• 최적화 프로그램: AdamW
• 기울기 검사점: 활성화
• 분산 전략: FSDP1, 8개의 A100 GPU

테스트 결과:

• Hugging Face 모델은 4K 컨텍스트 길이에서 메모리 오버플로가 발생하기 시작하지만 Hugging Face + Liger Kernel은 16K까지 확장 가능
• 훈련 처리량 20% 이상 향상
• 메모리 사용량 60% 감소

프레임워크 통합

Liger-Kernel은 여러 주요 훈련 프레임워크에 통합되었습니다.

• Axolotl
• LLaMa-Factory
• SFTTrainer
• Hugging Face Trainer
• SWIFT
• oumi

기술 원리

커널 융합 기술

여러 작업을 단일 커널로 융합하여 GPU 메모리 액세스 횟수를 줄이고 계산 효율성을 높입니다.

분할 계산

메모리 집약적인 작업의 경우 분할 처리 기술을 사용하여 대규모 계산을 작은 블록으로 분해하여 피크 메모리 사용량을 줄입니다.

제자리 작업

가능한 한 제자리 작업을 사용하여 추가 메모리 할당을 피하고 메모리 효율성을 더욱 최적화합니다.

요약

Liger-Kernel은 대규모 언어 모델 훈련 최적화의 중요한 진전을 나타냅니다. 정교하게 설계된 Triton 커널, 메모리 최적화 기술 및 광범위한 모델 지원을 통해 연구원과 엔지니어에게 훈련 효율성을 크게 향상시키고 계산 비용을 절감할 수 있는 강력하고 사용하기 쉬운 도구를 제공합니다. 프로젝트의 오픈 소스 특성과 활발한 커뮤니티 지원은 LLM 훈련 분야에서 중요한 리소스가 되도록 합니다.