Eine effiziente Triton-Kernel-Sammlung, entwickelt von LinkedIn, speziell optimiert für das Training großer Sprachmodelle, die die Trainingsgeschwindigkeit um 20 % erhöht und den Speicherverbrauch um 60 % reduziert.

Liger-Kernel Projekt – Detaillierte Beschreibung

Projektübersicht

Liger-Kernel ist eine von LinkedIn entwickelte Sammlung von Triton-Kerneln, die speziell für das Training großer Sprachmodelle (LLM) konzipiert wurde. Das Projekt kann den Multi-GPU-Trainingsdurchsatz effektiv um 20 % steigern und die Speichernutzung um 60 % reduzieren. Der Projektname "Liger" steht für "LinkedIn GPU Efficient Runtime" und spiegelt die Kernidee einer effizienten GPU-Laufzeit wider.

Kernfunktionen

Leistungsvorteile

• Verbesserte Trainingsgeschwindigkeit: Durch Kernel-Fusion, In-Place-Ersetzungen und Chunking-Techniken wird der Multi-GPU-Trainingsdurchsatz um 20 % gesteigert.
• Speichereffizienz: Die Speichernutzung wird um 60 % reduziert, wodurch längere Kontextlängen, größere Batch-Größen und riesige Vokabulare unterstützt werden.
• Post-Training-Optimierung: Post-Training-Kernel sparen bis zu 80 % Speicher bei Alignment- und Destillationsaufgaben.

Technische Umsetzung

• Präzise Berechnung: Keine Näherungsrechnungen, Vorwärts- und Rückwärtsausbreitung werden strengen Unit-Tests unterzogen.
• Geringe Abhängigkeiten: Benötigt nur Torch und Triton, keine zusätzlichen Bibliotheksabhängigkeiten.
• Hohe Kompatibilität: Sofort einsatzbereit, kompatibel mit Flash Attention, PyTorch FSDP und Microsoft DeepSpeed.

Unterstützte Modelle und Operationen

Unterstützte Modellarchitekturen

Das Projekt unterstützt verschiedene gängige Architekturen großer Sprachmodelle, darunter:

• LLaMA-Serie: LLaMA 2, LLaMA 3, LLaMA 3.2-Vision
• Mistral-Serie: Mistral, Mixtral
• Gemma-Serie: Gemma1, Gemma2, Gemma3
• Qwen-Serie: Qwen2, Qwen2.5, Qwen2-VL, Qwen3 usw.
• Andere Modelle: Phi3, Granite, OLMo2, GLM-4 usw.

Kern-Kernel-Operationen

Das Projekt implementiert verschiedene optimierte Kernel-Operationen:

Basiskernel

• LigerRMSNorm: RMS-Normalisierung
• LigerLayerNorm: Layer-Normalisierung
• liger_rotary_pos_emb: Rotatorische Positionskodierung (RoPE)
• LigerSwiGLUMLP: SwiGLU-Aktivierungsfunktion
• LigerGEGLUMLP: GeGLU-Aktivierungsfunktion
• LigerCrossEntropyLoss: Kreuzentropie-Verlust
• LigerFusedLinearCrossEntropyLoss: Verschmolzener linearer Kreuzentropie-Verlust

Post-Training-Kernel

Unterstützt verschiedene Alignment- und Präferenzoptimierungs-Verlustfunktionen:

• LigerFusedLinearDPOLoss: DPO-Verlust
• LigerFusedLinearORPOLoss: ORPO-Verlust
• LigerFusedLinearCPOLoss: CPO-Verlust
• LigerFusedLinearSimPOLoss: SimPO-Verlust
• LigerFusedLinearKTOLoss: KTO-Verlust

Verwendung

1. Automatische Integrationsmethode

from liger_kernel.transformers import AutoLigerKernelForCausalLM


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

2. Manuelle Patch-Methode

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. Low-Level-API-Methode

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. Beispiel für die Verwendung von Post-Training-Verlust

from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss

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

Installation

Installation der stabilen Version

pip install liger-kernel

Installation der Entwicklungsversion

pip install liger-kernel-nightly

Installation aus dem Quellcode

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

Installation der Entwicklungsumgebung

pip install -e ".[dev]"

Systemanforderungen

NVIDIA GPU-Umgebung

• torch >= 2.1.2
• triton >= 2.3.0

AMD GPU-Umgebung

• torch >= 2.5.0
• triton >= 3.0.0

Andere Abhängigkeiten

• transformers >= 4.x: Wenn die Transformers-Modell-Patch-API verwendet wird

Leistungsbenchmarks

Benchmark-Bedingungen:

• Modell: LLaMA 3-8B
• Batch-Größe: 8
• Datentyp: bf16
• Optimierer: AdamW
• Gradienten-Checkpointing: Aktiviert
• Verteilte Strategie: FSDP1, 8 A100 GPUs

Die Testergebnisse zeigen:

• Hugging Face-Modelle beginnen bei einer Kontextlänge von 4K mit Speicherauslastung, während Hugging Face + Liger Kernel auf 16K skaliert werden kann.
• Der Trainingsdurchsatz wird um über 20 % gesteigert.
• Die Speichernutzung wird um 60 % reduziert.

Framework-Integration

Liger-Kernel ist in mehrere gängige Trainingsframeworks integriert:

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

Technische Prinzipien

Kernel-Fusionstechnologie

Durch die Fusion mehrerer Operationen in einen einzigen Kernel werden die GPU-Speicherzugriffe reduziert und die Recheneffizienz gesteigert.

Chunking-Berechnung

Für speicherintensive Operationen wird die Chunking-Technik verwendet, um große Berechnungen in kleine Blöcke zu zerlegen und die maximale Speichernutzung zu reduzieren.

In-Place-Operationen

Verwenden Sie nach Möglichkeit In-Place-Operationen, um zusätzliche Speicherzuweisungen zu vermeiden und die Speichereffizienz weiter zu optimieren.

Zusammenfassung

Liger-Kernel stellt einen wichtigen Fortschritt in der Optimierung des Trainings großer Sprachmodelle dar. Durch sorgfältig entworfene Triton-Kernel, Speicheroptimierungstechniken und breite Modellunterstützung bietet es Forschern und Ingenieuren ein leistungsstarkes und einfach zu bedienendes Werkzeug, das die Trainingseffizienz erheblich steigern und die Rechenkosten senken kann. Die Open-Source-Natur des Projekts und die aktive Community-Unterstützung machen es zu einer wichtigen Ressource im Bereich des LLM-Trainings.