DiLoCo Découplé : Google révolutionne l'entraînement AI mondial

DiLoCo Découplé : une avancée dans l'entraînement AI

L'entraînement des modèles d'intelligence artificielle de pointe repose traditionnellement sur des systèmes où les puces doivent fonctionner en parfaite synchronisation. Cette méthode, bien qu'efficace pour les modèles actuels, devient de plus en plus complexe à maintenir à mesure que l'échelle des modèles augmente. Google propose une solution innovante avec DiLoCo Découplé (Distributed Low-Communication), qui divise les sessions d'entraînement en "îles" de calcul découplées, permettant une circulation asynchrone des données entre elles.

Cette nouvelle architecture offre une résilience accrue en isolant les perturbations locales, ce qui permet aux autres parties du système de continuer à apprendre efficacement. DiLoCo Découplé surmonte ainsi les limitations des méthodes distribuées précédentes, telles que le Data-Parallel, qui étaient freinées par les délais de communication à l'échelle mondiale.

Vers un entraînement asynchrone tolérant aux pannes

DiLoCo Découplé s'appuie sur deux avancées antérieures : Pathways, qui a introduit un système AI distribué basé sur un flux de données asynchrone, et DiLoCo, qui a considérablement réduit la bande passante requise entre les centres de données distribués. En combinant ces innovations, DiLoCo Découplé permet un entraînement asynchrone à travers des unités d'apprentissage séparées, garantissant qu'une panne de puce dans une zone n'interrompt pas le progrès des autres.

Lors des tests, une méthode appelée "ingénierie du chaos" a été utilisée pour introduire des pannes matérielles artificielles pendant les sessions d'entraînement. DiLoCo Découplé a poursuivi le processus d'entraînement après la perte d'unités d'apprentissage entières, puis les a réintégrées sans problème lorsqu'elles sont revenues en ligne. Cette capacité auto-réparatrice assure une continuité de l'apprentissage, même en cas de défaillances matérielles.

Performances et résultats des tests

Les essais réalisés avec le modèle Gemma 4 ont confirmé l'efficacité de DiLoCo Découplé. Le système maintient une plus grande disponibilité des clusters d'apprentissage que les méthodes d'entraînement plus traditionnelles, tout en délivrant finalement le même niveau de performance en apprentissage automatique. DiLoCo Découplé nécessite des ordres de grandeur moins de bande passante que les méthodes conventionnelles, le rendant extrêmement efficace.

Avec des niveaux croissants de pannes matérielles, DiLoCo Découplé continue de fournir un "goodput" élevé, ou d'entraînement utile, tandis que celui d'autres approches chute. Notamment, le système a réussi à entraîner un modèle de 12 milliards de paramètres à travers quatre régions distinctes des États-Unis en utilisant 2-5 Gbps de réseau à large bande. Ce résultat a été atteint plus de 20 fois plus rapidement que les méthodes de synchronisation conventionnelles.

Évolution de l'infrastructure d'entraînement AI

Chez Google, l'approche full-stack de l'entraînement AI, intégrant matériel, infrastructure logicielle et recherche, est illustrée par DiLoCo Découplé. Ce système permet d'exploiter des ressources de calcul inutilisées à travers le monde, transformant les capacités inutilisées en puissance de calcul utile.

DiLoCo Découplé permet également de mélanger différentes générations de matériel, comme les TPU v6e et TPU v5p, dans une même session d'entraînement. Cela prolonge la durée de vie du matériel existant et augmente la capacité totale de calcul disponible pour l'entraînement des modèles. Dans nos expériences, des puces de différentes générations fonctionnant à des vitesses différentes ont toujours égalé la performance des sessions d'entraînement à puce unique, garantissant que même le matériel plus ancien peut accélérer de manière significative l'entraînement AI.

De plus, comme les nouvelles générations de matériel n'arrivent pas partout en même temps, la possibilité de s'entraîner à travers les générations peut atténuer les goulets d'étranglement logistiques et de capacité récurrents. Alors que nous repoussons les frontières de l'infrastructure AI aujourd'hui, nous continuons à explorer des approches pour des systèmes résilients nécessaires pour débloquer la prochaine génération d'AI.

Contribution et soutien

Ce projet a été mené par une équipe de Google DeepMind et Google Research, avec des contributions clés de Arthur Douillard, Keith Rush, Yani Donchev, Zachary Charles, Ayush Dubey, Blake Woodworth, Ionel Gog, Josef Dean, Nova Fallen, Zachary Garrett. Le soutien opérationnel a été assuré par Nate Keating et Jenny Bishop, avec des conseils supplémentaires de Jeff Dean, Marc’Aurelio Ranzato, Raia Hadsell, Arthur Szlam, Edouard Yvinec, Henry Prior, Paul Barham, Michael Isard, Daniel Ramage, Brendan McMahan, Chase Hensel, et Zoltan Egyed.