Quel est le modèle d'IA le plus puissant que vous pouvez entraîner sur un ordinateur portable en cinq minutes ?

Par Sean Goedecke

Quel est le modèle d'IA le plus puissant que vous pouvez entraîner sur un ordinateur portable en cinq minutes ? par Sean Goedecke

Quel est le modèle le plus puissant que je peux entraîner sur mon MacBook Pro (2024 MacBook Pro, Apple M4, 24 GB de mémoire) en cinq minutes ?

Je vais vous donner la réponse d'emblée : le meilleur modèle que j'ai pu entraîner en 5 minutes était un transformateur de type GPT d'environ 1,8 million de paramètres, entraîné sur environ 20 millions de jetons TinyStories, atteignant une perplexité d'environ 9,6 sur un échantillon testé séparément. Voici un exemple du résultat, avec l'instruction génératives en gras :

Bon, ce n'est pas génial. Mais ce n'est pas mal pour cinq minutes !

Le défi

Cette question idiote m'intéresse depuis quelques jours. Elle est idiote pour deux raisons. Premièrement, toute personne qui peut s'offrir un MacBook peut se permettre de louer une demi-heure sur un H100 et de former un modèle plusieurs fois plus puissant. Deuxièmement, si vous étiez obligé de vous entraîner sur un appareil moins puissant, comme un ordinateur portable, il n'y a aucune raison de vous limiter à cinq minutes (et aucune raison de penser qu'il serait même possible de former un modèle puissant en si peu de temps).

D'autres défis de formation, comme BabyLM, limitent les données de formation, ce qui est logique : certains domaines peuvent disposer de très peu de données, il est donc utile de savoir comment former un modèle de la manière la plus efficace possible lorsque les données sont rares. Essayer de former le plus petit modèle puissant possible est également un objectif de recherche populaire, ce qui est également logique, car vous pouvez exécuter de petits modèles sur des téléphones et des appareils portables. Mais je peux utiliser autant de données d'entraînement que je le souhaite et un modèle aussi grand que je le souhaite. Ma principale limite est le temps.

En cinq minutes, vous ne pouvez tout simplement pas faire passer autant de jetons dans un modèle. Cela signifie que les grands modèles sont hors de question, car il faut plus de temps par jeton pour entraîner un modèle plus grand. Il vaut mieux entraîner un modèle à 1 million de paramètres sur 4 millions de jetons qu'un modèle à 1 milliard de paramètres sur 4 000 jetons. Mais bien sûr, vous ne pouvez pas aller trop petit. En cinq minutes, je peux faire passer beaucoup de jetons dans un petit modèle à 10 000 paramètres, mais ce modèle n'est tout simplement pas assez grand pour apprendre la grammaire anglaise. La perte d'entraînement plafonne dans les trente premières secondes et ne bouge plus après cela, et le modèle ne produit que du charabia.

Augmenter le nombre de jetons par seconde

Mon premier objectif était de déterminer quelles optimisations de performances seraient réellement utiles à cette échelle minuscule. Mon premier transformateur de type GPT-2, tel que décrit dans les manuels, s'entraînait à environ 3 000 jetons par seconde (en utilisant le MPS d'Apple). Il est intéressant de noter que les optimisations mathématiques m'ont soit ralenti, soit n'ont eu aucun effet significatif : utilisation de torch.compile sur le modèle, passage à float16, etc.

J'ai essayé l'accumulation de gradients (mise à jour des poids du modèle en plusieurs lots) en partant du principe que j'étais également limité en mémoire, mais comme on pouvait s'y attendre, cela m'a ralenti d'un ordre de grandeur. Il s'avère que le plus gros goulot d'étranglement pour l'entraînement sur un MacBook est simplement le « lancement » (c'est-à-dire demander au GPU d'aller faire quelque chose). Le type de modèle que vous pouvez entraîner sur un ordinateur portable doit pouvoir exécuter l'étape de mise à jour des poids en mémoire.

Il est intéressant de noter que le passage de PyTorch à MLX n'a pas amélioré les performances de manière significative. J'espérais obtenir un gain de vitesse de 2x, mais cela n'a pas été le cas. Donc, dans l'ensemble, mes recommandations pour gagner en vitesse sont les suivantes : utilisez MPS (évidemment), ne vous souciez pas de la compilation ou de la quantification, ne faites pas d'accumulation de gradients et choisissez le modèle le plus petit possible.

Choisir le bon ensemble de données

Si vous ne disposez que d'environ dix millions de jetons pour l'entraînement, de quel ensemble de données les extraire ? Cela représente environ 50 Mo de données textuelles, ce qui est loin d'être suffisant pour constituer un échantillon représentatif de tous les textes en anglais. Si 10 Mo de ces données proviennent de Shakespeare, 10 Mo d'autres de manuels de physique, etc., il n'y aura tout simplement pas assez de modèles cohérents pour le modèle.

