Como ajustar modelos Baichuan com LoRA e Hugging Face PEFT em conjuntos de dados proprietários

Em setores altamente regulamentados, como finanças, saúde, jurídico e suporte corporativo, o envio de informações sensíveis para APIs públicas de nuvem encontra barreiras legais e de segurança. A realização de ajuste fino (fine-tuning) em modelos de pesos abertos em infraestrutura proprietária tornou-se a abordagem padrão para engenheiros de PLN. A família de modelos Baichuan (incluindo Baichuan-7B/13B e Baichuan2-7B/13B nas variantes Base e Chat) conquistou ampla adoção devido ao seu pré-treinamento com 2,6 trilhões de tokens, vocabulário estendido de 125.696 tokens e inovações arquiteturais como o uso balanceado de RoPE/ALiBi e cabeçote com normalização de pesos (NormHead).

Utilizando Low-Rank Adaptation (LoRA) por meio da biblioteca Parameter-Efficient Fine-Tuning (peft) da Hugging Face, engenheiros podem especializar modelos Baichuan atualizando menos de 1% dos parâmetros. Isso viabiliza o treinamento em uma única GPU de 24 GB com QLoRA em 4 bits ou em instâncias A100/H100 de 80 GB com precisão bfloat16.

1. Particularidades arquiteturais do Baichuan

Antes de iniciar a configuração dos scripts de treinamento, é fundamental considerar as características singulares da arquitetura Baichuan:

1. Vocabulário BPE de 125.696 tokens: Possui excelente taxa de compressão para terminologias técnicas e caracteres asiáticos, gerando sequências de tokens 15% a 20% mais curtas em comparação com tokenizadores baseados em LLaMA. Isso reduz o consumo de memória VRAM durante o empacotamento de lotes.
2. Esquema de codificação posicional:
   • Modelos 7B: Utilizam Rotary Position Embeddings (RoPE) com janela de contexto padrão de 4.096 tokens.
   • Modelos 13B: Adotam Attention with Linear Biases (ALiBi). O ALiBi aplica penalidades lineares diretamente na matriz de pontuação de atenção em função da distância relativa, oferecendo extrapolação superior para textos longos, como contratos extensos e históricos médicos.
3. Cabeçote de projeção normalizado (NormHead): O Baichuan 2 normaliza os pesos da camada linear de saída (lm_head). Isso reduz a influência da magnitude de tokens raros e alinha o cálculo de logits à similaridade de cosseno. Durante o LoRA, recomenda-se não incluir o lm_head nos módulos treináveis, preservando a distribuição de probabilidades estabilizada no pré-treinamento.
4. Execução de código personalizado: Os modelos Baichuan possuem definições de rede próprias hospedadas no repositório. O parâmetro trust_remote_code=True é obrigatório no carregamento do modelo e do tokenizador.

2. Preparação de dados proprietários e tokens de controle

Dados empresariais geralmente estão organizados em registros de atendimento, bases de conhecimento ou pares de perguntas e respostas normativas. No ajuste supervisionado (SFT), o Baichuan utiliza delimitadores específicos para separar os turnos de diálogo:

• <reserved_106> (Token ID: 195): Indica o início da entrada do usuário (Human).
• <reserved_107> (Token ID: 196): Indica o início da resposta do assistente (Assistant).

Estrutura de dados (Esquema ShareGPT / Conversas)

Armazene os dados em formato JSON ou JSON Lines (.jsonl) estruturados por diálogo:

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "Quais são as exigências de conformidade para transferências internacionais de fundos acima de 50.000 dólares segundo a nossa política interna?"
      },
      {
        "from": "assistant",
        "value": "Conforme a Seção 4.2 do Manual de Compliance Financeiro, transferências acima de 50.000 USD exigem aprovação com dupla assinatura gerencial e liberação automática no módulo AML em até 24 horas úteis."
      }
    ]
  }
]

Mascaramento de perda no prompt (Prompt Loss Masking)

Para evitar que o modelo desperdice capacidade memorizando perguntas, deve-se mascarar os tokens correspondentes ao prompt. Todos os tokens entre <reserved_106> e <reserved_107> recebem o rótulo -100. A função de perda por entropia cruzada passa a ser calculada apenas sobre a resposta gerada pelo assistente e o delimitador final (</s> ou <reserved_107>).

3. Configuração do LoRA e hiperparâmetros recomendados

O LoRA decompõe a matriz de atualização de pesos $\Delta W$ no produto de duas matrizes de posto reduzido: $W = W_0 + \frac{\alpha}{r}(B \cdot A)$.

Seleção de módulos alvo (target_modules)

Em cenários corporativos complexos, aplicar LoRA apenas às projeções de consulta e valor (q_proj e v_proj) é insuficiente. A prática recomendada atual é adaptar todas as camadas lineares:

• Camadas de Atenção: q_proj, k_proj, v_proj, o_proj
• Camadas MLP: gate_proj, up_proj, down_proj

