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 RMS
• LigerLayerNorm: Normalización de capa
• liger_rotary_pos_emb: Codificación de posición rotatoria (RoPE)
• LigerSwiGLUMLP: Función de activación SwiGLU
• LigerGEGLUMLP: Función de activación GeGLU
• LigerCrossEntropyLoss: Pérdida de entropía cruzada
• LigerFusedLinearCrossEntropyLoss: 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 DPO
• LigerFusedLinearORPOLoss: Pérdida ORPO
• LigerFusedLinearCPOLoss: Pérdida CPO
• LigerFusedLinearSimPOLoss: Pérdida SimPO
• LigerFusedLinearKTOLoss: 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.