Reimplementando el modelo GPT-2 de forma económica con `llm.c`

• Explica cómo reproducir el modelo GPT-2 (124M) con llm.c en 90 minutos por $20
• GPT-2 (124M) es el modelo más pequeño que OpenAI publicó en 2019
• Usando un nodo 8X A100 80GB SXM en Lambda cuesta alrededor de $14 por hora, para un costo total de unos $20
  • También se puede entrenar con una sola GPU, pero tarda más (4-24 horas)

Comparación de resultados

• En el conjunto de validación FineWeb muestra mejor rendimiento que el checkpoint publicado por OpenAI
  • Aun así, GPT-2 fue entrenado con WebText, así que no es una comparación completamente justa
• También se midió la precisión en HellaSwag, logrando 29.9, cerca del 33.7 de GPT-3 Small (124M)
  • Ya supera el 29.4 de GPT-2 (124M)
  • Pero aquí se entrenó con 10B tokens, mientras que GPT-3 fue entrenado con 300B tokens

Configuración mínima del entorno

• Se requiere GPU (se recomienda Lambda Labs)
• Guía basada en Linux x86 de 64 bits Ubuntu 22.04 con CUDA 12
• Instalar miniconda y luego la versión nightly de PyTorch (opcional)
• Instalar los paquetes necesarios para el tokenizer
• Instalar cuDNN para mejorar la velocidad (opcional)
• Instalar MPI al usar varias GPU (opcional)
• Preprocesar el dataset FineWeb de 10B tokens (~1 hora)
• Compilar llm.c (precisión mixta, usando cuDNN FlashAttention)

Ejecución del entrenamiento

• Ejemplo de comando para usar una sola GPU
• Ejecución con mpirun al usar múltiples GPU (8)
• Explicación de los argumentos principales
  • -i, -j : rutas de los datos de entrenamiento/validación
  • -o : ruta para guardar logs y checkpoints
  • -e : inicialización del modelo (GPT-2 depth 12)
  • -b : tamaño de microbatch (reducir si falta memoria)
  • -t : longitud máxima de secuencia
  • -d : tamaño total del batch (ver el paper de GPT-3)
  • -r : configuración de Recompute (ahorro de memoria)
  • -z : ZeRO-1 (sharding del estado del optimizador)
  • Otras opciones para weight decay, learning rate, frecuencia de checkpoints, etc.

Proceso de entrenamiento

• Con 10B tokens de entrenamiento y un batch size de 0.5M, se estiman unas 20K iteraciones
• En una A100 40GB PCIe se muestran por iteración el tiempo requerido, MFU y throughput de tokens
• Al inicio del entrenamiento hay exploding gradients, pero se resuelve con clipping

Visualización

• Se proporciona un notebook de Jupyter para visualizar la curva de entrenamiento analizando el archivo de logs

Tokenizer

• Es necesario para convertir tokens enteros en texto
• Se puede generar con un script de PyTorch

Sampling

• Por ahora no está optimizado para inferencia
• Con algunas pequeñas modificaciones en el código se puede hacer sampling unconditional/conditional

Estructura del código

• La mayor parte de la implementación está en el archivo train_gpt2.cu
• Las primeras 500 líneas configuran MPI, NCCL, cuDNN, cuBLAS, etc.
• Las siguientes 1500 líneas contienen el forward/backward del Transformer
• Las siguientes 1000 líneas implementan el modelo GPT-2
• Las últimas 1000 líneas incluyen el loop de entrenamiento, parsing de argumentos, etc.

Modelo 350M

• 10B tokens no son suficientes y se usan 30B tokens
• En 8X A100 80GB toma 14 horas, con un costo aproximado de $200

FAQ

• ¿Se puede hacer sampling?: Sí, pero es ineficiente.
• ¿Se puede chatear?: Por ahora solo permite preentrenamiento; no se puede hacer fine-tuning para chat.
• Entrenamiento distribuido multinodo: Es posible, pero todavía no se ha probado.
• ¿Es determinístico a nivel de bits?: Es casi determinístico, pero hacen falta algunos parches en kernels.
• ¿Se puede entrenar con FP8?: Por ahora se entrena con BF16; FP8 llegará pronto.
• ¿Soporta GPU que no sean NVIDIA?: Por ahora solo soporta C/CUDA.

Opinión de GN⁺

• GPT-2 es un modelo muy importante como punto de partida de los LLM modernos. Después, GPT-3 y otros LLM no son tan distintos de GPT-2 en esencia.
• Este proyecto permite que cualquiera pueda entrenar directamente un modelo de nivel GPT-2 a un costo razonable. Parece muy útil para mejorar la comprensión sobre los LLM.
• Sin embargo, todavía no está optimizado para inferencia, así que su uso en servicios reales sigue siendo limitado. Tampoco soporta fine-tuning hacia modelos conversacionales.
• Actualmente solo soporta GPU de NVIDIA, pero se espera compatibilidad futura con plataformas diversas como AMD o Apple Silicon.
• Entre los proyectos open source con objetivos similares están Megatron-LM, DeepSpeed y FairSeq. Cada uno tiene sus ventajas y desventajas, así que conviene elegir según el caso de uso.
• Es un proyecto muy alentador para dinamizar el ecosistema de desarrollo de LLM. Dan ganas de ver qué viene después.