- 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
1 comentarios
Opiniones de Hacker News
Dice que le gustaría ver su charla sobre criptografía cuántica, y pregunta si alguien conoce el enlace de esa charla mencionada al inicio del video.
El entrenamiento optimizado de modelos pequeños es importante no solo por la accesibilidad, sino también para la investigación científica de los LLM; así como en biología se usan organismos simples como la levadura, hace falta estudiar el transformador más simple que muestre comportamientos interesantes de los modelos grandes para poder entenderlos y controlarlos.
Uno de los podcasts más interesantes que escuchó recientemente fue sobre el paper y dataset Tiny Stories; este dataset solo incluye palabras y conceptos simples, como cuentos infantiles, pero aun así permite que modelos pequeños generen inglés con gramática, diversidad y razonamiento. El podcast con el autor también explica muy bien las capacidades de los LLM con un ejemplo de investigación pequeño y controlado. Siguiendo la analogía biológica, cree que el dataset sería algo así como una placa de agar muy simple y controlada. Los enlaces relacionados son el episodio del podcast y el paper de TinyStories.
Muchas empresas pueden resolver problemas reales de negocio con modelos pequeños usando datasets privados, como el historial de compras de sus usuarios. Los avances logrados con modelos de lenguaje grandes también pueden aplicarse tal cual a problemas pequeños, si la secuencia de entrada puede expresarse en un lenguaje especializado.
Es bien sabido que los comportamientos y optimizaciones observados en modelos pequeños no se reproducen bien en modelos grandes.
Lo que hace el autor aquí es pretraining, algo que normalmente hacen creadores de modelos como Google o Meta; para negocios, el fine-tuning o, en menor medida, el pretraining adicional es mucho más práctico. Remarca que el autor lo intenta por razones académicas.
Le interesan los modelos que corren rápido en una laptop, pero el proceso de entrenamiento por sí solo puede tardar días o más.
Cree que estaría bien usar energía en vez de tiempo como criterio, es decir, entrenar el mejor modelo posible dentro de un presupuesto energético dado en joules; eso haría más justa la comparación entre una MBP y una H100.
El punto clave aquí no es la eficiencia sino la “accesibilidad”, porque una H100 no es un producto de uso cotidiano, mientras que una laptop sí.
Según entiende, las Mac son más competitivas en consumo eléctrico, ya que no jalan tanta energía como una GPU de Nvidia. Además, se puede rentar una H100 por menos de 10 dólares la hora, así que también sería interesante comparar qué tan bueno es un modelo que pueda entrenarse en menos de una hora.
Le parece bien cualquier criterio; aunque sea un poco arbitrario, no lo ve mal.
Si la idea es expresar una métrica práctica basada en laptop o en MacBook Pro, le gustaría que eso se aclarara mejor.
Dice en tono divertido que ya solo falta organizar unas olimpiadas de eficiencia de IA: en laptop, desktop o teléfono; en 5 minutos, 1 hora, 1 día o 1 semana; arriba de un barco o junto a una cabra, en cualquier lugar.
Lo de la cabra probablemente se refiere a Llama; la rima no da para mucho, pero con acento de Boston sería más chistoso.
Menciona una novela de Vernor Vinge en la que los humanos fabrican computadoras portátiles de ajedrez para usarlas como asistentes durante una partida; cree que sería curioso que en un torneo, además del reloj de ajedrez, también se diera energía eléctrica, para que los participantes pensaran jugadas que favorezcan a su propia IA.
Hace una broma exagerando el rendimiento de una Mac Studio M3 Ultra de 512GB, diciendo que con ese barco hasta se podría transportar una cabra.
Bromea con que la cabra tiene demasiados parámetros, casi al nivel de GPT-4.
Si sale GoatLM, estaría dispuesto a pagar por él.
Le gustó la expresión “officially major people” en el ejemplo absurdo de "Paris, France is a city in North Carolina...", y se pregunta cómo podría usarla en una conversación cotidiana.
Opina que esto le recordó al paper de “cramming” de hace unos años, y comparte el enlace al paper, que trata de intentar entrenar un modelo óptimo en una laptop moderna durante un día.
Cree que solo con aplicar algunos trucos tomados del intento de speedrun de GPT-2 —Muon, mejor inicialización de pesos y un ajuste cuidadoso de la tasa de aprendizaje— ya podría mejorar bastante más; deja el material relacionado aquí.
Cree que a la IA le falta una demoscene, esa cultura de presumir técnica con recursos mínimos como en las demos gráficas de los 90.
Siente que tiene valor crear modelos pequeños y más especializados, e incluso armarlos al momento según la necesidad; no hace falta un modelo grande que sepa de todo, sino uno que esté enfocado con precisión láser solo en el dominio en el que trabaja y que además corra muy rápido. Quiere poder pedirle a un LLM grande: “escríbeme un script para entrenar un modelo optimizado para <tarea necesaria>”, y luego ejecutar ese modelo. Pero mientras escribía el comentario, Google lanzó Gemma 3 270M.
Poniendo como ejemplo disparates como "Paris, France is a city in North Carolina...", dice que si existiera una técnica para poder decir “no sé”, los modelos pequeños serían mucho más útiles. Lo que hace falta en LLM enormes es precisamente evitar que inventen cosas en un rango tan amplio. Sería bueno poder entrenar un chatbot de atención al cliente en una laptop en un tiempo razonable, pero fuera de ese dominio probablemente tendría muchas posibilidades de responder cosas seriamente fuera de lugar.