Ajuste Fino de Modelos Baichuan com LoRA: Guia Prático para Assistentes em Domínios Verticais

Construir assistentes de inteligência artificial para setores altamente especializados — como triagem médica, análise contratual jurídica, conformidade regulatória financeira e suporte técnico corporativo — exige compreensão terminológica profunda, raciocínio metódico e aderência estrita a padrões formais de resposta. Embora modelos proprietários generalistas exibam notável fluência, exigências de privacidade de dados e restrições orçamentárias impulsionam as organizações para soluções abertas e auto-hospedadas. A linha de modelos Baichuan (notadamente Baichuan2-7B e Baichuan2-13B) consolidou-se como uma das fundações bilíngues mais robustas, treinada a partir do zero em 2,6 trilhões de tokens com excelente vocabulário e eficiência de processamento.

Entretanto, realizar o ajuste fino integral de parâmetros (Full Fine-Tuning) em arquiteturas de 7B ou 13B exige infraestrutura computacional pesada e comumente deflagra o "esquecimento catastrófico", degradando as habilidades de conversação geral do modelo. A técnica de Adaptação de Baixo Posto (LoRA, Low-Rank Adaptation) oferece o equilíbrio técnico ideal: ao congelar os pesos da base e injetar matrizes treináveis de decomposição linear nas camadas-chave, os engenheiros de machine learning conseguem especializar o modelo em hardware acessível com mínimo custo computacional.

1. Alinhamento Arquitetural: O Mecanismo W_pack na Atenção do Baichuan

Para configurar os adaptadores LoRA com eficácia, a equipe de engenharia precisa compreender a estrutura Transformer interna do Baichuan e selecionar as camadas lineares corretas.

Ao contrário de arquiteturas tradicionais baseadas no Llama, que separam as projeções de Query (Q), Key (K) e Value (V) em matrizes lineares independentes (q_proj, k_proj e v_proj), o Baichuan compacta essas operações em uma única camada de projeção combinada: W_pack.

Se uma configuração padrão de LoRA for aplicada sem considerar essa arquitetura, o mecanismo de autoatenção não receberá nenhum ajuste de gradiente. Para adaptar adequadamente tanto a atenção quanto as transformações não lineares do perceptron multicamadas (MLP), os módulos-alvo (target_modules) devem contemplar:

• Camadas de atenção: W_pack (fusão de Q, K e V) e o_proj (projeção de saída).
• Camadas MLP: gate_proj, up_proj e down_proj.

A indicação explícita dessas projeções no PEFT — ou o uso do comando target_modules="all-linear" — assegura a capacidade representacional requerida para assimilar os fluxos do domínio específico.

2. Preparação de Dados e Padronização de Tokens Especiais

A consistência técnica do assistente vertical decorre diretamente do rigor aplicado na curadoria do conjunto de dados e na obediência ao formato de diálogo.

Estrutura de Tokens de Diálogo

As versões de conversação do Baichuan2 (Chat) empregam tokens reservados para demarcar os turnos interativos:

• <reserved_106>: Marca o início da fala do usuário (<|User|>).
• <reserved_107>: Dispara o início da geração do assistente (<|Assistant|>).

O formato padrão em JSON estrutura-se em turnos conversacionais:

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "Quais são os prazos regimentais e os requisitos comprobatórios para a apresentação de um pedido de nulidade de patente?"
      },
      {
        "from": "gpt",
        "value": "Nos termos da legislação de propriedade industrial, após a concessão oficial da patente, qualquer interessado legítimo pode instaurar o processo de nulidade mediante..."
      }
    ]
  }
]

Regra Crítica: Mascaramento de Perda (Loss Masking)

Um erro crítico comum no desenvolvimento de scripts personalizados é calcular a perda de entropia cruzada em toda a sequência de tokens. No ajuste supervisionado (SFT), o gradiente de perda deve ser computado estritamente sobre as respostas geradas pelo assistente.

Os engenheiros devem preencher as mensagens do usuário, instruções de sistema e tokens delimitadores com o valor -100 no vetor de rótulos (labels), correspondente ao índice padrão de exclusão no PyTorch. Sem esse mascaramento, o modelo despende capacidade de aprendizado memorizando como formular perguntas, resultando em respostas redundantes e instabilidade sintática.

Critérios de Curadoria de Dados

