Comment TGS et AWS ont ramené l'entraînement d'un modèle d'IA sismique de six mois à cinq jours

TGS, avec AWS, affiche un résultat rare pour de grands projets d'IA : l'entraînement d'un modèle de fondation sismique, qui prenait auparavant environ six mois, a été ramené à cinq jours. Parallèlement, l'équipe a agrandi la fenêtre de contexte, c'est-à-dire le volume de données géologiques tridimensionnelles que le modèle peut analyser en une seule passe. Pour les entreprises d'exploration du sous-sol, cela signifie des itérations plus rapides et une image plus complète des structures souterraines.

TGS est un fournisseur de données géoscientifiques pour le secteur de l'énergie. L'entreprise utilise des modèles de fondation sismiques pour analyser des tableaux 3D complexes et repérer les structures géologiques importantes pour l'exploration et la production. L'architecture de base du modèle repose sur Vision Transformer et sur le schéma d'entraînement Masked AutoEncoder. Le principal problème ne tient pas seulement à la taille du modèle, mais à la nature des données : les volumes sismiques comptent des milliards de points, sont stockés dans des formats spécialisés et doivent alimenter les GPU en continu, sans temps morts.

Avec l'AWS Generative AI Innovation Center, l'entreprise a transféré l'entraînement sur Amazon SageMaker HyperPod et monté un cluster de 16 nœuds Amazon EC2 P5. Chaque nœud comprenait 8 GPU NVIDIA H200 dotés de 141 Go de mémoire HBM3e, 192 vCPU, 2 To de RAM et un réseau EFAv3 à 3200 Gbit/s. Au total, cela fait 128 GPU. D'après AWS, cette configuration a produit une mise à l'échelle quasi linéaire : l'efficacité du calcul parallèle s'est maintenue autour de 90–95 % en passant d'un nœud à 16.

Le pipeline de données a fait l'objet d'une attention particulière. Plutôt que le schéma classique avec Amazon FSx for Lustre, l'équipe a choisi le streaming direct depuis Amazon S3. La raison est simple : quand le cluster grandit, S3 permet à chaque nœud d'ajouter sa propre bande passante, alors qu'un système de fichiers partagé devient plus vite un goulot d'étranglement.

Le jeu de données d'entraînement de TGS est stocké au format MDIO, développé par l'entreprise sur la base de Zarr et optimisé pour les grandes données scientifiques dans le cloud. Après réglage du chargement multithread et du prefetching, chaque nœud recevait de façon stable 4–5 Go/s, et l'ensemble du cluster 64–80 Go/s. Cela a aussi réduit de plus de 90 % les coûts de l'infrastructure de stockage.

Pour l'entraînement distribué, l'équipe a comparé plusieurs approches : DeepSpeed ZeRO-2, ZeRO-3 et FSDP2. ZeRO-2 a offert le meilleur compromis entre vitesse et économie de mémoire : 1974 échantillons par seconde, contre 1833 pour FSDP2 et 869 pour ZeRO-3. C'est important, car dans ce type de tâches, l'économie de mémoire maximale ne donne pas toujours le meilleur résultat final : si les échanges entre GPU deviennent trop coûteux, le débit s'effondre. Dans ce projet, l'équipe n'a pas parié sur l'option théoriquement la plus compacte, mais sur la configuration qui exécute réellement l'entraînement le plus vite en conditions de production.

La deuxième avancée majeure concerne l'élargissement de la fenêtre de contexte. Pour les modèles sismiques, cela influe directement sur la qualité de l'analyse : plus le volume de roche vu par le modèle en une passe est grand, mieux il saisit à la fois les détails locaux, comme de petites fractures, et les grands motifs, par exemple les systèmes de failles à l'échelle d'un bassin entier. Grâce au context parallelism et à l'adaptation de ring attention à l'architecture ViT pour les données 3D, la taille maximale d'entrée est passée de 640 x 640 x 1024 à 1536 x 1536 x 2048 voxels.

La longueur du contexte a augmenté de 102 400 à 1,17 million de tokens, et le volume analysé d'environ 4,5 fois.

L'intérêt pratique du cas est que TGS peut désormais mettre à jour ses modèles quasiment chaque semaine, et non plus tous les six mois, intégrer plus vite de nouvelles données géologiques et offrir aux clients un contexte d'analyse plus large. Pour AWS, c'est un exemple de plus montrant que les modèles de fondation spécialisés en science et dans l'industrie dépendent non seulement de l'architecture du modèle, mais aussi d'une bonne organisation des données, du réseau et de l'entraînement distribué. Pour le marché dans son ensemble, le signal est clair : l'IA très spécialisée commence à s'imposer là où dominaient jusqu'ici de longs cycles de calcul et une infrastructure trop coûteuse.