Présentation détaillée du projet Liger-Kernel

Aperçu du projet

Liger-Kernel est une collection de kernels Triton développée par LinkedIn, spécialement conçue pour l'entraînement de grands modèles de langage (LLM). Ce projet permet d'améliorer efficacement le débit d'entraînement multi-GPU de 20 % et de réduire l'utilisation de la mémoire de 60 %. Le nom du projet "Liger" signifie "LinkedIn GPU Efficient Runtime", reflétant son concept central de runtime GPU efficace.

Caractéristiques principales

Avantages en termes de performances

• Amélioration de la vitesse d'entraînement : Grâce à la fusion de kernels, au remplacement sur place et aux techniques de découpage, le débit d'entraînement multi-GPU est amélioré de 20 %.
• Efficacité de la mémoire : L'utilisation de la mémoire est réduite de 60 %, ce qui permet de prendre en charge des longueurs de contexte plus importantes, des tailles de lots plus grandes et des vocabulaires massifs.
• Optimisation post-entraînement : Pour les tâches d'alignement et de distillation, les kernels de post-entraînement peuvent économiser jusqu'à 80 % de la mémoire.

Implémentation technique

• Calcul précis : Pas de calculs approximatifs, les passes avant et arrière sont rigoureusement testées unitairement.
• Dépendances légères : Nécessite uniquement Torch et Triton, sans dépendances de bibliothèques supplémentaires.
• Forte compatibilité : Prêt à l'emploi, compatible avec Flash Attention, PyTorch FSDP et Microsoft DeepSpeed.

Modèles et opérations pris en charge

Architectures de modèles prises en charge

Le projet prend en charge plusieurs architectures de grands modèles de langage courants, notamment :

• Série LLaMA : LLaMA 2, LLaMA 3, LLaMA 3.2-Vision
• Série Mistral : Mistral, Mixtral
• Série Gemma : Gemma1, Gemma2, Gemma3
• Série Qwen : Qwen2, Qwen2.5, Qwen2-VL, Qwen3, etc.
• Autres modèles : Phi3, Granite, OLMo2, GLM-4, etc.

Opérations de kernel principales

Le projet a implémenté plusieurs opérations de kernel optimisées :

Kernels de base

• LigerRMSNorm : Normalisation RMS
• LigerLayerNorm : Normalisation de couche
• liger_rotary_pos_emb : Encodage de position rotatif (RoPE)
• LigerSwiGLUMLP : Fonction d'activation SwiGLU
• LigerGEGLUMLP : Fonction d'activation GeGLU
• LigerCrossEntropyLoss : Perte d'entropie croisée
• LigerFusedLinearCrossEntropyLoss : Perte d'entropie croisée linéaire fusionnée

Kernels de post-entraînement

Prise en charge de plusieurs fonctions de perte d'alignement et d'optimisation des préférences :

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

Méthodes d'utilisation

1. Méthode d'intégration automatique

from liger_kernel.transformers import AutoLigerKernelForCausalLM


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

2. Méthode de patch manuel

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. Méthode d'API de bas niveau

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. Exemple d'utilisation de la perte de post-entraînement

from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss

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

Méthodes d'installation

Installation de la version stable

pip install liger-kernel

Installation de la version de développement

pip install liger-kernel-nightly

Installation à partir du code source

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

Installation de l'environnement de développement

pip install -e ".[dev]"

Configuration système requise

Environnement GPU NVIDIA

• torch >= 2.1.2
• triton >= 2.3.0

Environnement GPU AMD

• torch >= 2.5.0
• triton >= 3.0.0

Autres dépendances

• transformers >= 4.x : Si vous utilisez l'API de patch de modèle transformers

Tests de performance

Conditions des tests de performance :

• Modèle : LLaMA 3-8B
• Taille du lot : 8
• Type de données : bf16
• Optimiseur : AdamW
• Checkpoint de gradient : Activé
• Stratégie distribuée : FSDP1, 8 GPU A100

Les résultats des tests montrent :

• Les modèles Hugging Face commencent à manquer de mémoire à une longueur de contexte de 4K, tandis que Hugging Face + Liger Kernel peuvent s'étendre à 16K.
• Amélioration du débit d'entraînement de plus de 20 %.
• Réduction de l'utilisation de la mémoire de 60 %.

Intégration du framework

Liger-Kernel a été intégré à plusieurs frameworks d'entraînement courants :

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

Principes techniques

Technologie de fusion de kernels

En fusionnant plusieurs opérations en un seul kernel, le nombre d'accès à la mémoire GPU est réduit, ce qui améliore l'efficacité du calcul.

Calcul par blocs

Pour les opérations gourmandes en mémoire, une technique de traitement par blocs est utilisée, décomposant les grands calculs en petits blocs, ce qui réduit l'utilisation maximale de la mémoire.

Opérations sur place

Dans la mesure du possible, des opérations sur place sont utilisées pour éviter l'allocation de mémoire supplémentaire, ce qui optimise davantage l'efficacité de la mémoire.

Conclusion

Liger-Kernel représente une avancée importante dans l'optimisation de l'entraînement des grands modèles de langage. Grâce à des kernels Triton soigneusement conçus, des techniques d'optimisation de la mémoire et une large prise en charge des modèles, il offre aux chercheurs et aux ingénieurs un outil puissant et facile à utiliser, capable d'améliorer considérablement l'efficacité de l'entraînement et de réduire les coûts de calcul. La nature open source du projet et le soutien actif de la communauté en font une ressource importante dans le domaine de l'entraînement des LLM.