Cómo ajustar modelos Baichuan con LoRA y Hugging Face PEFT en conjuntos de datos propietarios

En sectores verticales como finanzas, medicina, legal y gestión del conocimiento corporativo, el uso directo de APIs en la nube suele toparse con restricciones de soberanía de datos, normativas de privacidad y la necesidad de inyectar terminología interna especializada. Adaptar modelos fundacionales de pesos abiertos dentro de una infraestructura privada se ha convertido en el estándar de la ingeniería de PLN. La familia de modelos Baichuan (incluyendo Baichuan-7B/13B y Baichuan2-7B/13B en versiones Base y Chat) destaca en escenarios multilingües y bilingües gracias a su entrenamiento sobre 2,6 billones de tokens, un vocabulario ampliado de 125.696 tokens y diseños arquitectónicos optimizados para textos extensos y estabilidad numérica (NormHead).

Al emplear Low-Rank Adaptation (LoRA) a través de la biblioteca Parameter-Efficient Fine-Tuning (peft) de Hugging Face, los ingenieros pueden adaptar un modelo Baichuan actualizando menos del 1 % de sus parámetros. Esto permite ejecutar el entrenamiento en una sola GPU de 24 GB mediante QLoRA en 4 bits, o en una GPU A100/H100 de 80 GB en precisión nativa bfloat16.

1. Particularidades arquitectónicas de Baichuan

Antes de configurar el pipeline de entrenamiento, es fundamental tener en cuenta las decisiones de diseño que diferencian a Baichuan de una arquitectura LLaMA estándar:

1. Vocabulario BPE de 125.696 tokens: Posee una alta tasa de compresión sobre caracteres ideográficos y términos técnicos, reduciendo la longitud de secuencia entre un 15 % y un 20 % en comparación con tokenizadores occidentales estándar. Esto disminuye notablemente el consumo de memoria VRAM durante el procesamiento de lotes.
2. Esquema de codificación posicional:
   • Modelos 7B: Utilizan Rotary Position Embedding (RoPE) con una ventana de contexto de 4.096 tokens.
   • Modelos 13B: Implementan Attention with Linear Biases (ALiBi). En lugar de sumar embeddings posicionales a los vectores Query y Key, ALiBi aplica una penalización lineal en la matriz de atención según la distancia, lo que ofrece una extrapolación de contexto superior para contratos extensos o historiales clínicos.
3. Capa de salida normalizada (NormHead): Baichuan 2 introduce la normalización de pesos en la proyección lineal de salida (lm_head). Esta técnica mitiga la distorsión producida por la norma de tokens infrecuentes y orienta el cálculo de logits hacia la similitud de coseno. En un ajuste con LoRA, se recomienda mantener congelada la capa lm_head para no desestabilizar la distribución de salida preentrenada.
4. Código personalizado obligatorio: Baichuan requiere scripts de modelado específicos alojados en su repositorio, por lo que es obligatorio pasar trust_remote_code=True al cargar el modelo y el tokenizador en Hugging Face.

2. Preparación de datos propietarios y tokens especiales de control

Los conjuntos de datos empresariales suelen organizarse en registros de consultas, transcripciones de soporte técnico o documentación interna. Para el ajuste supervisado (SFT), Baichuan utiliza tokens de control específicos para delimitar los turnos conversacionales:

• <reserved_106> (Token ID: 195): Señala el inicio de la instrucción o consulta del usuario.
• <reserved_107> (Token ID: 196): Señala el inicio de la respuesta generada por el asistente.

Formato de datos (Esquema ShareGPT / Conversaciones)

Estructure sus datos en archivos JSON o JSON Lines (.jsonl) con la jerarquía de turnos:

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "¿Cuáles son los requisitos de aprobación para transferencias transfronterizas de fondos superiores a 50.000 dólares según nuestra normativa interna?"
      },
      {
        "from": "assistant",
        "value": "De acuerdo con el Artículo 4.2 del Manual de Cumplimiento Financiero, las transferencias que superen los 50.000 USD requieren la firma conjunta de dos directores y la validación automática del módulo AML en el sistema dentro de las 24 horas hábiles."
      }
    ]
  }
]

