Como TGS e AWS Reduziram o Treinamento de Modelo de IA Sísmica de Seis Meses para Cinco Dias

TGS em conjunto com AWS apresentou um resultado raro para grandes projetos de IA: o treinamento de um modelo foundation sísmico, que anteriormente levava aproximadamente seis meses, foi reduzido para cinco dias. Ao mesmo tempo, a equipe aumentou o tamanho da janela de contexto, ou seja, o volume de dados geológicos tridimensionais que o modelo pode analisar em uma única passagem. Para empresas que trabalham com exploração de subsuperfície, isso significa iterações mais rápidas e uma visão mais completa das estruturas subterrâneas.

TGS é fornecedora de dados geocientíficos para o setor energético. A empresa utiliza modelos foundation sísmicos para analisar matrizes 3D complexas para encontrar estruturas geológicas importantes para exploração e produção. A arquitetura básica do modelo é construída em Vision Transformer e esquema de treinamento Masked AutoEncoder.

O principal problema aqui não é apenas o tamanho do modelo em si, mas também a natureza dos dados: volumes sísmicos consistem em bilhões de pontos, são armazenados em formatos especializados e requerem entrega contínua para GPU sem tempo de inatividade. Em conjunto com o AWS Generative AI Innovation Center, a empresa migrou o treinamento para Amazon SageMaker HyperPod e montou um cluster de 16 nós Amazon EC2 P5. Cada nó incluía 8 GPUs NVIDIA H200 com 141 GB de memória HBM3e, 192 vCPU, 2 TB de RAM e rede EFAv3 com velocidade de 3200 Gbps.

No total, são 128 GPUs. De acordo com AWS, essa configuração proporcionou escalabilidade quase linear: a eficiência da operação paralela ao escalar de um nó para 16 manteve-se em aproximadamente 90–95%. Um foco separado foi no pipeline de dados.

Em vez do esquema clássico com Amazon FSx for Lustre, a equipe escolheu streaming direto do Amazon S3. O motivo é simples: conforme o cluster cresce, S3 permite que cada nó adicione sua própria largura de banda, enquanto um sistema de arquivos compartilhado fica um gargalo mais rapidamente. O conjunto de dados de treinamento da TGS é armazenado no formato MDIO desenvolvido pela empresa com base em Zarr, otimizado para grandes dados científicos na nuvem.

Após configurar carregamento multi-thread e prefetching, cada nó obtinha consistentemente 4–5 GB/s, e todo o cluster — 64–80 GB/s. Além disso, isso reduziu os custos de infraestrutura de armazenamento em mais de 90%. Para treinamento distribuído, a equipe comparou várias abordagens: DeepSpeed ZeRO-2, ZeRO-3 e FSDP2.

ZeRO-2 mostrou o melhor equilíbrio entre velocidade e economia de memória: 1.974 amostras por segundo em comparação com 1.833 para FSDP2 e 869 para ZeRO-3.

Isso é importante porque em tais tarefas, a máxima economia de memória nem sempre oferece o melhor resultado geral: se a comunicação de GPU fica muito cara, o throughput cai drasticamente. Neste projeto, o foco não foi na opção teoricamente mais densa, mas na configuração que realmente processa o treinamento mais rápido em um cenário de produção. A segunda conquista-chave diz respeito à expansão da janela de contexto.

Para modelos sísmicos, isso afeta diretamente a qualidade da análise: quanto maior o volume de rocha que o modelo vê em uma única passagem, melhor ele captura tanto detalhes locais, como pequenas fraturas, quanto padrões grandes — por exemplo, sistemas de falhas na escala de uma bacia inteira. Devido ao context parallelism e adaptação de ring attention para arquitetura Vision Transformer para dados 3D, o tamanho máximo de entrada aumentou de 640 x 640 x 1.024 para 1.

536 x 1.536 x 2.048 voxels.

O comprimento do contexto aumentou de 102.400 para 1,17 milhão de tokens, e o volume analisado — aproximadamente 4,5 vezes. O significado prático do caso é que TGS agora pode atualizar modelos não uma vez a cada seis meses, mas efetivamente semanalmente, conectar mais rapidamente novos dados geológicos e fornecer aos clientes um contexto de análise mais amplo.

Para AWS, este é outro exemplo demonstrativo de que modelos foundation especializados em ciência e indústria são limitados não apenas pela arquitetura do modelo, mas também pela organização adequada de dados, redes e treinamento distribuído. Para o mercado como um todo, o sinal é claro: IA especializada em nicho está começando a vencer onde ciclos computacionais longos e infraestrutura cara demais costumavam dominar.