SmolGPT: implementación mínima en PyTorch para entrenar un LLM pequeño desde cero

• Proyecto sencillo de entrenamiento de LLM basado en PyTorch, diseñado con fines educativos y aplicando entrenamiento eficiente y técnicas modernas de muestreo

Características

• Base de código mínima: implementado en PyTorch, sin abstracciones complejas
• Arquitectura moderna: modelo GPT con Flash Attention, RMSNorm, SwiGLU y técnicas eficientes de muestreo
• Funciones de entrenamiento:
  • Precisión mixta (bfloat16/float16)
  • Acumulación de gradientes
  • Reducción de la tasa de aprendizaje con warmup
  • Weight decay y gradient clipping
• Soporte de datasets: procesamiento integrado del dataset TinyStories
• Tokenizer personalizado: integración del entrenamiento del tokenizer SentencePiece

Instalación

• Python 3.8+
• PyTorch 2.0+ y CUDA
• Se recomienda una GPU moderna

Inicio rápido 🚀

Opción 1: ciclo completo de entrenamiento

1. Preparar el dataset
2. Iniciar el entrenamiento
3. Generar texto

Opción 2: usar un modelo preentrenado

1. Descargar los assets
2. Ejecutar la inferencia

Detalles del modelo preentrenado

• Entrenado con el dataset TinyStories
• Vocabulario de 4096 tokens, transformador de 8 heads y 8 capas, dimensión de embedding de 512
• Entrenado durante 18.5 horas con aproximadamente 400 millones de tokens
• Pérdida de validación: 1.0491

Especificaciones del equipo de entrenamiento

• GPU: NVIDIA L4 Tensor Core
• vCPUs: 16
• RAM: 64 GB
• VRAM: 24 GB

Nota: esta implementación fue creada con fines educativos tomando como referencia prácticas modernas de entrenamiento de LLM. En entornos de producción, se recomienda escalar el tamaño del modelo y del dataset.

[GN⁺]

Comentarios de Hacker News

• Un usuario mencionó que, al trabajar en una implementación de GPT en C, aprendió la importancia de la gestión de memoria y de datos. El proyecto consta de unas 1500 líneas de código y compartió un enlace de GitHub

• Otro usuario explicó que implementar todo directamente con numpy es el siguiente paso. Destacó que se puede empezar usando abstracciones potentes y luego ir eliminándolas gradualmente para comprender por completo el sistema

• Compartió un proyecto que tarda alrededor de 2 horas en una GPU A100 usando un cuaderno de Google Colab. Mencionó que también podría funcionar con una cuenta gratuita

• Explicó que en GitHub hay varios proyectos, como nanoGPT de Andrej Karpathy, y que también existen otros proyectos con MoE implementado

• Un usuario preguntó si era posible entrenar un modelo específico con un conjunto de datos específico y probar los resultados. Mencionó que está buscando un proyecto de código abierto, que indique si se puede usar GPU y que ofrezca un contenedor Docker que pueda reemplazarse por CPU

• Otro usuario compartió su experiencia implementándolo con su propio tokenizador multicanal. Comentó que el nombre del protagonista se repite con frecuencia, por lo que sospecha que podría haber un bug

• Om Alve agradeció y comentó que, como la publicación se volvió viral, siente que el esfuerzo valió la pena

• Mencionó que resulta interesante que la tecnología pueda implementarse en unos pocos cientos de líneas de código. Propuso adivinar cuántas líneas de código tienen los modelos de vanguardia

• Aclaró que no está relacionado con SmolLM y compartió un enlace al blog de HuggingFace

• Preguntó si podían compartir ejemplos del dataset de entrenamiento y casos de uso