Entrenando un LLM de 3.8B por $998: Optimización de Throughput y Convergencia en Hardware de Consumo

El entrenamiento de modelos de lenguaje a gran escala (LLMs) tradicionalmente ha sido dominio de grandes organizaciones con presupuestos de cómputo sustanciales. Sin embargo, este trabajo postula que, mediante una ingeniería de sistemas rigurosa y la aplicación de principios de optimización de rendimiento, es posible entrenar LLMs con un número significativo de parámetros (en este caso, 3.8B) a un costo accesible para un individuo ($998). El problema fundamental que aborda es cómo maximizar la eficiencia del cómputo por unidad monetaria, trasladando las lecciones de la optimización de sistemas distribuidos a la carga de trabajo intensiva de entrenamiento de modelos de aprendizaje profundo.

La relevancia actual de esta tesis radica en la democratización del acceso a la capacidad de entrenar modelos de IA, lo que permite la experimentación y el desarrollo fuera de los ecosistemas de grandes empresas. Históricamente, la ley de Moore y las innovaciones en hardware han reducido el costo del cómputo, pero la optimización del software y la arquitectura de entrenamiento son igualmente críticas para aprovechar estas ganancias. Este proyecto es un ejemplo de cómo la aplicación de principios de ingeniería de software a la infraestructura de ML puede desbloquear capacidades significativas con recursos limitados.

Arquitectura del Sistema

El sistema de entrenamiento, denominado 'little-lm', se concibe como un framework configurable para LLMs decoder-only. Su arquitectura se basa en una separación de preocupaciones clara, donde la configuración del modelo, el dataset, el optimizador, el scheduler y los callbacks se especifican a través de un archivo YAML. Los componentes se auto-registran en un registro global, permitiendo su resolución por nombre y facilitando la experimentación mediante cambios mínimos en la configuración.

El modelo en sí sigue una arquitectura tipo Llama, incorporando componentes como RMSNorm para la normalización, RoPE (Rotary Position Embeddings) para la codificación posicional, Grouped-Query Attention (GQA) con una relación de 3:1 (24 query heads, 8 KV heads) para eficiencia de memoria y cómputo, MLPs con activación relu² (no-gated), QK-norm, logit softcap, y escalares residuales aprendibles por capa. Una característica notable es el uso de 'value embeddings' (14 tablas, alternando capas), que representan el 19% del conteo total de parámetros y contribuyen a una mejora en la calidad del modelo con un costo de throughput casi nulo. La estrategia de paralelización distribuida utilizada es DistributedDataParallel (DDP) estándar de PyTorch, considerada suficiente para el tamaño del modelo en un solo nodo, evitando la complejidad de optimizadores shardeados como ZeRO-2.

Las decisiones clave de diseño incluyen el uso de un scheduler de tasa de aprendizaje trapezoidal (warmup, hold flat, linear cooldown) para mantener la convergencia hasta el final del entrenamiento, y un optimizador compuesto: Muon para los parámetros de matriz (debido a su mejor eficiencia por token en esta escala) y AdamW para el resto de los parámetros (embeddings, lm_head, value_embeds, scalars). Se implementó entrenamiento en FP8 utilizando torch._scaled_mm con escalado dinámico tensorwise en todas las operaciones GEMM, y se realizó padding del vocabulario a un múltiplo de 64 para optimizar el uso de los Tensor Cores. La función de pérdida LigerFusedLinearCrossEntropyLoss se utilizó para fusionar la multiplicación de la capa lm_head con el cálculo de la pérdida, reduciendo el uso de VRAM y permitiendo tamaños de batch mayores. Los datasets se descargan localmente para evitar latencias y caídas de red asociadas con el streaming.

model:
  hidden_size: 3072
  intermediate_size: 12288
  num_hidden_layers: 28
  num_attention_heads: 24
  num_key_value_heads: 8
  vocab_pad_to: 64
  max_position_embeddings: 2048
  dtype: bf16
optimizer:
  matrix: Muon lr=0.02 momentum=0.95 wd=0.0
  embeddings: AdamW lr=0.1414 betas=(0.8, 0.995) eps=1e-10 wd=0.001
scheduler:
  trapezoidal:
    warmup_ratio: 0.05
    warmdown_ratio: 0.50
    final_lr_frac: 0.05

Fundamentos Teóricos

Este trabajo se conecta directamente con los fundamentos del aprendizaje profundo y la optimización de algoritmos. La elección de arquitecturas como Llama, con sus componentes (RMSNorm, RoPE, GQA), refleja la evolución de las redes neuronales transformer, cuya base teórica se estableció en el paper "Attention Is All You Need" (Vaswani et al., 2017). La optimización de la tasa de aprendizaje y la selección de optimizadores como AdamW y Muon se basan en décadas de investigación en optimización convexa y no convexa, con Adam (Kingma & Ba, 2014) siendo un punto de referencia clave.

La discusión sobre el impacto de la longitud del contexto en la evaluación de tareas como SQuAD y BoolQ subraya la importancia de la "long-range dependency" en los modelos de lenguaje, un problema central en el procesamiento del lenguaje natural que los transformers buscan resolver. La observación de que un modelo puede "empeorar" en una tarea específica debido a la truncación del contexto y la pérdida de ejemplos de pocas tomas (few-shot examples) ilustra cómo la interacción entre el diseño del modelo, el preprocesamiento de datos y la metodología de evaluación es crítica para obtener resultados significativos. Finalmente, la búsqueda de eficiencia computacional y la maximización del "capability per dollar" se alinea con la investigación en eficiencia energética y cómputo de alto rendimiento, un campo con raíces en la optimización de algoritmos y la arquitectura de computadoras.