Comment TGS et AWS Ont Réduit l'Entraînement du Modèle IA Sismique de Six Mois à Cinq Jours

TGS en collaboration avec AWS a démontré un résultat rare pour les grands projets d'IA : l'entraînement d'un modèle foundation sismique, qui prenait auparavant environ six mois, a été réduit à cinq jours. En même temps, l'équipe a augmenté la taille de 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 travaillant dans l'exploration en profondeur, 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 énergétique. L'entreprise utilise des modèles foundation sismiques pour analyser des matrices 3D complexes afin de trouver les structures géologiques importantes pour l'exploration et la production. L'architecture de base du modèle repose sur Vision Transformer et le schéma d'entraînement Masked AutoEncoder.

Le principal problème ici n'est pas seulement la taille du modèle lui-même, mais aussi la nature des données : les volumes sismiques sont composés de milliards de points, sont stockés dans des formats spécialisés et nécessitent une livraison continue au GPU sans temps d'arrêt. En collaboration avec AWS Generative AI Innovation Center, l'entreprise a migré l'entraînement vers Amazon SageMaker HyperPod et a assemblé un cluster de 16 nœuds Amazon EC2 P5. Chaque nœud comprenait 8 GPU NVIDIA H200 avec 141 GB de mémoire HBM3e, 192 vCPU, 2 TB de RAM et un réseau EFAv3 avec une vitesse de 3200 Gbps.

Au total, il y a 128 GPU. Selon AWS, cette configuration a fourni une mise à l'échelle quasi linéaire : l'efficacité du fonctionnement parallèle lors du passage d'un nœud à 16 s'est maintenue à environ 90–95%. Un accent particulier a été mis sur le pipeline de données.

Au lieu du schéma classique avec Amazon FSx for Lustre, l'équipe a choisi la diffusion en continu directe depuis Amazon S3. La raison est simple : à mesure que le cluster se développe, S3 permet à chaque nœud d'ajouter sa propre bande passante, alors qu'un système de fichiers partagé devient rapidement un goulot d'étranglement. L'ensemble de données d'entraînement de TGS est stocké au format MDIO développé par l'entreprise basé sur Zarr, optimisé pour les grandes données scientifiques dans le cloud.

Après la configuration du chargement multithread et du prefetching, chaque nœud a atteint de manière stable 4–5 GB/s, et l'ensemble du cluster — 64–80 GB/s. De plus, cela a réduit les coûts d'infrastructure de stockage de plus de 90%. Pour l'entraînement distribué, l'équipe a comparé plusieurs approches : DeepSpeed ZeRO-2, ZeRO-3 et FSDP2.

ZeRO-2 a montré le meilleur équilibre entre vitesse et économies de mémoire : 1 974 échantillons par seconde par rapport à 1 833 pour FSDP2 et 869 pour ZeRO-3. C'est important car dans de telles tâches, l'économie maximale de mémoire ne donne pas toujours le meilleur résultat global : si la communication entre GPU devient trop coûteuse, le débit chute fortement. Dans ce projet, l'accent a été mis non pas sur l'option théoriquement la plus compacte, mais sur la configuration qui traite réellement l'entraînement plus rapidement dans un scénario de production.

Le deuxième succès clé concerne l'expansion de la fenêtre de contexte. Pour les modèles sismiques, cela affecte directement la qualité de l'analyse : plus le volume de roche que le modèle voit en une seule passe est grand, mieux il capte à la fois les détails locaux, comme les petites fractures, et les grands motifs — par exemple, les systèmes de failles à l'échelle de tout un bassin. Grâce au context parallelism et à l'adaptation de ring attention pour l'architecture Vision Transformer pour les données 3D, la taille maximale d'entrée a augmenté de 640 x 640 x 1 024 à 1 536 x 1 536 x 2 048 voxels.

La longueur du contexte a augmenté de 102 400 à 1,17 million de tokens, et le volume analysé — environ 4,5 fois. Le sens pratique du cas est que TGS peut maintenant mettre à jour les modèles non pas une fois tous les six mois, mais effectivement chaque semaine, connecter plus rapidement les nouvelles données géologiques et fournir aux clients un contexte d'analyse plus large. Pour AWS, c'est un autre exemple démonstratif que les modèles foundation spécialisés dans la science et l'industrie sont limités non seulement par l'architecture du modèle, mais aussi par une organisation appropriée des données, des réseaux et de l'entraînement distribué.

Pour le marché dans son ensemble, le signal est clair : l'IA étroitement spécialisée commence à gagner là où auparavant dominaient les longs cycles de calcul et l'infrastructure trop coûteuse.