Le calcul en 4 bits dès le préentraînement d’Inkling 2 Small

Les couches d’experts d’Inkling 2 Small combinent calcul en 4 et 8 bits. Thinking Machines annonce des gains de vitesse sur GB300.

Du premier pas de préentraînement au dernier épisode d’apprentissage par renforcement, une même recette de calcul en basse précision accompagne Inkling 2 Small. Dans un billet signé Jack Cook, Thinking Machines détaille l’emploi du format NVFP4 pour le passage avant de couches d’experts : l’étape où le modèle calcule ses prédictions. Le passage arrière, qui calcule les corrections nécessaires à l’apprentissage, utilise le format MXFP8.

La formule « entraînement en 4 bits de bout en bout » désigne donc une continuité au fil des étapes, avec plusieurs précisions numériques. Le dernier sixième des couches reste calculé en MXFP8, tandis que des activations sont conservées en BF16 et des poids de référence en FP32. Les premiers essais appliquant NVFP4 également au passage arrière avaient, selon le laboratoire, trop dégradé l’efficacité de l’apprentissage.

Sur ses plus grands entraînements, le laboratoire mesure un gain de vitesse de 8 % au préentraînement et de 20 % en apprentissage par renforcement, par rapport à une configuration MXFP8. Ces chiffres concernent la recette complète, avec ses adaptations de calcul et de communication. Ils doivent être distingués des accélérations théoriques des multiplications matricielles offertes par le matériel.

Les 68 heures mises en avant par Mira Murati correspondent à un autre périmètre. L’annonce du modèle précise que cette durée couvre la phase initiale H1, portant sur 8 400 milliards de tokens de langue naturelle, de code et de données d’images, sur des racks GB300 NVL72. D’autres phases ont suivi, avant l’apprentissage par renforcement. Le modèle compte 594 milliards de paramètres au total, dont 17 milliards actifs. Le passage de plusieurs semaines à 68 heures est attribué à un ensemble d’améliorations, sans que le billet indique le nombre total de GPU mobilisés.

Réduire le nombre de bits diminue la place occupée par les données et accélère certaines opérations, au prix d’arrondis plus importants. Le format NVFP4 associe des valeurs sur quatre bits à un facteur d’échelle partagé par seize valeurs, soit 4,5 bits par valeur en comptant ce facteur. Le tableau publié annonce un débit maximal de calcul jusqu’à six fois supérieur à BF16 sur GB300. Ce rapport local ne se traduit pas mécaniquement par un entraînement complet six fois plus rapide.

La difficulté apparaît particulièrement lors du passage arrière. Les matrices utilisées à l’aller doivent être transposées, ce qui regroupe autrement les valeurs partageant un même facteur d’échelle. Une nouvelle quantification peut alors produire des valeurs différentes et fausser le calcul des gradients. Un arrondi qui favorise systématiquement une direction crée un biais susceptible de s’accumuler pendant l’entraînement.

Pour limiter cet effet, l’équipe déquantifie les valeurs utilisées à l’aller, les transpose avec une précision supérieure, puis les quantifie à nouveau en MXFP8 avec un arrondi stochastique. Cet arrondi choisit probabilistiquement la valeur représentable supérieure ou inférieure, en fonction de leur distance. Le but est de préserver en moyenne la valeur utilisée lors du passage avant, plutôt que d’exiger une représentation identique à chaque opération. Cela ne restaure pas toute l’information perdue par la première quantification.

Une fois les calculs accélérés, les échanges entre GPU deviennent un autre facteur limitant. Dans une architecture à mélange d’experts, les données sont envoyées vers les processeurs hébergeant les experts sélectionnés, puis leurs résultats sont rassemblés. Le laboratoire expédie les entrées en NVFP4. Pour le retour, une transformation de Hadamard répartit les valeurs extrêmes avant une quantification dans le format S2F6, avec des facteurs d’échelle. Le volume transmis passe alors de seize à 8,25 bits par valeur.

Des programmes de calcul GPU, ou kernels, développés en interne sous le nom BBBK combinent plusieurs de ces opérations pour limiter les attentes entre calcul, conversion et transfert. L’équipe décrit une fusion de la transformation et de la quantification avec la fin des multiplications matricielles. Le résultat dépend ainsi de l’organisation du travail sur les GB300 autant que du choix d’un format numérique.

Le compromis conserve un coût d’apprentissage. Au préentraînement, le laboratoire rapporte environ 2 % de perte d’efficacité de calcul : un entraînement MXFP8 aurait atteint le même niveau de perte prédictive avec environ 2 % d’opérations en moins. En apprentissage par renforcement, l’équipe ne relève pas de différence discernable de récompense dans ses comparaisons. Le calcul du passage arrière en FP4 reste un chantier de recherche, distinct de la recette utilisée pour Inkling 2 Small.