Enmascaramiento de pérdida en el prompt (Prompt Loss Masking)

Para evitar que el modelo gaste capacidad memorizando las preguntas, es una buena práctica asignar el valor -100 a las etiquetas (Labels) de todos los tokens del prompt (desde <reserved_106> hasta <reserved_107>). De este modo, la pérdida por entropía cruzada se calcula exclusivamente sobre los tokens de respuesta generados por el asistente y el token de finalización (</s> o <reserved_107>).

3. Configuración de LoRA y selección de hiperparámetros

LoRA descompone la matriz de actualización de pesos $\Delta W$ en el producto de dos matrices de bajo rango: $W = W_0 + \frac{\alpha}{r}(B \cdot A)$.

Selección de módulos objetivo (target_modules)

Aunque las configuraciones iniciales de LoRA solo adaptaban q_proj y v_proj, la práctica actual en adaptación de dominios complejos exige cubrir todas las capas lineales:

• Atención: q_proj, k_proj, v_proj, o_proj
• Capas MLP: gate_proj, up_proj, down_proj

En versiones modernas de peft, especificar target_modules="all-linear" detecta e instrumenta automáticamente todas las proyecciones lineales del decodificador Transformer de Baichuan.

Hiperparámetros recomendados

4. Script de implementación completo con Transformers y PEFT

El siguiente script implementa el flujo de trabajo de QLoRA en 4 bits listo para producción en una sola 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. Cargar el tokenizador
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 cuantización en 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. Cargar el 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 el modelo para entrenamiento k-bit
model = prepare_model_for_kbit_training(model)

# 5. Configurar el adaptador 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. Preprocesamiento de datos con enmascaramiento de pérdida
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 entrenamiento
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. Instanciar Trainer y ejecutar entrenamiento
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. Guardar pesos del adaptador LoRA
model.save_pretrained("./baichuan2-domain-lora-adapter")
tokenizer.save_pretrained("./baichuan2-domain-lora-adapter")

5. Fusión de pesos y despliegue en producción

Mantener los adaptadores LoRA desacoplados en producción genera sobrecarga de cómputo e impide que motores de inferencia de alto rendimiento como vLLM, Triton o LMDeploy ejecuten optimizaciones de fusión de kernels. La solución estándar es fusionar el adaptador en los pesos base.

Fusión sin conexión (Merge and Unload)

Advertencia: Nunca intente fusionar adaptadores directamente sobre un modelo cargado en 4 bits. Debe cargar el modelo base original sin cuantizar en torch.bfloat16 o torch.float16:

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",  # Usar "cuda" si hay suficiente VRAM libre
    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)

Una vez exportado, el modelo puede servirse inmediatamente en vLLM: vllm serve ./baichuan2-7b-domain-merged --trust-remote-code

6. Criterios de selección y directrices de ingeniería

1. Selección del punto de partida (Base frente a Chat):
   • Utilice Baichuan2-Base para tareas de extracción estructurada, generación de consultas SQL o clasificación formal, evitando sesgos conversacionales previos.
   • Utilice Baichuan2-Chat cuando desarrolle asistentes de atención al cliente o bots conversacionales que requieran modales de diálogo prealineados.
2. Elección de escala (7B frente a 13B):
   • Para el procesamiento de documentos que superen los 4.000 tokens, el modelo 13B con ALiBi ofrece una capacidad de extrapolación superior frente al 7B (RoPE).
   • El modelo 7B ofrece una menor latencia de inferencia y menor consumo de memoria, siendo la opción ideal para servicios de alto volumen transaccional.
3. Control del olvido catastrófico:
   • Para evitar la degradación de capacidades lógicas generales, se recomienda incorporar entre un 5 % y un 10 % de datos conversacionales generales en el corpus de ajuste fino.

Referencias

• 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