Una colección eficiente de kernels Triton desarrollada por LinkedIn, optimizada para el entrenamiento de grandes modelos de lenguaje, que puede aumentar la velocidad de entrenamiento en un 20% y reducir el uso de memoria en un 60%.
Descripción detallada del proyecto Liger-Kernel
Resumen del proyecto
Liger-Kernel es una colección de kernels de Triton desarrollada por LinkedIn, diseñada específicamente para el entrenamiento de modelos de lenguaje grandes (LLM). Este proyecto puede mejorar eficazmente el rendimiento del entrenamiento multi-GPU en un 20% y reducir el uso de memoria en un 60%. El nombre del proyecto "Liger" significa "LinkedIn GPU Efficient Runtime", lo que refleja su concepto central de tiempo de ejecución de GPU eficiente.
Características principales
Ventajas de rendimiento
- Mejora de la velocidad de entrenamiento: A través de la fusión de kernels, el reemplazo in situ y las técnicas de fragmentación, el rendimiento del entrenamiento multi-GPU aumenta en un 20%.
- Eficiencia de memoria: El uso de memoria se reduce en un 60%, lo que permite longitudes de contexto más largas, tamaños de lote más grandes y vocabularios masivos.
- Optimización posterior al entrenamiento: Para las tareas de alineación y destilación, los kernels posteriores al entrenamiento pueden ahorrar hasta un 80% de memoria.
Implementación técnica
- Cálculo preciso: Sin cálculos aproximados, la propagación hacia adelante y hacia atrás se someten a pruebas unitarias rigurosas.
- Dependencias ligeras: Solo requiere Torch y Triton, sin dependencias de bibliotecas adicionales.
- Alta compatibilidad: Listo para usar, compatible con Flash Attention, PyTorch FSDP y Microsoft DeepSpeed.
Modelos y operaciones compatibles
Arquitecturas de modelos compatibles
El proyecto admite varias arquitecturas de modelos de lenguaje grandes convencionales, que incluyen:
- Serie LLaMA: LLaMA 2, LLaMA 3, LLaMA 3.2-Vision
- Serie Mistral: Mistral, Mixtral
- Serie Gemma: Gemma1, Gemma2, Gemma3
- Serie Qwen: Qwen2, Qwen2.5, Qwen2-VL, Qwen3, etc.
- Otros modelos: Phi3, Granite, OLMo2, GLM-4, etc.
Operaciones de kernel centrales
El proyecto implementa varias operaciones de kernel optimizadas:
Kernels básicos
LigerRMSNorm
: Normalización RMSLigerLayerNorm
: Normalización de capaliger_rotary_pos_emb
: Codificación de posición rotatoria (RoPE)LigerSwiGLUMLP
: Función de activación SwiGLULigerGEGLUMLP
: Función de activación GeGLULigerCrossEntropyLoss
: Pérdida de entropía cruzadaLigerFusedLinearCrossEntropyLoss
: Pérdida de entropía cruzada lineal fusionada
Kernels posteriores al entrenamiento
Admite varias funciones de pérdida de optimización de alineación y preferencia:
LigerFusedLinearDPOLoss
: Pérdida DPOLigerFusedLinearORPOLoss
: Pérdida ORPOLigerFusedLinearCPOLoss
: Pérdida CPOLigerFusedLinearSimPOLoss
: Pérdida SimPOLigerFusedLinearKTOLoss
: Pérdida KTO
Modo de empleo
1. Método de integración automática
from liger_kernel.transformers import AutoLigerKernelForCausalLM
model = AutoLigerKernelForCausalLM.from_pretrained("path/to/some/model")
2. Método de parche manual
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étodo de API de bajo nivel
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. Ejemplo de uso de pérdida posterior al entrenamiento
from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss
orpo_loss = LigerFusedLinearORPOLoss()
y = orpo_loss(lm_head.weight, x, target)
Método de instalación
Instalación de la versión estable
pip install liger-kernel
Instalación de la versión de desarrollo
pip install liger-kernel-nightly
Instalación desde el código fuente
git clone https://github.com/linkedin/Liger-Kernel.git
cd Liger-Kernel
pip install -e .
Instalación del entorno de desarrollo
pip install -e ".[dev]"
Requisitos del sistema
Entorno NVIDIA GPU
torch >= 2.1.2
triton >= 2.3.0
Entorno AMD GPU
torch >= 2.5.0
triton >= 3.0.0
Otras dependencias
transformers >= 4.x
: Si utiliza la API de parche de modelos de transformers
Pruebas comparativas de rendimiento
Condiciones de prueba comparativa:
- Modelo: LLaMA 3-8B
- Tamaño del lote: 8
- Tipo de datos: bf16
- Optimizador: AdamW
- Punto de control de gradiente: habilitado
- Estrategia distribuida: FSDP1, 8 GPU A100
Los resultados de las pruebas muestran:
- Los modelos de Hugging Face comienzan a experimentar desbordamiento de memoria con una longitud de contexto de 4K, mientras que Hugging Face + Liger Kernel se puede escalar a 16K.
- El rendimiento del entrenamiento aumenta en más del 20%.
- El uso de memoria se reduce en un 60%.
Integración de frameworks
Liger-Kernel se ha integrado en varios frameworks de entrenamiento convencionales:
- Axolotl
- LLaMa-Factory
- SFTTrainer
- Hugging Face Trainer
- SWIFT
- oumi
Principio técnico
Tecnología de fusión de kernels
Al fusionar varias operaciones en un solo kernel, se reduce el número de accesos a la memoria de la GPU y se mejora la eficiencia computacional.
Cálculo de fragmentación
Para las operaciones con uso intensivo de memoria, se adopta la tecnología de procesamiento de fragmentación, que descompone los cálculos grandes en fragmentos pequeños, lo que reduce el uso máximo de memoria.
Operaciones in situ
Utilice operaciones in situ tanto como sea posible para evitar la asignación de memoria adicional y optimizar aún más la eficiencia de la memoria.
Resumen
Liger-Kernel representa un importante avance en la optimización del entrenamiento de modelos de lenguaje grandes. A través de kernels de Triton cuidadosamente diseñados, técnicas de optimización de memoria y un amplio soporte de modelos, proporciona a investigadores e ingenieros una herramienta potente y fácil de usar que puede mejorar significativamente la eficiencia del entrenamiento y reducir los costos computacionales. La naturaleza de código abierto del proyecto y el soporte activo de la comunidad lo convierten en un recurso importante en el campo del entrenamiento de LLM.