Ajuste fino de modelos Baichuan con LoRA: Guía práctica para asistentes en dominios verticales

El desarrollo de asistentes de inteligencia artificial para sectores altamente especializados —como el triaje médico, el análisis de contratos legales, el cumplimiento financiero o el soporte técnico empresarial— exige una comprensión rigurosa de la terminología técnica, razonamiento estructurado y formatos de salida predecibles. Si bien los modelos comerciales de propósito general ofrecen fluidez conversacional, los imperativos de privacidad de datos y los costes de infraestructura impulsan a las organizaciones hacia soluciones basadas en pesos abiertos y autoalojadas. En este ámbito, la familia de modelos Baichuan (especialmente Baichuan2-7B y Baichuan2-13B) destaca como una de las bases bilingües más sólidas, respaldada por un entrenamiento previo sobre 2,6 billones de tokens y un vocabulario optimizado.

No obstante, realizar un ajuste fino de todos los parámetros (Full Fine-Tuning) sobre modelos de 7B o 13B exige costosos clústeres de GPUs y tiende a provocar el temido "olvido catastrófico" de las capacidades conversacionales generales. La técnica de Adaptación de Bajo Rango (LoRA, Low-Rank Adaptation) resuelve este problema congelando los pesos del modelo base e introduciendo matrices de descomposición de bajo rango en las capas de proyección críticas. De este modo, los ingenieros de aprendizaje automático pueden adaptar Baichuan a sectores regulados empleando hardware accesible con una sobrecarga computacional mínima.

1. Alineación arquitectónica: el mecanismo W_pack en la atención de Baichuan

Para configurar adecuadamente los adaptadores LoRA, el equipo de ingeniería debe comprender las particularidades del diseño Transformer de Baichuan e intervenir en las capas lineales correctas.

A diferencia de arquitecturas tipo Llama convencionales, que descomponen la atención propia (Self-Attention) en proyecciones individuales q_proj, k_proj y v_proj, Baichuan agrupa las transformaciones lineales de Query, Key y Value en una única proyección empaquetada: W_pack.

Si se aplican configuraciones genéricas de LoRA sin contemplar esta estructura, el mecanismo de atención no recibirá ninguna actualización de gradiente. Con el objetivo de adaptar tanto la atención como las transformaciones del perceptrón multicapa (MLP), los módulos objetivo (target_modules) de LoRA deben incluir:

• Capas de atención: W_pack (fusión de Q, K y V) y o_proj (proyección de salida).
• Capas MLP: gate_proj, up_proj y down_proj.

Especificar explícitamente estas capas o utilizar la directiva all-linear en PEFT garantiza la capacidad expresiva requerida para asimilar el conocimiento del dominio.

2. Preparación de conjuntos de datos y tokens especiales

La precisión de un asistente vertical depende de la calidad de sus datos de entrenamiento y de la disciplina en la estructura de los diálogos.

Estructura de tokens de rol

Las variantes de diálogo de Baichuan2 (Chat) utilizan identificadores reservados para separar las intervenciones de la conversación:

• <reserved_106>: Señaliza el inicio de la intervención del usuario (<|User|>).
• <reserved_107>: Señaliza el inicio de la generación del asistente (<|Assistant|>).

El formato estándar en JSON sigue la estructura habitual de conversaciones por turnos:

[
  {
    "conversations": [
      {
        "from": "human",
        "value": "¿Cuáles son los plazos estatutarios y los requisitos probatorios para solicitar la nulidad de una patente?"
      },
      {
        "from": "gpt",
        "value": "Conforme a la normativa de propiedad industrial, tras la concesión oficial de la patente, cualquier interesado legítimo puede incoar el procedimiento de nulidad acreditando..."
      }
    ]
  }
]

Regla crítica: Enmascaramiento de pérdida (Loss Masking)

