Axolotl - Ferramenta Poderosa de Pós-Treinamento para Modelos de IA

Visão Geral do Projeto

Axolotl é uma ferramenta de código aberto projetada especificamente para simplificar vários processos de pós-treinamento de modelos de IA. Pós-treinamento refere-se a qualquer modificação ou treinamento adicional realizado em um modelo pré-treinado, incluindo ajuste fino completo do modelo, ajuste fino eficiente em termos de parâmetros (como LoRA e QLoRA), ajuste fino supervisionado (SFT), ajuste de instruções e técnicas de alinhamento. A ferramenta suporta diversas arquiteturas de modelo e configurações de treinamento, permitindo que os usuários comecem facilmente a usar essas tecnologias avançadas de treinamento.

Principais Características

Suporte a Modelos

• Arquiteturas de Modelo Diversificadas: Suporta o treinamento de vários modelos Hugging Face, incluindo LLaMA, Pythia, Falcon, MPT, Mistral, Mixtral e outros modelos de linguagem grandes populares.
• Métodos de Treinamento Flexíveis: Suporta ajuste fino completo (Full Fine-tuning), LoRA, QLoRA, ReLoRA, GPTQ e outros métodos de treinamento.

Gerenciamento de Configuração

• Arquivos de Configuração YAML: Usa arquivos YAML simples para incluir todas as configurações necessárias, como pré-processamento de conjuntos de dados, treinamento/ajuste fino de modelos, inferência ou avaliação de modelos.
• Substituição via CLI: Suporta a substituição de configurações nos arquivos de configuração por meio de argumentos de linha de comando.
• Configuração Flexível: Permite personalizar vários parâmetros de treinamento e configurações de modelo.

Capacidade de Processamento de Dados

• Conjuntos de Dados em Múltiplos Formatos: Suporta o carregamento de conjuntos de dados locais, HuggingFace e na nuvem (S3, Azure, GCP, OCI).
• Formato Personalizado: Pode usar formatos personalizados ou importar diretamente conjuntos de dados já tokenizados.
• Pré-processamento de Conjuntos de Dados: Possui funcionalidades poderosas de pré-processamento de dados integradas.

Otimização de Desempenho

• Técnicas Avançadas de Otimização: Integra xformers, Flash Attention, Liger Kernel, Rope Scaling e técnicas de multi-packing.
• Suporte a Múltiplos GPUs: Suporta treinamento em GPU única ou múltipla via FSDP ou DeepSpeed.
• Treinamento Eficiente: Otimizado para GPUs NVIDIA (Ampere ou mais recentes, com suporte para bf16 e Flash Attention) e GPUs AMD.

Implantação e Monitoramento

• Pronto para a Nuvem: Fornece imagens Docker e pacotes PyPI, que podem ser usados em plataformas de nuvem e hardware local.
• Registro de Resultados: Suporta o registro de resultados e checkpoints no WandB, MLflow ou Comet.
• Suporte a Monitoramento: Integra várias ferramentas de rastreamento e monitoramento de experimentos.

Requisitos Técnicos

Requisitos de Hardware

• GPU NVIDIA (Ampere ou mais recentes, para bf16 e Flash Attention) ou GPU AMD.
• Memória GPU suficiente para o treinamento do modelo.

Requisitos de Software

• Python 3.11
• PyTorch ≥2.4.1
• Pacotes de dependência relacionados.

Instalação

Instalação Rápida

pip3 install -U packaging==23.2 setuptools==75.8.0 wheel ninja
pip3 install --no-build-isolation axolotl[flash-attn,deepspeed]

# Baixar arquivos de configuração de exemplo
axolotl fetch examples
axolotl fetch deepspeed_configs  # Opcional

Instalação a partir do Código Fonte

git clone https://github.com/axolotl-ai-cloud/axolotl.git
cd axolotl
pip3 install -U packaging setuptools wheel ninja
pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'

Via Docker

docker run --gpus '"all"' --rm -it axolotlai/axolotl:main-latest

Como Usar

Fluxo de Uso Básico

1. Obter Configurações de Exemplo:

   axolotl fetch examples
   

2. Treinar Modelo:

   axolotl train examples/llama-3/lora-1b.yml
   

3. Personalizar Configuração: Modifique os parâmetros no arquivo de configuração YAML conforme necessário.

Estrutura do Arquivo de Configuração

Axolotl usa arquivos de configuração YAML para controlar todo o fluxo de treinamento, incluindo:

• Seleção e parâmetros do modelo
• Configuração e pré-processamento do conjunto de dados
• Hiperparâmetros de treinamento
• Configurações do otimizador
• Monitoramento e registro de logs

Matriz de Compatibilidade de Modelos Suportados

Modelo	fp16/fp32	LoRA	QLoRA	GPTQ	Flash Attn	xformers
LLaMA	✅	✅	✅	✅	✅	✅
Mistral	✅	✅	✅	✅	✅	✅
Mixtral-MoE	✅	✅	✅	❓	❓	❓
Pythia	✅	✅	✅	❌	❌	❌
Falcon	✅	✅	✅	❌	❌	❌
Qwen	✅	✅	✅	❓	❓	❓
Gemma	✅	✅	✅	❓	❓	✅

✅: Suportado ❌: Não Suportado ❓: Não Testado

Cenários de Aplicação

Área de Pesquisa

• Pesquisa de ajuste fino de modelos de linguagem grandes
• Experimentação com métodos de treinamento eficientes em termos de parâmetros
• Pesquisa de alinhamento e segurança de modelos

Aplicações Industriais

• Customização de modelos de nível empresarial
• Treinamento de modelos específicos de domínio
• Otimização de modelos de funcionalidade de produtos

Educação e Treinamento

• Ensino de cursos de IA/ML
• Desenvolvimento de projetos práticos
• Aprimoramento de habilidades técnicas

Vantagens do Projeto

1. Facilidade de Uso: Controla fluxos de treinamento complexos por meio de arquivos de configuração YAML simples.
2. Flexibilidade: Suporta diversas arquiteturas de modelo e métodos de treinamento.
3. Desempenho: Integra as mais recentes tecnologias de otimização, proporcionando uma experiência de treinamento eficiente.
4. Escalabilidade: Suporta várias escalas de treinamento, de GPU única a multi-nós.
5. Código Aberto: Licença Apache 2.0, totalmente de código aberto e gratuito para uso.