J'ai commencé avec un ensemble de données tiré de Wikipédia en anglais simple, ce qui semblait judicieux : l'utilisation d'une grammaire et d'un vocabulaire simples signifie moins d'apprentissage pour le modèle. Cela a bien fonctionné, mais il y avait un autre problème. Voici un exemple du résultat que j'ai obtenu (le texte en gras est mon instruction générative) :

Je trouve cela assez impressionnant pour cinq minutes d'entraînement. Il sait qu'une ville se trouve à un certain endroit et qu'une ville peut également être la capitale de cet endroit. Il se dégrade un peu après cela, mais la grammaire anglaise est correcte et le texte qu'il génère porte toujours sur les villes et les pays. Cependant, il est clair qu'il ne génère pas de contenu cohérent, car il est obsédé par les noms propres. Chaque phrase parle du nom X se trouvant dans le nom Y. C'est l'inconvénient de former un modèle à partir d'une encyclopédie : il est difficile de produire un contenu impressionnant lorsque l'on récite une liste de faits inventés.

Comme la plupart des expériences avec des petits modèles, j'ai fini par utiliser TinyStories : un ensemble de données synthétiques composé de nouvelles d'un paragraphe, écrites à un niveau de lecture adapté à des enfants de 4 ans. Cela présente de nombreux avantages pour l'entraînement de petits modèles de langage. Les histoires sont toutes cohérentes, avec une relation de cause à effet et même généralement une morale. Il y a très peu de noms propres à mémoriser pour le modèle, et le langage est encore plus simple que celui de Wikipédia en anglais simplifié. L'article TinyStories montre qu'il est possible de former un modèle de langage de 1 million de paramètres pour produire des histoires de ce type assez correctement (même si cela n'a pas été fait en cinq minutes). Un modèle de cette taille est suffisamment petit pour être formé sur mon ordinateur portable !

Tokenisation

Une remarque sur la tokenisation. Je n'ai pas inclus le temps passé à former le tokeniseur dans mon budget de cinq minutes, mais pour un modèle de cette taille, cela n'a vraiment aucune importance. Peu importe le soin que vous apportez à l'optimisation, vous ne pouvez tout simplement pas traiter plus de quelques centaines de mégaoctets de données pendant le processus d'entraînement en cinq minutes. Tant que vous ne faites rien de stupide, comme tokeniser à l'avance l'ensemble des données TinyStories alors que vous n'allez en utiliser qu'une fraction, la tokenisation ne prend que quelques secondes.

C'est pourquoi je n'ai pas passé beaucoup de temps à optimiser l'étape de tokenisation. J'aurais pu essayer de m'entraîner sur des caractères bruts plutôt que sur des tokens multioctets, par exemple, mais les tokens multioctets sont très probablement plus faciles à apprendre pour un modèle.

Architecture

Quel type d'architecture de modèle fonctionne le mieux pour un entraînement de cinq minutes ? J'ai essayé plusieurs choses différentes, plus ou moins au hasard.

Transformateurs et LSTMS

Le choix évident pour une architecture de modèle est un transformateur de type GPT-2. Il est simple à mettre en œuvre, très bien pris en charge et fonctionne avec un faible nombre de paramètres. Je n'ai pas pris la peine d'essayer un mélange d'experts, cela me semblait ridicule pour un modèle à 1 million de paramètres. Je n'ai pas utilisé de dropout, car je ne m'inquiétais pas du surajustement et j'aurais été heureux de voir n'importe quel type d'ajustement.

J'ai passé la plupart de mon temps à ajuster les hyperparamètres du transformateur (même si, comme vous le verrez, j'ai essayé d'autres types de modèles). En termes de structure réelle du modèle, j'ai constaté que SwiGLU faisait une différence notable et que j'obtenais les meilleurs résultats avec des modèles à deux ou trois couches. J'ai essayé plusieurs taux d'apprentissage différents et tout ce qui se situait entre 0,001 et 0,002 semblait convenir : c'est assez élevé, mais lorsque vous essayez de converger en cinq minutes, c'est ce qu'il vous faut. J'ai également testé l'apprentissage par programme, qui consiste à commencer par des séquences plus courtes et à augmenter la taille à mesure que le modèle apprend des modèles plus simples, mais comme vous pouvez vous en douter, cinq minutes étaient trop courtes pour que cela en vaille vraiment la peine. En termes d'attention, j'ai essayé à la fois avec RoPE et avec des intégrations positionnelles apprises, et les intégrations positionnelles semblaient nettement meilleures.