Nas versões recentes do peft, a flag target_modules="all-linear" identifica e configura automaticamente todos os módulos lineares do decodificador Transformer do Baichuan.

Tabela de hiperparâmetros

4. Script de implementação prática (Transformers e PEFT)

O script abaixo executa o pipeline completo de ajuste com QLoRA em 4 bits em uma única GPU.

import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq
)
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType
)

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

# 1. Carregar Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True,
    use_fast=False
)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# 2. Configurar quantização em 4 bits (QLoRA)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True
)

# 3. Carregar modelo base
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
)

# 4. Preparar para treinamento quantizado
model = prepare_model_for_kbit_training(model)

# 5. Configurar LoRA
peft_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ]
)

model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

# 6. Pré-processamento com mascaramento de prompt
USER_TOKEN_ID = 195  # <reserved_106>
ASSISTANT_TOKEN_ID = 196  # <reserved_107>

def preprocess_function(example):
    input_ids = []
    labels = []
    
    for message in example["conversations"]:
        if message["from"] == "human":
            prompt_tokens = [USER_TOKEN_ID] + tokenizer.encode(message["value"], add_special_tokens=False)
            input_ids.extend(prompt_tokens)
            labels.extend([-100] * len(prompt_tokens))
        elif message["from"] == "assistant":
            resp_tokens = [ASSISTANT_TOKEN_ID] + tokenizer.encode(message["value"], add_special_tokens=False) + [tokenizer.eos_token_id]
            input_ids.extend(resp_tokens)
            labels.extend(resp_tokens)
            
    return {"input_ids": input_ids[:2048], "labels": labels[:2048]}

dataset = load_dataset("json", data_files="domain_data.jsonl")
tokenized_dataset = dataset["train"].map(preprocess_function, remove_columns=["conversations"])

# 7. Argumentos de treinamento
training_args = TrainingArguments(
    output_dir="./baichuan2-lora-output",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.05,
    logging_steps=10,
    num_train_epochs=3,
    gradient_checkpointing=True,
    bf16=True,
    save_strategy="epoch",
    report_to="none"
)

# 8. Inicializar Trainer e treinar
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8, return_tensors="pt")
)

trainer.train()

# 9. Salvar adaptador LoRA
model.save_pretrained("./baichuan2-domain-lora-adapter")
tokenizer.save_pretrained("./baichuan2-domain-lora-adapter")

5. Fusão de pesos e disponibilização em produção

Manter adaptadores LoRA desacoplados em ambientes de alta demanda introduz latência adicional e impede otimizações de fusão de kernels em motores como vLLM, Triton e LMDeploy. O procedimento padrão de engenharia consiste em fundir o adaptador diretamente nos pesos base.

Script de fusão offline (Merge and Unload)

Importante: Não execute a fusão diretamente em instâncias quantizadas em 4 bits. Carregue o modelo base original em torch.bfloat16 ou torch.float16 na memória principal ou em GPU:

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

base_model_path = "baichuan-inc/Baichuan2-7B-Base"
adapter_path = "./baichuan2-domain-lora-adapter"
export_path = "./baichuan2-7b-domain-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",  # Utilizar "cuda" se houver VRAM disponível
    trust_remote_code=True
)

peft_model = PeftModel.from_pretrained(base_model, adapter_path)
merged_model = peft_model.merge_and_unload()

merged_model.save_pretrained(export_path, safe_serialization=True)
tokenizer.save_pretrained(export_path)

O diretório export_path gerado pode ser disponibilizado imediatamente no vLLM: vllm serve ./baichuan2-7b-domain-merged --trust-remote-code

6. Diretrizes e recomendações práticas

1. Escolha da base (Base vs. Chat):
   • Prefira Baichuan2-Base para tarefas que exijam saídas estruturadas (como JSON rígido, SQL ou extração de entidades), eliminando vícios discursivos prévios.
   • Adote Baichuan2-Chat quando o objetivo for implementar assistentes conversacionais corporativos ou bots de triagem que dependam de alinhamento prévio de diálogo.
2. Escolha do tamanho (7B vs. 13B):
   • Para processamento de documentos técnicos ou jurídicos extensos (acima de 4.000 tokens), o modelo 13B com ALiBi apresenta melhor estabilidade de atenção de longo alcance.
   • O modelo 7B oferece menor latência e economia de infraestrutura, sendo a escolha ideal para aplicações de alto volume e requisições curtas.
3. Prevenção de esquecimento catastrófico:
   • É recomendável misturar de 5% a 10% de dados de instrução geral em linguagem natural no conjunto de treinamento para preservar o raciocínio básico do modelo.

Referências

• GitHub - baichuan-inc/Baichuan2: A series of large language models developed by Baichuan Intelligent Technology
• GitHub - huggingface/peft: 🤗 PEFT: State-of-the-art Parameter-Efficient Fine-Tuning.
• [2309.10305] Baichuan 2: Open Large-scale Language Models
• [2106.09685] LoRA: Low-Rank Adaptation of Large Language Models