Ajustando um LLM de código aberto com Axolotl usando Otimização de Preferência Direta (DPO)

LLMs desbloquearam inúmeras novas oportunidades para aplicações de IA. Se você já quis ajustar seu próprio modelo, este guia mostrará como fazer isso facilmente e sem escrever nenhum código. Usando ferramentas como Axolotl e DPO, iremos passar pelo processo passo a passo.

O Que É um LLM?

Um Modelo de Linguagem Grande (LLM) é um modelo de IA poderoso treinado em vastas quantidades de dados de texto—dezenas de trilhões de caracteres—para prever o próximo conjunto de palavras em uma sequência. Isso só se tornou possível nos últimos 2-3 anos com os avanços feitos em computação GPU, que permitiram que esses modelos enormes fossem treinados em questão de algumas semanas.

Provavelmente você já interagiu com LLMs através de produtos como ChatGPT ou Claude antes e experimentou em primeira mão sua capacidade de entender e gerar respostas semelhantes às humanas.

Por Que Ajustar um LLM?

Não podemos simplesmente usar o GPT-4o para tudo? Bem, embora seja o modelo mais poderoso que temos no momento em que este artigo está sendo escrito, nem sempre é a escolha mais prática. Ajustar um modelo menor, variando de 3 a 14 bilhões de parâmetros, pode produzir resultados comparáveis a uma pequena fração do custo. Além disso, o ajuste fino permite que você possua sua propriedade intelectual e reduz sua dependência de terceiros.

Entendendo os Modelos Base, Instrução e Chat

Antes de mergulhar no ajuste fino, é essencial entender os diferentes tipos de LLMs que existem:

Modelos Base: Estes são pré-treinados em grandes quantidades de texto não estruturado, como livros ou dados da internet. Embora tenham uma compreensão intrínseca da linguagem, não são otimizados para inferência e produzirão saídas incoerentes. Os modelos base são desenvolvidos para servir como ponto de partida para o desenvolvimento de modelos mais especializados.
Modelos Instrutivos: Construídos em cima dos modelos base, os modelos instrutivos são ajustados usando dados estruturados, como pares de prompt e resposta. Eles são projetados para seguir instruções específicas ou responder perguntas.
Modelos de Conversação: Também construídos em cima dos modelos base, mas ao contrário dos modelos instrutivos, os modelos de conversação são treinados em dados de conversação, permitindo-lhes engajar em diálogos de ida e volta.

O Que É Aprendizado por Reforço e DPO?

O Aprendizado por Reforço (RL) é uma técnica onde os modelos aprendem recebendo feedback sobre suas ações. É aplicado a modelos instrutivos ou de conversação para refinar ainda mais a qualidade de suas saídas. Tipicamente, o RL não é realizado em cima de modelos base, pois utiliza uma taxa de aprendizado muito menor que não moverá a agulha o suficiente.

O DPO é uma forma de RL onde o modelo é treinado usando pares de respostas boas e ruins para o mesmo prompt/conversa. Ao apresentar esses pares, o modelo aprende a favorecer os bons exemplos e evitar os ruins.

Quando Usar DPO

O DPO é particularmente útil quando você deseja ajustar o estilo ou comportamento do seu modelo, por exemplo:

Ajustes de Estilo: Modificar o comprimento das respostas, o nível de detalhe ou o grau de confiança expresso pelo modelo.
Medidas de Segurança: Treine o modelo para recusar responder a solicitações potencialmente inseguras ou inadequadas.

No entanto, o DPO não é adequado para ensinar ao modelo novos conhecimentos ou fatos. Para esse propósito, técnicas de Aprendizado Supervisionado (SFT) ou Geração Aumentada por Recuperação (RAG) são mais apropriadas.

Criando um Conjunto de Dados DPO

Em um ambiente de produção, você normalmente geraria um conjunto de dados DPO usando feedback de seus usuários, por exemplo:

Feedback do Usuário: Implementando um mecanismo de avaliação positiva/negativa nas respostas.
Escolhas Comparativas: Apresentando aos usuários duas saídas diferentes e pedindo que escolham a melhor.

Se você não tiver dados de usuários, também pode criar um conjunto de dados sintético aproveitando LLMs maiores e mais capazes. Por exemplo, você pode gerar respostas ruins usando um modelo menor e, em seguida, usar o GPT-4 para corrigi-las.

Para simplificar, usaremos um conjunto de dados pronto da HuggingFace: olivermolenschot/alpaca_messages_dpo_test. Se você inspecionar o conjunto de dados, notará que ele contém solicitações com respostas escolhidas e rejeitadas—estes são os bons e maus exemplos. Esses dados foram criados sinteticamente usando GPT-3.5-turbo e GPT-4.

Geralmente, você precisará de entre 500 e 1.000 pares de dados, no mínimo, para ter um treinamento eficaz sem overfitting. Os maiores conjuntos de dados DPO contêm até 15.000–20.000 pares.

Ajustando o Qwen2.5 3B Instruct com Axolotl

Vamos usar o Axolotl para ajustar o modelo Qwen2.5 3B Instruct, que atualmente ocupa a primeira posição no OpenLLM Leaderboard para sua classe de tamanho. Com o Axolotl, você pode ajustar um modelo sem escrever uma única linha de código – apenas um arquivo de configuração YAML. Abaixo está o config.yml que usaremos:

