Guía práctica para afinar LFM2 en Google Colab con QLoRA, LoRA y DPO

Un tutorial práctico muestra cómo afinar el modelo LFM2 en Google Colab siguiendo un flujo abierto que combina QLoRA, SFT (entrenamiento supervisado ligero), LoRA y DPO (Direct Preference Optimization). El material incluye fragmentos de código ejecutable que cubren el proceso completo: instalación de dependencias, verificación de GPU, carga del checkpoint base, ajuste con LoRA y QLoRA, fusión del adaptador en el modelo y una etapa opcional de alineación por preferencias. El objetivo declarado es ofrecer una ruta eficiente en memoria para ajustar modelos medianos en entornos con GPU limitadas sin rehacer el entrenamiento desde cero.

La guía comienza con la instalación y verificación del entorno: paquetes clave como transformers, trl, peft, datasets, accelerate y bitsandbytes; y una comprobación explícita de GPU en Colab via torch.cuda.is_available(). También se definen parámetros principales al inicio del script: MODEL_ID = "LiquidAI/LFM2-1.2B", USE_4BIT = True, RUN_DPO = True, SFT_SAMPLES = 500, SFT_STEPS = 60, DPO_STEPS = 40 y MAX_LEN = 1024, de modo que el lector pueda reproducir exactamente las condiciones del experimento.

Para minimizar el uso de memoria el tutorial explica la carga del modelo con cuantización a 4 bits mediante BitsAndBytesConfig: load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True y bnb_4bit_compute_dtype asignado al DTYPE elegido (bf16 si el hardware lo soporta, en caso contrario fp16). Se incluye la corrección del tokenizer cuando pad_token es None y se propone una función de chat para probar respuestas antes y después del ajuste, facilitando comparaciones rápidas del comportamiento del modelo en cada etapa.

En la fase de SFT el autor utiliza un subconjunto de datos conversacionales: "HuggingFaceTB/smoltalk" con el split train[:500], seleccionando únicamente la columna messages para crear ejemplos de entrenamiento. Para mantener el ajuste ligero y eficiente se configura LoRA con parámetros concretos: r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM" y target_modules="all-linear". Esta configuración añade adaptadores compactos al modelo base en lugar de actualizar todos los pesos, reduciendo costos de memoria y almacenamiento del checkpoint.

La configuración de entrenamiento SFT se define mediante un SFTConfig con valores reproducibles: output_dir="outputs/sft/lfm2_demo", max_length=1024, per_device_train_batch_size=2, gradient_accumulation_steps=4, learning_rate=2e-5, warmup_ratio=0.03, lr_scheduler_type="cosine" y max_steps=60. El flujo recomienda activar gradient_checkpointing para ahorrar memoria, mantener un registro mínimo y seleccionar el optimizador en función de la cuantización (por ejemplo paged_adamw_8bit si USE_4BIT es True). El entrenamiento se ejecuta con SFTTrainer de TRL y el adaptador resultante se guarda en outputs/sft/lfm2_adapter.

Tras completar el SFT, el tutorial libera memoria y recarga el checkpoint base en fp16 o bf16 según soporte hardware, aplicando PeftModel.from_pretrained("outputs/sft/lfm2_adapter").merge_and_unload() para fusionar el adaptador LoRA con los pesos del modelo. El modelo fusionado y el tokenizer se guardan en outputs/sft/lfm2_merged, generando un checkpoint independiente que incorpora los cambios del adaptador y que puede distribuirse o probarse sin depender del pipeline de LoRA en tiempo de inferencia.

Si RUN_DPO está activado, la guía extiende el flujo con una etapa de ajuste por preferencias (DPO). Describe cómo preparar ejemplos de preferencia compuestos por un prompt y dos respuestas etiquetadas como "chosen" y "rejected" (elegida y rechazada), y cómo ejecutar DPO durante DPO_STEPS = 40 para alinear el generador con preferencias observadas en estas parejas. El tutorial incluye las llamadas y parámetros necesarios para integrar DPO tras la fusión del adaptador, de modo que el modelo final refleje señales humanas de preferencia sin reentrenar desde cero.

La combinación QLoRA + LoRA + merge + DPO que presenta el tutorial ofrece una ruta práctica para afinar modelos medianos en entornos con recursos limitados. QLoRA reduce el uso de memoria durante la carga y entrenamiento, LoRA aporta adaptadores de bajo coste que se pueden fusionar posteriormente y DPO permite introducir señales de preferencia humana. El autor documenta las rutas de salida (outputs/sft/...) y los comandos necesarios para reproducir el pipeline, facilitando la replicabilidad del experimento en Google Colab.

El tutorial también señala limitaciones importantes: el script depende de la presencia de GPU en Colab (fallará si torch.cuda.is_available() devuelve False), el uso de bf16 depende del soporte del hardware y la cuantización a 4 bits requiere configuraciones específicas de bitsandbytes y optimizadores compatibles. Además, aunque el artículo entrega código y parámetros reproducibles, advierte que este flujo no sustituye pruebas adicionales de seguridad, evaluación de sesgos ni validación en entornos de producción; dichas verificaciones deben realizarse antes de desplegar modelos ajustados.

En conjunto, la guía proporciona un recorrido detallado y reproducible para investigadores y practicantes que buscan ajustar LFM2 en Colab con restricciones de GPU: desde la instalación y los parámetros exactos hasta la fusión del adaptador y la alineación por preferencias, manteniendo un enfoque práctico y orientado a la eficiencia de memoria.