¿Cuál es el modelo de IA más potente que se puede entrenar en una laptop en 5 minutos?

• En una MacBook Pro, se entrenó en 5 minutos un modelo transformer estilo GPT de aproximadamente 1.8M de parámetros con unos 20M de tokens de TinyStories, logrando una perplejidad de ~9.6
• Las principales restricciones del entrenamiento en menos de 5 minutos son el tamaño del modelo y la cantidad de tokens que se pueden procesar; si el modelo es grande, converge más lento y rinde menos con pocos datos
• En la optimización del rendimiento, usar MPS, compilación/cuantización/acumulación de gradientes o reemplazar PyTorch resulta menos efectivo que elegir un modelo pequeño
• Un dataset simple y consistente como TinyStories tiene un impacto más positivo en modelos pequeños que datos enciclopédicos
• La arquitectura transformer mostró mejores resultados que los enfoques LSTM o diffusion bajo condiciones de tamaño pequeño y tiempo de entrenamiento corto

---

Resumen general

Este artículo presenta los resultados de un experimento para encontrar el modelo de lenguaje de IA de mayor rendimiento que pueda entrenarse en una laptop (MacBook Pro) en 5 minutos, junto con ideas sobre la estrategia de entrenamiento óptima, la selección de datasets y la arquitectura del modelo.

Resumen de resultados experimentales

• Se entrenó un modelo transformer estilo GPT de aproximadamente 1.8M de parámetros con cerca de 20M de datos de TinyStories, registrando una perplejidad de 9.6
• Los ejemplos generados son historias cortas pero consistentes, con una gramática inglesa en general correctamente mantenida
• Se enfatiza que los resultados obtenidos por un modelo en 5 minutos superan lo esperable a un nivel práctico

Contexto y límites del experimento

• El experimento partió de una curiosidad poco realista: entrenar rápidamente un modelo potente en un entorno de laptop
• En la práctica, se pueden entrenar modelos mucho más potentes en la nube con GPUs de alto rendimiento (como H100), pero la condición limitante del experimento fue el tiempo: 5 minutos
• A mayor tamaño del modelo, más lenta es la velocidad de procesamiento de tokens, lo que dificulta obtener buenos resultados en 5 minutos
  • Los modelos demasiado pequeños (por ejemplo, de 10K parámetros) no logran aprender suficiente complejidad
  • El rango práctico es de aproximadamente 1M a 2M de parámetros

Optimización del rendimiento

• Lo más efectivo fue usar MPS (Metal Performance Shaders de Apple)
• Diversas optimizaciones matemáticas como torch.compile, float16 y MLX mostraron mejoras mucho menores a las esperadas o incluso un peor rendimiento
• La acumulación de gradientes sirve para gestionar memoria, pero en la práctica provoca una fuerte caída de velocidad
• Para ser eficiente, el modelo debe poder actualizar sus weights rápidamente dentro de la memoria interna

Selección del dataset

• Al usar primero datos simples de wiki en inglés, como Simple English Wikipedia, con una cantidad limitada de tokens (aprox. 10~20M), se logró cierta consistencia gramatical, pero faltó consistencia semántica
  • Al centrarse en nombres propios y enumeraciones de hechos que parecían forzadas, hubo límites para generar contenido significativo
• Con el dataset TinyStories, como la estructura narrativa es clara y el lenguaje simple, el resultado fue mucho más consistente y con más sentido
  • Al ser historias de nivel de 4 años, incluso un modelo pequeño puede aprenderlas bien

Tokenizer y tokenización

• El entrenamiento del tokenizer no está incluido dentro de los 5 minutos y, como la escala de datos es pequeña, la necesidad de optimización es baja
• Aprender tokens multibyte es más fácil para el entrenamiento del modelo

Experimentos de arquitectura del modelo

• Se usó una arquitectura transformer (estilo GPT-2)

  • Se ajustaron hiperparámetros como 2~3 capas, funciones de activación como SwiGLU y positional embedding
  • El LSTM tuvo un rendimiento cercano, pero el transformer fue superior en términos de perplejidad
  • Dropout y mixture-of-experts resultan ineficientes a esta escala pequeña
  • El curriculum learning tuvo poco efecto porque el tiempo de entrenamiento era demasiado corto

• Se probó un modelo de diffusion (D3PM)

  • Como el lenguaje natural está compuesto por tokens discretos, el proceso de difusión solo generó tokens aleatorios sin sentido y fracasó
  • Fue difícil formar estructura de oraciones rápidamente en comparación con transformer o LSTM

Relación entre tamaño del modelo y rendimiento de tokens/segundo

• Los modelos de 1M a 2M de parámetros son el punto ideal
  • Si son demasiado grandes, no convergen dentro de 5 minutos; si son demasiado pequeños, alcanzan su límite de rendimiento apenas empiezan a entrenarse
• La ley de escalado de Chinchilla coincide en términos generales con los resultados experimentales
  • El tamaño ideal del modelo es aproximadamente el total de tokens de entrenamiento dividido entre 20, y eso también se confirmó en este experimento

Conclusión e implicaciones

• Incluso con muy poco tiempo y hardware pequeño, sí es posible entrenar un modelo de storytelling consistente
• Entrenar durante 5 minutos no es adecuado para desarrollar modelos potentes, pero sí tiene valor para diseñar modelos pequeños y ultraligeros, además de experimentar con optimización de hardware y arquitectura
• A medida que avancen las GPUs para laptops y la estructura de los modelos, existe potencial para que siga mejorando el rendimiento de modelos entrenables en apenas unos minutos