Apresentação Detalhada do Projeto Liger-Kernel

Visão Geral do Projeto

Liger-Kernel é uma coleção de kernels Triton desenvolvida pelo LinkedIn, projetada especificamente para o treinamento de grandes modelos de linguagem (LLM). Este projeto é capaz de aumentar efetivamente o throughput de treinamento multi-GPU em 20% e reduzir o uso de memória em 60%. O nome do projeto "Liger" representa "LinkedIn GPU Efficient Runtime", refletindo sua filosofia central de runtime GPU eficiente.

Principais Características

Vantagens de Desempenho

• Aumento da Velocidade de Treinamento: Através de fusão de kernels, substituição in-place e técnicas de chunking, o throughput de treinamento multi-GPU é aumentado em 20%.
• Eficiência de Memória: O uso de memória é reduzido em 60%, suportando comprimentos de contexto mais longos, tamanhos de lote maiores e vocabulários massivos.
• Otimização Pós-Treinamento: Para tarefas de alinhamento e destilação, os kernels de pós-treinamento podem economizar até 80% de memória.

Implementação Técnica

• Cálculo Preciso: Sem cálculos aproximados, o forward e o backward propagation são rigorosamente testados unitariamente.
• Dependências Leves: Requer apenas Torch e Triton, sem dependências de bibliotecas adicionais.
• Alta Compatibilidade: Pronto para uso, compatível com Flash Attention, PyTorch FSDP e Microsoft DeepSpeed.

Modelos e Operações Suportadas

Arquiteturas de Modelo Suportadas

O projeto suporta várias arquiteturas de grandes modelos de linguagem, incluindo:

• 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.
• Outros Modelos: Phi3, Granite, OLMo2, GLM-4, etc.

Operações de Kernel Essenciais

O projeto implementou várias operações de kernel otimizadas:

Kernels Básicos

• LigerRMSNorm: Normalização RMS
• LigerLayerNorm: Normalização de Camada
• liger_rotary_pos_emb: Codificação de Posição Rotativa (RoPE)
• LigerSwiGLUMLP: Função de Ativação SwiGLU
• LigerGEGLUMLP: Função de Ativação GeGLU
• LigerCrossEntropyLoss: Perda de Entropia Cruzada
• LigerFusedLinearCrossEntropyLoss: Perda de Entropia Cruzada Linear Fundida

Kernels de Pós-Treinamento

Suporta várias funções de perda de alinhamento e otimização de preferência:

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

Como Usar

1. Método de Integração Automática

from liger_kernel.transformers import AutoLigerKernelForCausalLM


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

2. Método de Patch 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 Baixo Nível

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. Exemplo de Uso de Perda de Pós-Treinamento

from liger_kernel.chunked_loss import LigerFusedLinearORPOLoss

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

Instalação

Instalação da Versão Estável

pip install liger-kernel

Instalação da Versão de Desenvolvimento

pip install liger-kernel-nightly

Instalação a Partir do Código Fonte

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

Instalação do Ambiente de Desenvolvimento

pip install -e ".[dev]"

Requisitos do Sistema

Ambiente NVIDIA GPU

• torch >= 2.1.2
• triton >= 2.3.0

Ambiente AMD GPU

• torch >= 2.5.0
• triton >= 3.0.0

Outras Dependências

• transformers >= 4.x: Se usar a API de patch de modelo transformers

Benchmarks de Desempenho

Condições de benchmark:

• Modelo: LLaMA 3-8B
• Tamanho do Lote: 8
• Tipo de Dados: bf16
• Otimizador: AdamW
• Checkpoint de Gradiente: Ativado
• Estratégia Distribuída: FSDP1, 8 GPUs A100

Os resultados dos testes mostram:

• Os modelos Hugging Face começam a apresentar estouro de memória em comprimentos de contexto de 4K, enquanto Hugging Face + Liger Kernel podem ser escalados para 16K.
• Aumento do throughput de treinamento em mais de 20%.
• Redução do uso de memória em 60%.

Integração com Frameworks

Liger-Kernel foi integrado em vários frameworks de treinamento principais:

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

Princípios Técnicos

Técnica de Fusão de Kernel

Ao fundir várias operações em um único kernel, o número de acessos à memória da GPU é reduzido, aumentando a eficiência computacional.

Computação em Chunk

Para operações com uso intensivo de memória, a técnica de processamento em chunk é usada para decompor grandes cálculos em pequenos chunks, reduzindo o uso máximo de memória.

Operações In-Place

Sempre que possível, operações in-place são usadas para evitar alocação de memória adicional, otimizando ainda mais a eficiência da memória.

Conclusão

Liger-Kernel representa um avanço importante na otimização do treinamento de grandes modelos de linguagem. Através de kernels Triton cuidadosamente projetados, técnicas de otimização de memória e amplo suporte a modelos, ele fornece aos pesquisadores e engenheiros uma ferramenta poderosa e fácil de usar que pode aumentar significativamente a eficiência do treinamento e reduzir os custos computacionais. A natureza de código aberto do projeto e o suporte ativo da comunidade o tornam um recurso importante no campo do treinamento de LLM.