Como TGS e AWS reduziram o treinamento de um modelo de IA sísmico de seis meses para cinco dias

A TGS, junto com a AWS, apresentou um resultado raro em grandes projetos de IA: o treinamento de um modelo fundacional sísmico, que antes levava cerca de seis meses, caiu para cinco dias. Ao mesmo tempo, a equipe ampliou a janela de contexto, ou seja, o volume de dados geológicos tridimensionais que o modelo consegue analisar de uma só vez. Para empresas que atuam na exploração do subsolo, isso significa iterações mais rápidas e uma visão mais completa das estruturas subterrâneas.

A TGS fornece dados geocientíficos para o setor de energia. A empresa usa modelos fundacionais sísmicos para analisar matrizes 3D complexas e encontrar estruturas geológicas importantes para exploração e produção. A arquitetura base do modelo é construída sobre Vision Transformer e o esquema de treinamento Masked AutoEncoder. O principal problema não é apenas o tamanho do modelo, mas a natureza dos dados: volumes sísmicos têm bilhões de pontos, ficam em formatos especializados e precisam alimentar as GPUs continuamente, sem ociosidade.

Em conjunto com o AWS Generative AI Innovation Center, a empresa levou o treinamento para o Amazon SageMaker HyperPod e montou um cluster de 16 nós Amazon EC2 P5. Cada nó tinha 8 GPUs NVIDIA H200 com 141 GB de memória HBM3e, 192 vCPUs, 2 TB de RAM e rede EFAv3 de 3200 Gbit/s. No total, são 128 GPUs. Segundo a AWS, essa configuração entregou escalabilidade quase linear: a eficiência do processamento paralelo ficou em cerca de 90–95% ao passar de um nó para 16.

O pipeline de dados recebeu atenção à parte. Em vez do esquema clássico com Amazon FSx for Lustre, a equipe optou por streaming direto do Amazon S3. O motivo é simples: conforme o cluster cresce, o S3 permite que cada nó adicione sua própria vazão, enquanto um sistema de arquivos compartilhado vira gargalo mais rápido. O dataset de treinamento da TGS fica armazenado no MDIO, formato desenvolvido pela empresa sobre o Zarr e otimizado para grandes dados científicos na nuvem. Depois de ajustar o carregamento multithread e o prefetching, cada nó passou a receber de forma estável 4–5 GB/s, e o cluster inteiro, 64–80 GB/s. Além disso, isso reduziu em mais de 90% os custos de infraestrutura de armazenamento.

Para o treinamento distribuído, a equipe comparou várias abordagens: DeepSpeed ZeRO-2, ZeRO-3 e FSDP2. O ZeRO-2 mostrou o melhor equilíbrio entre velocidade e economia de memória: 1974 amostras por segundo, contra 1833 do FSDP2 e 869 do ZeRO-3. Isso importa porque, nesse tipo de tarefa, a máxima economia de memória nem sempre gera o melhor resultado final: se a troca de dados entre GPUs fica cara demais, o throughput despenca. Neste projeto, a aposta não foi na opção teoricamente mais compacta, mas na configuração que de fato executa o treinamento mais rápido em um cenário de produção.

A segunda conquista importante envolve a ampliação da janela de contexto. Em modelos sísmicos, isso afeta diretamente a qualidade da análise: quanto maior o volume de rocha que o modelo enxerga de uma vez, melhor ele capta tanto detalhes locais, como pequenas fraturas, quanto grandes padrões, como sistemas de falhas na escala de uma bacia inteira. Com context parallelism e a adaptação do ring attention à arquitetura ViT para dados 3D, o tamanho máximo de entrada subiu de 640 x 640 x 1024 para 1536 x 1536 x 2048 voxels. O comprimento do contexto passou de 102 400 para 1,17 milhão de tokens, e o volume analisado cresceu cerca de 4,5 vezes.

O ganho prático é que a TGS agora pode atualizar seus modelos na prática toda semana, e não a cada seis meses, incorporar novos dados geológicos mais rápido e dar aos clientes um contexto de análise mais amplo. Para a AWS, é mais um exemplo de que modelos fundacionais especializados em ciência e indústria dependem não só da arquitetura do modelo, mas também da organização correta de dados, rede e treinamento distribuído. Para o mercado em geral, o sinal é claro: a IA altamente especializada começa a vencer onde antes dominavam longos ciclos de computação e infraestrutura cara demais.