Uno de los fallos más perjudiciales al programar scripts de entrenamiento consiste en computar la pérdida de entropía cruzada sobre toda la secuencia. En el ajuste fino supervisado (SFT), la pérdida debe calcularse única y exclusivamente sobre los tokens generados por el asistente.

Es imprescindible enmascarar los mensajes del usuario, las directrices del sistema y los tokens delimitadores asignándoles el valor -100 en el vector de etiquetas (labels), correspondiente al índice de exclusión estándar en PyTorch. Si no se aplica este enmascaramiento, el modelo consumirá capacidad de gradiente memorizando cómo formular preguntas en lugar de perfeccionar la precisión técnica de sus respuestas.

Directrices de ingeniería de datos

1. Casos y terminología especializada: Recopilar entre 5.000 y 20.000 pares de instrucción-respuesta centrados en diagnósticos, jurisprudencia o análisis financiero.
2. Datos de anclaje general (Prevención del olvido): Incorporar entre un 5% y un 10% de datos conversacionales generales. Esto evita la pérdida de fluidez, sentido común y cortesía conversacional básica.
3. Muestras de rechazo y límites de seguridad: Incluir ejemplos en los que el asistente decline formular diagnósticos concluyentes o garantías de rentabilidad, orientando al usuario hacia profesionales acreditados.

3. Implementación práctica de LoRA con Hugging Face PEFT

A continuación se detalla una receta reproducible para entrenar Baichuan2-7B-Base o Baichuan2-7B-Chat empleando transformers, peft y PyTorch.

Paso 1: Instalación de dependencias

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

Paso 2: Configuración del modelo y 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"

# Carga del tokenizador con trust_remote_code obligatorio
tokenizer = AutoTokenizer.from_pretrained(
    model_id, 
    use_fast=False, 
    trust_remote_code=True
)

# Carga del modelo base en bfloat16 para máxima estabilidad numérica
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True
)

# Configuración de hiperparámetros LoRA
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,                         # Dimensión del rango intrínseco
    lora_alpha=32,                # Factor de escala (típicamente 2x el rango)
    lora_dropout=0.05,            # Mitiga el sobreajuste en datos limitados
    target_modules=[
        "W_pack",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj"
    ],
    bias="none"
)

# Creación del modelo PEFT
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Salida típica: trainable params: ~20M || all params: 7B || trainable%: ~0.28%

Paso 3: Ejecución del entrenamiento y ahorro de memoria

El uso de gradient_checkpointing reduce drásticamente la ocupación de memoria de vídeo (VRAM), facilitando tamaños de lote operativos:

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 en arquitecturas modernas
    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")

En entornos con restricciones de hardware (por ejemplo, una única GPU de 24GB como la RTX 3090/4090), se puede activar QLoRA añadiendo BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16) en la llamada a from_pretrained.

4. Fusión de pesos para despliegue en producción sin sobrecarga

Cargar adaptadores LoRA de forma dinámica en entornos de inferencia intensiva añade latencia de ejecución y complejidad operativa. Una vez validada la precisión técnica del modelo, se aconseja fusionar los pesos adaptados directamente en la 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
)

# Carga del adaptador y fusión permanente
model = PeftModel.from_pretrained(base_model, adapter_path)
merged_model = model.merge_and_unload()

# Exportación del modelo independiente
merged_model.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)

El artefacto resultante es idéntico a un modelo base estándar y puede desplegarse en motores como vLLM o TGI, garantizando un servicio de alto rendimiento sin penalizaciones de latencia.

5. Criterios de diseño: LoRA vs. Generación Aumentada por Recuperación (RAG)

Al estructurar asistentes corporativos, a menudo surge el dilema entre aplicar LoRA o implementar RAG. En la práctica, ambas tecnologías son complementarias:

Recomendación en producción: La estrategia óptima radica en una arquitectura híbrida: utilizar LoRA para instruir a Baichuan en los patrones de razonamiento y el tono formal del sector, complementándolo con un mecanismo RAG que suministre el contexto documental exacto en tiempo real.

Referencias

• 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