base_model: Qwen/Qwen2.5-3B-Instruct
strict: false

# Axolotl will automatically map the dataset from HuggingFace to the prompt template of Qwen 2.5
chat_template: qwen_25
rl: dpo
datasets:
  - path: olivermolenschot/alpaca_messages_dpo_test
    type: chat_template.default
    field_messages: conversation
    field_chosen: chosen
    field_rejected: rejected
    message_field_role: role
    message_field_content: content

# We pick a directory inside /workspace since that's typically where cloud hosts mount the volume
output_dir: /workspace/dpo-output

# Qwen 2.5 supports up to 32,768 tokens with a max generation of 8,192 tokens
sequence_len: 8192

# Sample packing does not currently work with DPO. Pad to sequence length is added to avoid a Torch bug
sample_packing: false
pad_to_sequence_len: true

# Add your WanDB account if you want to get nice reporting on your training performance
wandb_project:
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:

# Can make training more efficient by batching multiple rows together
gradient_accumulation_steps: 1
micro_batch_size: 1

# Do one pass on the dataset. Can set to a higher number like 2 or 3 to do multiple
num_epochs: 1

# Optimizers don't make much of a difference when training LLMs. Adam is the standard
optimizer: adamw_torch

# DPO requires a smaller learning rate than regular SFT
lr_scheduler: constant
learning_rate: 0.00005

# Train in bf16 precision since the base model is also bf16
bf16: auto

# Reduces memory requirements
gradient_checkpointing: true

# Makes training faster (only suported on Ampere, Ada, or Hopper GPUs)
flash_attention: true

# Can save multiple times per epoch to get multiple checkpoint candidates to compare
saves_per_epoch: 1

logging_steps: 1
warmup_steps: 0

Configurando o Ambiente na Nuvem

Para executar o treinamento, usaremos um serviço de hospedagem na nuvem como Runpod ou Vultr. Aqui está o que você vai precisar:

Imagem Docker: Clone a imagem Docker winglian/axolotl-cloud:main fornecida pela equipe do Axolotl.
*Requisitos de Hardware: Uma GPU com 80GB de VRAM (como um nó 1×A100 PCIe) será mais do que suficiente para este tamanho de modelo.
Armazenamento: 200GB de armazenamento em volume para acomodar todos os arquivos de que precisamos.
Versão do CUDA: Sua versão do CUDA deve ser no mínimo 12.1.

*Esse tipo de treinamento é considerado um ajuste completo do LLM, e é muito intensivo em VRAM. Se você deseja executar um treinamento localmente, sem depender de hosts na nuvem, poderia tentar usar o QLoRA, que é uma forma de ajuste supervisionado. Embora seja teoricamente possível combinar DPO & QLoRA, isso é muito raramente feito.

Passos para Iniciar o Treinamento

Definir o Diretório de Cache da HuggingFace:

export HF_HOME=/workspace/hf

Isso garante que o modelo original seja baixado para o nosso armazenamento em volume, que é persistente.

Criar Arquivo de Configuração: Salvar o arquivo config.yml que criamos anteriormente em /workspace/config.yml.

Iniciar o Treinamento:

python -m axolotl.cli.train /workspace/config.yml

E eis! Seu treinamento deve começar. Depois que Axolotl baixar o modelo e os dados de treinamento, você deverá ver uma saída semelhante a esta:

[2024-12-02 11:22:34,798] [DEBUG] [axolotl.train.train:98] [PID:3813] [RANK:0] loading model

[2024-12-02 11:23:17,925] [INFO] [axolotl.train.train:178] [PID:3813] [RANK:0] Starting trainer...

O treinamento deve levar apenas alguns minutos para ser concluído, pois este é um pequeno conjunto de dados com apenas 264 linhas. O modelo ajustado será salvo em /workspace/dpo-output.

Enviando o Modelo para o HuggingFace

Você pode enviar seu modelo para o HuggingFace usando a CLI:

Instale a CLI do HuggingFace Hub:

pip install huggingface_hub[cli]

Envie o Modelo:

huggingface-cli upload /workspace/dpo-output yourname/yourrepo

Substitua yourname/yourrepo pelo seu nome de usuário real do HuggingFace e nome do repositório.

Avaliando Seu Modelo Ajustado

Para avaliação, é recomendado hospedar tanto os modelos originais quanto os ajustados usando uma ferramenta como Text Generation Inference (TGI). Em seguida, realize inferências em ambos os modelos com uma configuração de temperatura de 0 (para garantir saídas determinísticas) e compare manualmente as respostas dos dois modelos.

Essa abordagem prática fornece insights melhores do que depender exclusivamente de métricas de perda de avaliação do treinamento, que podem não capturar as nuances da geração de linguagem em LLMs.

Conclusão

O ajuste fino de um LLM usando DPO permite personalizar modelos para atender melhor às necessidades da sua aplicação, tudo isso mantendo os custos gerenciáveis. Ao seguir os passos descritos neste artigo, você pode aproveitar o poder de ferramentas e conjuntos de dados de código aberto para criar um modelo que se alinhe com seus requisitos específicos. Se você está procurando ajustar o estilo das respostas ou implementar medidas de segurança, o DPO oferece uma abordagem prática para refinar seu LLM.

Bom ajuste fino!