Collecte et entraînement continu : les limites du MacBook

  • llm
  • entrainement

J’ai été limité par les capacités de mon MacBook. J’ai donc dû choisir entre PyTorch, qui est portable, et MLX, qui est rapide. PyTorch a gagné la partie car il permet d’entraîner des modèles sur différentes machines sans avoir à modifier le code. La portabilité était cruciale pour mon projet.

Le MacBook sert de laboratoire pour le projet Thadeus. Il permet de mettre au point les modèles, de déboguer et d’ajuster les hyperparamètres. Le choix de PyTorch facilite l’utilisation du même code sur le MacBook et sur le GPU H100, évitant ainsi les divergences numériques.

L’architecture retenue

Pour répondre aux besoins de collecte et d’entraînement continu sur un MacBook, j’ai conçu une architecture en plusieurs étages. Chaque composant de cette chaîne est indépendant et peut être rejoué sans refaire les précédents. Cela permet de comparer deux variantes d’un étage, toutes choses égales par ailleurs.

Chaque composant s’enregistre sous un nom et peut être sélectionné depuis un fichier TOML. Ce système est très flexible : passer d’une configuration à une autre ne prend qu’une ligne, sans refactoring.

Les tokens produits servent à entraîner le modèle. Chaque étape de cette chaîne est enregistrée sur disque, ce qui permet de réutiliser les données intermédiaires sans refaire les étapes précédentes. Les sources sont mélangées pour former des shards, qui alimentent le tokenizer.

Voici un extrait du script de collecte quotidienne, qui élargit le corpus pendant la journée sans GPU.

Ce que ça change en pratique

Le modèle de 188 M paramètres a été entraîné sur un MacBook, avec une collecte de corpus étalée sur 110 heures. Cette période a permis de voir 3,74 milliards de tokens, une base solide pour l’entraînement. Les résultats montrent que la densité d’information par FLOP est cruciale : un corpus technique spécialisé bat un corpus web généraliste, même si le modèle est plus petit. Le tokenizer sur mesure a permis d’économiser 27,2 % de tokens par rapport à GPT-2. La méthode muP a également été validée : le taux d’apprentissage optimal se transfère entre différentes largeurs de modèle, assurant une utilisation efficace des crédits H100. Cependant, Muon n’a pas été confirmé comme attendu : il perd 2 % sur notre banc. Ces gains sont mesurés et assumés, avec des erreurs consignées plutôt qu’effacées.

La collecte s’est adaptée pour suivre la consommation réelle de l’entraînement. Cela évite d’engorger le disque dur et assure une gestion plus efficace des ressources. La collecte quotidienne est désormais calibrée pour ne pas empiéter sur la fenêtre d’entraînement nocturne. La graine de chaque source dérive maintenant du jour, ce qui change réellement la récolte et évite les recouvrements. La validation est prélevée en blocs répartis sur tout le corpus, ce qui assure une meilleure représentativité et évite les biais de distribution.

L’entraînement continu est géré par une configuration qui permet de reprendre exactement où l’on s’était arrêté. Cela permet d’entraîner par nuits sur le Mac et de finir au cloud sans perdre de temps. Les checkpoints sont nommés par l’identité du run, garantissant que la reprise se fait toujours au bon endroit. J’utilise le planificateur WSD, qui maintient un taux d’apprentissage constant pendant une période prolongée, puis déclenche la décroissance.

Les tradeoffs, assumés

J’ai d’abord envisagé d’entraîner le modèle sur un H100. Pour des raisons de coût, j’ai abandonné cette idée et opté pour un plan hybride avec entraînement sur le Mac la nuit. Ce choix a été dicté par les contraintes budgétaires et les ressources disponibles. Le MacBook, utilisé comme laboratoire, permet de régler les hyperparamètres et de valider la chaîne d’entraînement sans coûts supplémentaires. Le H100, quant à lui, est utilisé pour le run final avec des hyperparamètres déjà transférés. Cette approche hybride maximise l’utilisation des ressources disponibles tout en minimisant les coûts. Le MacBook sert à déboguer, régler les hyperparamètres et valider la chaîne d’entraînement. Le H100 est réservé pour le run final, une fois les paramètres optimisés.

Un autre choix crucial a été de privilégier PyTorch sur MLX. PyTorch permet une portabilité entre le MacBook et le H100, évitant ainsi les divergences numériques. MLX, bien que plus rapide sur Apple Silicon en fp32, ne tourne pas sur H100 et n’offre aucun avantage significatif en bf16, ce qui a scellé son sort. Le benchmark a montré que PyTorch était aussi performant en bf16, confirmant ainsi le choix de la portabilité. Ce choix a été crucial pour éviter les erreurs numériques silencieuses qui peuvent survenir lors du passage d’une implémentation à une autre.

Le corpus a été conçu pour être dense et spécialisé. Il se compose de 55 % de français, 25 % de technique en anglais, 15 % de code et 5 % de notes personnelles. Ce choix a été fait pour maximiser la densité d’information par FLOP, en privilégiant un corpus technique spécialisé plutôt qu’un corpus web généraliste. Le modèle a été conçu pour être spécialisé sur un domaine précis, ce qui permet d’obtenir une meilleure performance avec moins de paramètres.

Quand ça vaut le coup, et quand non

L’approche suivie ici est utile pour quelqu’un qui souhaite apprendre. Le projet Thadeus est un laboratoire d’expérimentation, pas une solution de production.

J’ai envisagé d’entraîner le modèle sur un GPU H100.

Si j’avais à refaire ce projet, je le ferais sans hésitation. Cela montre qu’il est utile pour l’apprentissage, mais pas nécessairement pour une mise en production.

Ce que je changerais si je recommençais

Si je devais recommencer le projet Thadeus, je le ferais sans hésitation. Cela montre qu’il est utile pour l’apprentissage.

Ce projet a été conçu comme un laboratoire d’expérimentation. Il ne vise pas à être une solution de production.

La stack

L’ensemble du projet repose sur Python 3.12 et PyTorch, qui permet d’entraîner des modèles sur différentes machines sans modifier le code.

Les configurations sont gérées par des fichiers TOML, qui permettent de définir les paramètres du modèle et de passer d’une configuration à une autre sans refactoring. Les chemins en dur sont absents, les variables d’environnement renseignées dans un fichier .env sont utilisées.

Les scripts de collecte, d’entraînement et d’évaluation sont écrits en Python et utilisent des outils comme uv, pytest et ruff.

Une question sur ce sujet ?

Si votre cas ressemble à celui décrit ici et que vous voulez un avis avant de vous lancer, écrivez-moi. Une réponse courte et honnête ne coûte rien.