La familia de modelos Llama 3

Guía de lectura

Problema y motivación

Los sistemas modernos de inteligencia artificial dependen en gran medida de modelos fundacionales masivos; sin embargo, los modelos más capaces (como GPT-4) han permanecido cerrados. La construcción de un modelo de acceso abierto y de clase frontera requiere superar enormes cuellos de botella de ingeniería en tres pilares: curación de datos, escala de cómputo y gestión de la complejidad.

Los autores desarrollan Llama 3, una familia de modelos de lenguaje multilingües que abarca 8B, 70B y 405B parámetros, admitiendo ventanas de contexto de hasta 128K, uso de herramientas zero-shot y multietapa, razonamiento avanzado y generación de código.

Ideas clave

• Gestión de la complejidad con arquitectura densa: En lugar de cambiar a arquitecturas Mixture-of-Experts (MoE), los autores se apegan a un Transformer denso estándar (126 capas, dimensión de 16,384, 128 cabezales de atención para 405B) con Grouped-Query Attention (GQA) y RoPE con frecuencia base $\theta =500,000$ para maximizar la estabilidad del preentrenamiento.
• Escalado de datos y filtrado riguroso: El preentrenamiento se expande desde los 1.8T tokens de Llama 2 hasta aproximadamente 15.6T tokens a lo largo de dominios diversos (~50% conocimiento general, 25% matemáticas/razonamiento, 17% código, 8% multilingüe).
• Leyes de escalado predecibles: Un marco de leyes de escalado en dos etapas vincula el cómputo con el rendimiento en tareas posteriores (como ARC Challenge) a través de cuatro órdenes de magnitud antes de comenzar el entrenamiento.
• Post-entrenamiento iterativo: Seis rondas de alineación que combinan Supervised Fine-Tuning (SFT), Rejection Sampling (RS) y Direct Preference Optimization (DPO) con generación dirigida de datos sintéticos.
• Paralelismo 4D e infraestructura: Preentrenamiento de 405B en hasta 16K GPUs H100 utilizando una combinación de Tensor Parallelism (TP), Context Parallelism (CP), Pipeline Parallelism (PP) y Fully Sharded Data Parallelism (FSDP/DP).

Método (con las ecuaciones clave)

1. Formulación de las leyes de escalado

La asignación de tokens óptima en términos de cómputo se modela como una función de ley de potencias del presupuesto de cómputo $C$:

$$N^{\star }(C)=A{C}^{\alpha }$$

Los ajustes empíricos sobre presupuestos de cómputo entre $6\times 10^{18}$ y $10^{22}$ FLOPs identificaron $(\alpha ,A)=(0.53,0.29)$. La extrapolación al presupuesto objetivo de $3.8\times 10^{25}$ FLOPs sugirió entrenar un modelo de 402B parámetros en 16.55T tokens, guiando la arquitectura final de 405B parámetros.

2. Extensión de la ventana de contexto y enmascaramiento de atención

El escalado del contexto de 8K a 128K tokens ocurre en seis etapas continuas de preentrenamiento a lo largo de ~800B tokens. Para evitar la contaminación cruzada entre documentos empaquetados dentro de la misma secuencia de contexto largo, se aplica una máscara de atención de documentos personalizada.

El Context Parallelism (CP) divide las secuencias de entrada en $2\times \text{CP}$ fragmentos, ejecutando un patrón de comunicación all-gather sobre las representaciones de clave ($K$) y valor ($V$). Debido a que GQA reduce los tamaños de $K$ y $V$ en relación con la consulta ($Q$), la sobrecarga de comunicación all-gather $O(S)$ es insignificante en comparación con el cómputo de la atención $O(S^2)$.

3. Adaptaciones de Direct Preference Optimization (DPO)

En la optimización de preferencias, los tokens de formato (encabezados y marcadores de terminación) se enmascaran explícitamente de la pérdida de DPO para evitar colas repetitivas o detenciones prematuras. Además, se añade un término auxiliar de regularización de log-verosimilitud negativa (NLL) a la secuencia elegida:

$${\mathcal{L}}_{\text{reg}}=0.2\cdot {\mathcal{L}}_{\text{NLL}}(\text{chosen})$$

Resultados (con las cifras más importantes)

• Benchmarks del modelo insignia 405B:
  • MMLU (5-shot): 87.3% (0-shot CoT: 88.6%)
  • GSM8K (8-shot, CoT): 96.8%
  • MATH (0-shot, CoT): 73.8%
  • HumanEval (0-shot): 89.0%
  • ARC Challenge (0-shot): 96.9%
  • IFEval: 88.6%
• Modelos más pequeños (8B y 70B): Ambos modelos superan a sus pares de código abierto comparables (p. ej., Gemma 2 9B, Mistral 7B, Mixtral 8x22B). Por ejemplo, Llama 3 8B alcanza un 73.0% en 0-shot CoT MMLU y un 84.5% en GSM8K.
• Contexto largo: Recuperación perfecta del 100% en Needle-In-A-Haystack simple de hasta 128K tokens, y 98.1% en Multi-needle recall para el modelo 405B.
• Eficiencia del preentrenamiento: La utilización de FLOPs del modelo (MFU) en BF16 alcanzó el 38–43% en todas las configuraciones sobre 16,384 GPUs H100, manteniendo más del 90% de tiempo de entrenamiento efectivo durante una captura de 54 días a pesar de 466 interrupciones.

Limitaciones y preguntas abiertas

• Fiabilidad del hardware y fluctuación a escala: A una escala de 16K GPUs, las interrupciones de hardware promediaron varias veces al día (el 78% de las interrupciones inesperadas se debieron a problemas de hardware/GPU). Las fluctuaciones de energía simultáneas de decenas de megavatios durante el guardado de checkpoints o pasos colectivos desafían las redes eléctricas de los centros de datos.
• Brecha de robustez adversarial: Aunque es resistente en benchmarks de paráfrasis (PAWS), los benchmarks adversariales de razonamiento matemático y preguntas y respuestas (GSM-Plus, Adv SQuAD) demuestran una degradación significativa en comparación con los benchmarks no adversariales.
• Disponibilidad multimodal: Los adaptadores de visión y voz demostraron capacidades competitivas en entornos de investigación, pero se mantuvieron en desarrollo activo y se excluyeron del lanzamiento público inicial.