Google’s Tunix system lifts the choke point that has kept large-scale agentic reinforcement learning (RL) from using TPUs efficiently. By separating the work of generating interaction data from the work of updating the policy, Tunix drives TPU utilization from single-digit percentages up to near-full capacity, cutting compute waste dramatically.

Le goulot d'étranglement de l'apprentissage par renforcement (RL) agentique

Le RL agentique diffère de l'entraînement plus familier des modèles de langage par « prochain jeton » (next-token). Un agent doit envoyer des appels API, exécuter du code ou progresser dans un environnement simulé, puis réagir au résultat. La boucle d'entraînement est donc synchrone : le modèle produit une action, l'environnement s'exécute, le résultat est renvoyé, et ce n'est qu'ensuite que le modèle reçoit une mise à jour du gradient. Lorsqu'une seule étape de l'environnement prend plusieurs secondes, le matériel TPU coûteux reste inactif, et l'utilisation rapportée peut descendre sous les 10 %. Cette inefficacité se traduit directement par des factures cloud plus élevées et des cycles de recherche plus lents.

L'architecture découplée de Tunix

Tunix s'attaque au problème en répartissant les deux étapes — la génération de trajectoires et l'optimisation de la politique — sur des pools de matériel distincts.

  • Les acteurs asynchrones s'exécutent sur des CPU ou GPU peu coûteux. Chaque acteur interagit continuellement avec l'environnement qui lui est assigné, enregistre les actions et les observations, et diffuse les trajectoires résultantes vers un stockage partagé.
  • Les apprenants continus occupent des Pods TPU dédiés. L'apprenant extrait des lots (batches) du tampon central et effectue des mises à jour de gradient sans attendre qu'un acteur termine un rollout.
  • Un tampon à haut débit se situe au milieu, agissant comme une zone de transit pour les trajectoires. Comme l'apprenant peut lire aussi vite que le tampon peut fournir des données, le TPU ne s'arrête jamais.

L'effet net est un pipeline d'entraînement où les TPU restent occupés presque tout le temps, poussant l'utilisation vers les 100 %.

Obstacles techniques et solutions apportées par Tunix

Épisodes de longueur variable et recompilation XLA

Le compilateur XLA de JAX optimise pour des formes de tenseurs fixes. Les tâches agentiques produisent cependant des séquences de longueurs différentes, ce qui déclencherait normalement des recompilations coûteuses. Tunix regroupe les séquences plus courtes et organise les épisodes de longueur similaire dans des compartiments (buckets), maintenant des formes stables suffisamment longtemps pour que XLA puisse réutiliser les noyaux (kernels) compilés. Le résultat est un débit constant sans la surcharge du compilateur qui, autrement, paralyserait les performances.

Mise à l'échelle de modèles massifs sur de nombreuses puces TPU

L'entraînement d'agents de plus de 70 milliards de paramètres nécessite de répartir les poids et les données sur plusieurs nœuds TPU. Tunix utilise la primitive ShardMap de JAX pour partitionner (shard) à la fois les paramètres du modèle et les activations, permettant à l'apprenant de conserver l'intégralité du modèle en mémoire tout en le nourrissant de données à grande vitesse. Cette stratégie de partitionnement permet d'entraîner des modèles qui étaient auparavant hors de portée pour un seul pod TPU.

Gradients obsolètes issus de pipelines découplés

Lorsque les acteurs sont en avance sur l'apprenant, les données qu'ils fournissent peuvent devenir « obsolètes » par rapport à la politique actuelle. Tunix atténue cette dérive grâce à deux mécanismes : l'échantillonnage préférentiel (importance sampling) repondère les échantillons plus anciens pour refléter leur pertinence, et un seuil d'obsolescence configurable qui rejette les trajectoires dépassant un âge prédéfini. Ensemble, ils maintiennent la stabilité de l'apprentissage même lorsque le pipeline fonctionne de manière asynchrone.

Points de vigilance pour les adoptants

  • Audit de la latence – L'avantage du découplage dépend du temps de réponse de l'environnement. Les équipes doivent mesurer la latence de bout en bout et s'assurer que les pools d'acteurs sont dimensionnés pour maintenir le tampon bien rempli.
  • Conception du pool de travailleurs – Des CPU ou GPU peu coûteux peuvent héberger de nombreux acteurs, mais la surallocation peut provoquer des conflits sur le réseau ou le stockage. Un pool équilibré, correspondant au taux d'ingestion du tampon, est essentiel.
  • Robustesse du tampon – Le stockage central doit gérer des taux d'écriture et de lecture élevés sans devenir un nouveau goulot d'étranglement. Le choix d'un système de stockage présentant une faible latence de queue (tail latency) et une bande passante suffisante est un élément non négociable de l'architecture.

Inconvénients potentiels

L'architecture divisée introduit plus d'éléments mobiles : des flottes de matériel distinctes, un tampon persistant et une logique de coordination pour appliquer les limites d'obsolescence.

À retenir

Tunix démontre que le coût dominant du RL agentique n'est pas le modèle lui-même, mais le temps d'inactivité causé par les boucles d'interaction synchrones. En déchargeant le travail de rollout sur du matériel peu coûteux et en alimentant un pod TPU en apprentissage continu à partir d'un tampon à haut débit, Google a transformé un problème d'utilisation inférieure à 10 % en un flux de travail à capacité quasi maximale.