Quel est le modèle d’IA le plus puissant qu’on peut entraîner sur un ordinateur portable en 5 minutes ?

• Environ 1,8M de paramètres pour un modèle transformer de style GPT entraîné en 5 minutes sur un MacBook Pro avec quelque 20M de tokens TinyStories, pour atteindre une perplexité d’environ 9,6
• La principale contrainte d’un entraînement en moins de 5 minutes est la taille du modèle et le nombre de tokens traitables ; plus le modèle est grand, plus la convergence est lente et moins un petit volume de données est efficace
• Côté optimisation des performances, le plus efficace est de choisir un petit modèle plutôt que d’utiliser MPS, la compilation/quantification/l’accumulation de gradient ou des alternatives à PyTorch
• Des jeux de données simples et cohérents comme TinyStories ont un impact plus positif sur les petits modèles que des données de type encyclopédique
• L’architecture transformer donne de meilleurs résultats que les approches LSTM ou diffusion dans des conditions de petite taille et de temps d’entraînement très court

---

Vue d’ensemble

Cet article présente les résultats d’une expérience visant à déterminer le modèle de langage IA le plus performant qu’il est possible d’entraîner en 5 minutes sur un ordinateur portable (MacBook Pro), ainsi que des enseignements sur la stratégie d’entraînement optimale, le choix du jeu de données et l’architecture du modèle.

Résumé des résultats expérimentaux

• Un modèle transformer de style GPT d’environ 1,8M de paramètres a été entraîné sur environ 20M de données TinyStories, avec une perplexité de 9,6
• Les exemples générés sont courts mais prennent la forme d’histoires cohérentes, avec un anglais globalement correct sur le plan grammatical
• L’auteur souligne que, pour un modèle entraîné en 5 minutes, le niveau obtenu est plus pratique et convaincant qu’attendu

Contexte de l’expérience et limites

• L’expérience est née d’une curiosité peu réaliste : entraîner rapidement un modèle puissant dans un environnement d’ordinateur portable
• En pratique, il est possible d’entraîner des modèles bien plus puissants dans le cloud avec des GPU haut de gamme (comme le H100), mais ici la contrainte de l’expérience est le temps : 5 minutes
• À mesure que la taille du modèle augmente, la vitesse de traitement des tokens diminue, ce qui rend difficile l’obtention de bons résultats en 5 minutes
  • Les modèles trop petits (par exemple 10K paramètres) n’apprennent pas une complexité suffisante
  • La plage réellement exploitable se situe autour de 1M à 2M de paramètres

Optimisation du débit

• L’utilisation de MPS (Metal Performance Shaders d’Apple) est ce qui fonctionne le mieux
• Diverses optimisations mathématiques comme torch.compile, float16, MLX, etc., apportent moins de gains que prévu, voire dégradent les performances
• L’accumulation de gradient peut servir à gérer la mémoire, mais en pratique elle ralentit fortement l’exécution
• Pour être efficace, le modèle doit pouvoir mettre à jour rapidement ses poids dans la mémoire interne

Choix du jeu de données

• Avec un nombre limité de tokens (environ 10 à 20M), des données de type Simple English Wikipedia ont d’abord été utilisées ; elles permettaient d’obtenir une certaine cohérence grammaticale, mais pas de réelle cohérence sémantique
  • La prédominance des noms propres et l’énumération de faits au rendu artificiel limitaient la génération de contenu réellement pertinent
• Avec le jeu de données TinyStories, la structure narrative est claire et la langue est simple, ce qui produit des résultats bien plus cohérents et plus riches de sens
  • Comme il s’agit d’histoires de niveau enfant de 4 ans, même un petit modèle les apprend bien

Tokenizer et tokenisation

• L’entraînement du tokenizer n’est pas inclus dans les 5 minutes et, vu la petite taille des données, il y a peu de besoin d’optimisation
• L’apprentissage de tokens multioctets est plus facile pour le modèle

Expérimentations sur l’architecture du modèle

• Utilisation d’une architecture transformer (style GPT-2)

  • Réglage de divers hyperparamètres comme 2 à 3 couches, des fonctions d’activation comme SwiGLU, les positional embeddings, etc.
  • Les LSTM obtiennent des performances proches, mais le transformer reste meilleur en termes de perplexité
  • Le dropout et le mixture-of-experts sont inefficaces à cette petite échelle
  • Le curriculum learning a peu d’effet car la durée d’entraînement est trop courte

• Tentative avec un modèle de diffusion (D3PM)

  • Comme le langage naturel est composé de tokens discrets, le processus de diffusion ne génère que des tokens aléatoires sans signification, ce qui mène à un échec
  • Il est plus difficile d’y former rapidement une structure de phrase que dans un transformer ou un LSTM

Relation entre taille du modèle et débit en tokens/s

• Les modèles de 1M à 2M de paramètres constituent le meilleur sweet spot
  • Trop grands, ils ne convergent pas en 5 minutes ; trop petits, ils atteignent immédiatement leurs limites de performance
• La loi de scaling de Chinchilla correspond globalement aux résultats observés
  • Une taille de modèle idéale d’environ le nombre total de tokens d’entraînement / 20 a également été confirmée dans cette expérience

Conclusion et implications

• Même avec très peu de temps et un matériel limité, il est possible d’entraîner un modèle de storytelling cohérent
• Un entraînement de 5 minutes n’est pas adapté au développement de modèles puissants, mais il a un intérêt pour la conception de modèles petits et ultra-légers ainsi que pour les expérimentations d’optimisation matérielle et architecturale
• Avec les progrès futurs des GPU pour ordinateurs portables et des architectures de modèles, les performances de modèles entraînables en quelques minutes pourraient encore progresser