J'ai également essayé de former quelques modèles LSTM (longue mémoire à court terme) de taille similaire, mais je n'ai pas obtenu de résultats très satisfaisants. C'était assez similaire à mes expériences avec les transformateurs, comme on pouvait s'y attendre, mais la perplexité n'était pas aussi faible qu'avec des transformateurs de même taille.

Diffusion

Juste pour le plaisir, j'ai entraîné quelques petits modèles de langage de diffusion avec D3PM. Les modèles de langage de diffusion sont difficiles à mettre en œuvre car les jetons linguistiques sont discrets, contrairement aux valeurs de pixels dans une image, de sorte que toute quantité de bruit peut transformer un jeton valide en un jeton invalide. D3PM tente de résoudre ce problème en appliquant un bruit « structuré » (par exemple, en remplaçant un jeton par un jeton [MASK] ou un jeton différent).

Je n'ai absolument pas réussi avec cette approche. Voici un exemple du type de résultat que j'ai obtenu :

Comme vous pouvez le constater, il s'agit simplement de jetons aléatoires. Le processus de débruitage n'a produit aucune structure significative. Je ne m'attendais pas à ce que cela soit compétitif - intuitivement, apprendre à construire une phrase en plaçant des mots au hasard est plus difficile qu'apprendre à la construire dans l'ordre - mais je suis tout de même un peu surpris que le résultat soit aussi mauvais. Les transformateurs et les LSTM produisaient des résultats grammaticalement plausibles après la première minute d'entraînement.

Taille du modèle

Enfin, la question la plus importante : quelle taille de modèle peut-on entraîner de manière utile en cinq minutes ? J'ai essayé plusieurs tailles différentes :


Vous pouvez voir que le point idéal se situe autour de 2 millions de paramètres : au-delà, le modèle est trop lent pour converger en cinq minutes, en dessous, il cesse de s'améliorer après la première minute (probablement parce qu'il n'y a pas assez de paramètres pour encoder les modèles dans les données d'entraînement).

Il est intéressant de noter que cela coïncide plus ou moins avec le célèbre article sur les lois d'échelle de Chinchilla, qui indique que la taille optimale de votre modèle correspond au nombre total de jetons d'entraînement divisé par 203. Pour ce défi, le nombre de jetons d'entraînement dépend également de la taille du modèle (car vous pouvez faire passer plus de jetons dans un modèle plus petit en cinq minutes). Pour un modèle de 2,6 millions de paramètres, je peux entraîner 56 000 jetons par seconde, ce qui, sur cinq minutes, correspond à une taille de modèle idéale de 0,84 million : le modèle Chinchilla prédit donc que c'est excessif, ce qui correspond à mon expérience. Pour un modèle de 1 million, je peux entraîner 100 000 jetons par seconde, ce qui correspond à une taille de modèle idéale de 1,5 million, beaucoup plus proche de l'optimum, et mes résultats sont en effet meilleurs. Je n'ai pas fait le calcul pour déterminer la taille optimale exacte à laquelle mes jetons par seconde correspondent à la taille réelle du modèle, mais elle se situe clairement entre 1 million et 1,5 million de paramètres.

C'était cool de voir des lois de mise à l'échelle réelles bien connues s'appliquer à ce défi !

Conclusion

J'ai pris beaucoup de plaisir à faire cela. J'ai beaucoup appris sur l'entraînement de très petits modèles pendant de très courtes périodes. J'ai essayé de comprendre chaque ligne de code que j'ai écrite (et j'ai déjà écrit des transformateurs à partir de zéro), mais je n'aurais certainement pas essayé des choses comme les modèles de diffusion sans l'aide de LLM.

Je ne pense pas que ce défi soit particulièrement utile pour former des modèles puissants en général. La plupart des comportements intéressants se produisent après les cinq premières minutes de formation. Mais j'ai tout de même été agréablement surpris de la facilité avec laquelle j'ai pu former un modèle de narration globalement cohérent. À mesure que les architectures (et les GPU des ordinateurs portables) s'améliorent, je me demande quel type de modèles nous pourrons finalement former en cinq minutes.

Source : "What's the strongest AI model you can train on a laptop in five minutes?"

Et vous ?

Pensez-vous que cette expérience est crédible ou pertinente ?
Quel est votre avis sur le sujet ?

Voir aussi :

Pourquoi DeepSeek est bon marché à grande échelle mais coûteux à exploiter localement, par Sean Goedecke

La "distillation pas-à-pas", un nouveau paradigme d'apprentissage pour surpasser les performances des grands LLM, avec moins de données d'entraînement et des modèles d'IA de plus petite taille

Comment Meta entraîne les grands modèles de langage en permettant aux chercheurs d'utiliser PyTorch et d'autres nouveaux développements open-source