L’apprentissage par renforcement sur les TPU Google Cloud

L’apprentissage par renforcement (reinforcement learning ou RL) a de nombreuses applications : robot, voiture autonome, chatbot… Cependant, son adoption pour les applications du monde réel est entravée par des défis d’ingénierie majeurs. Les pods TPU de Google Cloud offrent une solution innovante pour surmonter ces obstacles.

Les premières victoires notables de l’RL ont eu lieu dans l’univers des jeux vidéo, où des systèmes basés sur l’apprentissage par renforcement ont triomphé de champions mondiaux dans des jeux tels que Go, StarCraft 2 et Dota2. Cependant, malgré ces succès, l’application généralisée du RL reste un défi.

Les TPU : des accélérateurs transformateurs pour l’IA

Les TPU (Tensor Processing Unit) introduits en 2016 par Google Cloud, émergent comme une réponse aux défis techniques de l’RL. Ces unités de traitement tensorielles sont spécialement conçues pour les applications de calcul haute performance à grande échelle, offrant une architecture optimisée pour accélérer les opérations tensorielles essentielles aux réseaux neuronaux modernes. La quatrième génération de ces TPU, annoncée en 2021, a été organisée en pods, des grappes de puces TPU interconnectées. Les pods TPU v4 peuvent ainsi combiner jusqu’à 4 096 puces, offrant une performance de pointe de 1,1 ExaFLOPS.

L’architecture Sebulba pour la mise à l’échelle du RL

L’architecture Sebulba, introduite en 2021 dans l’article Podracer de Google DeepMind. est une réponse à la mise à l’échelle des systèmes RL. En co-localisant les acteurs et les apprenants sur les mêmes TPU, Sebulba minimise les goulots d’étranglement associés au transfert de données, tout en capitalisant sur les capacités de traitement parallèle des TPU. La figure  ci-dessous illustre la configuration des cœurs d’acteur et d’apprenant sur un TPU, qui peut être répliquée sur plusieurs TPU pour augmenter le débit et la performance des agents RL.

Architecture de Sebulba illustrant le placement des cœurs d’apprentissage (jaune) et des cœurs d’acteur (rouge) sur le TPU, ainsi que les environnements sur les processeurs hôtes (vert). Les environnements envoient leurs observations et leurs récompenses aux cœurs d’acteurs qui envoient ensuite en réponse les prochaines actions à entreprendre. Ce faisant, les acteurs construisent un tampon de trajectoires. Simultanément, l’apprenant extrait des lots de trajectoires de la mémoire tampon des acteurs et effectue des mises à jour d’apprentissage, en poussant de nouveaux paramètres de modèle vers les appareils des acteurs. Crédit Google Cloud

Un cas d’utilisation concret : le routage des cartes de circuits imprimés

InstaDeep propose DeepPCB, un produit basé sur l’IA permettant aux ingénieurs électriciens d’acheminer leurs  cartes de circuits imprimés de manière optimale sans intervention humaine. Dans ce cadre, son équipe produit a créé un moteur de simulation rapide pour entraîner efficacement des agents de RL. En connectant l’architecture Sebulba au simulateur DeepPCB, elle a obtenu une accélération de 15 à 235 fois par rapport à l’ancienne architecture distribuée.

Lors de la conférence Next ’23, Google Cloud a présenté la 5ème  génération de ses TPU : les v5e, indiquant sur son blog :

“Cloud TPU v5 est spécialement conçu pour apporter la rentabilité et les performances requises pour l’entraînement et l’inférence à moyenne et grande échelle. TPU v5e offre des performances d’entraînement jusqu’à 2 fois supérieures par dollar et jusqu’à 2,5 fois plus performantes par dollar pour les LLM et les modèles d’IA gen par rapport à Cloud TPU v4. Avec un coût inférieur de moitié à celui de TPU v4, TPU v5e permet à un plus grand nombre d’organisations d’entraîner et de déployer des modèles d’IA plus grands et plus complexes”.

Références de l’article : Blog Google Cloud, rédigé par Armand Picard, Ingénieur logiciel, InstaDeep et Donal Byrne, Ingénieur de recherche senior, InstaDeep.

Recevez gratuitement l'actualité de l'intelligence artificielle

Suivez la Newsletter de référence sur l'intelligence artificielle (+ de 18 000 membres), quotidienne et 100% gratuite.


Tout comme vous, nous n'apprécions pas le spam. Vos coordonnées ne seront transmises à aucun tiers.
Partager l'article
1 an d'abonnement au magazine de référence sur l'IA pour 27,60€ Boutique ActuIA Recherche, Business, Impact : Restez à la page en matière d'intelligence artificielle.
intelligence artificielle
À PROPOS DE NOUS
Le portail francophone consacré à l'intelligence artificielle et à la datascience, à destination des chercheurs, étudiants, professionnels et passionnés.