1. Casos reais e terminologia técnica: Reúna de 5.000 a 20.000 pares de instrução e resposta abordando diagnósticos clínicos, pareceres jurídicos ou demonstrações financeiras.
2. Dados de ancoragem geral (Prevenção de esquecimento): Misture de 5% a 10% de dados de instrução conversacional geral. Isso preserva a fluidez linguística, a coerência dialógica e o raciocínio comum.
3. Amostras de recusa e conformidade ética: Inclua interações em que o assistente decline assertivas médicas taxativas ou garantias financeiras sem base factual, orientando o usuário a buscar profissionais habilitados.

3. Implementação Prática com Hugging Face PEFT

Abaixo apresentamos o fluxo completo de ajuste fino para Baichuan2-7B-Base ou Baichuan2-7B-Chat, utilizando transformers, peft e PyTorch.

Passo 1: Instalação das dependências

pip install torch>=2.0.0 transformers>=4.31.0 peft>=0.5.0 accelerate bitsandbytes datasets

Passo 2: Inicialização do modelo e parâmetros LoRA

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
from peft import LoraConfig, get_peft_model, TaskType

model_id = "baichuan-inc/Baichuan2-7B-Base"

# Carregamento do tokenizador com trust_remote_code
tokenizer = AutoTokenizer.from_pretrained(
    model_id, 
    use_fast=False, 
    trust_remote_code=True
)

# Carregamento do modelo em bfloat16 para alta fidelidade numérica
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)

# Definição dos hiperparâmetros LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,                         # Dimensão do posto intrínseco
    lora_alpha=32,                # Fator de escala (usualmente 2x rank)
    lora_dropout=0.05,            # Evita sobreajuste aos dados do nicho
    target_modules=[
        "W_pack",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj"
    ],
    bias="none"
)

# Aplicação do adaptador PEFT
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Exibição típica: trainable params: ~20M || all params: 7B || trainable%: ~0.28%

Passo 3: Treinamento com otimização de memória

A ativação de gradient_checkpointing reduz substancialmente o uso de memória de vídeo (VRAM), viabilizando lotes adequados:

training_args = TrainingArguments(
    output_dir="./baichuan2-lora-vertical",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=10,
    num_train_epochs=3,
    save_strategy="epoch",
    fp16=False,
    bf16=True,                    # Recomendado para arquiteturas Ampere em diante
    gradient_checkpointing=True,
    report_to="none"
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_domain_dataset,
    tokenizer=tokenizer
)

trainer.train()
model.save_pretrained("./baichuan2-vertical-adapter")

Para restrições severas de hardware (por exemplo, placas de 24GB como RTX 3090/4090), adicione BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16) para executar via QLoRA.

4. Fusão de Pesos para Produção com Latência Zero

Em ambientes corporativos de alto rendimento, carregar adaptadores LoRA dinamicamente acarreta latência adicional e maior complexidade de orquestração. Após a consolidação das métricas de teste, recomenda-se fundir os pesos LoRA diretamente ao modelo base.

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

base_model_path = "baichuan-inc/Baichuan2-7B-Base"
adapter_path = "./baichuan2-vertical-adapter"
output_path = "./baichuan2-7b-vertical-merged"

tokenizer = AutoTokenizer.from_pretrained(base_model_path, trust_remote_code=True)
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.bfloat16,
    device_map="cpu",
    trust_remote_code=True
)

# Fusão dos pesos do adaptador
model = PeftModel.from_pretrained(base_model, adapter_path)
merged_model = model.merge_and_unload()

# Salvamento do modelo autônomo resultante
merged_model.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)

O modelo resultante opera como um checkpoint padrão e pode ser imediatamente carregado em motores de inferência de alto rendimento como vLLM ou TGI, eliminando qualquer sobrecarga em produção.

5. Tomada de Decisão Arquitetural: LoRA vs. RAG

Na engenharia de assistentes verticais, discute-se com frequência a escolha entre LoRA e Recuperação Aumentada por Geração (RAG). Na prática, ambas as técnicas exercem papéis complementares:

Recomendação para Produção: Adote uma arquitetura híbrida: configure o modelo Baichuan via LoRA para dominar o linguajar formal e a cadeia de raciocínio da área, e acople um pipeline de RAG para embasar as respostas em documentos verídicos e atualizados.

Referências

• GitHub - baichuan-inc/Baichuan2: A series of large language models developed by Baichuan Intelligent Technology · GitHub
• [2309.10305] Baichuan 2: Open Large-scale Language Models
• baichuan-inc/Baichuan2-7B-Base · Hugging Face
• baichuan-inc/Baichuan2-13B-Chat · Hugging Face
• GitHub - huggingface/peft: 🤗 PEFT: State-of-the-art Parameter-Efficient Fine-Tuning. · GitHub