AWS SageMaker HyperPod acelera treinamento de modelos sísmicos com redução de 6 meses para 5 dias

A TGS, fornecedora de dados geocientíficos para o setor de energia, em parceria com o AWS Generative AI Innovation Center (GenAIIC), lançou uma solução inovadora para acelerar o treinamento de seus modelos sísmicos baseados em Vision Transformer (ViT). Utilizando o Amazon SageMaker HyperPod, a equipe conseguiu reduzir o tempo de treinamento de seis meses para apenas cinco dias, ao mesmo tempo em que expandiu a capacidade analítica dos modelos para volumes sísmicos maiores do que antes.

O que foi lançado e para quem serve

O destaque da solução é a aplicação do SageMaker HyperPod para treinamento distribuído resiliente e escalável, aliado a técnicas avançadas de paralelismo de contexto que ampliam a janela de análise dos modelos sísmicos. Essa inovação é especialmente útil para empresas do setor de energia que dependem da análise de grandes volumes de dados 3D sísmicos para identificar estruturas geológicas essenciais à exploração e produção.

Desafios enfrentados no treinamento de modelos sísmicos

Escala e complexidade dos dados: A TGS trabalha com grandes volumes de dados sísmicos 3D proprietários armazenados em formatos específicos, exigindo estratégias eficientes de streaming para evitar ociosidade das GPUs.
Eficiência no treinamento: O processamento volumétrico 3D demanda alta capacidade computacional, tornando fundamental a aceleração dos ciclos de treinamento para incorporar dados e aprimoramentos com maior frequência.
Ampliação da capacidade analítica: Expandir a janela de contexto permite analisar simultaneamente detalhes locais e padrões geológicos amplos, aumentando a qualidade das interpretações.

Como funciona a solução técnica

A arquitetura da solução integra o SageMaker HyperPod para orquestração de clusters de treinamento, configurados com 16 instâncias Amazon EC2 P5, cada uma com 8 GPUs NVIDIA H200, totalizando 128 GPUs com alta memória e rede de baixa latência.

O pipeline de dados foi otimizado para streaming direto do Amazon S3 usando o formato MDIO, desenvolvido pela TGS, que permite alta taxa de transferência (64-80 GBps no cluster) sem necessidade de armazenamento intermediário, reduzindo custos em mais de 90% comparado a soluções com Amazon FSx for Lustre.

O treinamento distribuído utiliza o framework DeepSpeed ZeRO-2, que equilibra eficiência de memória e desempenho, alcançando throughput de 1.974 amostras por segundo e escalabilidade quase linear entre os 16 nós.

Ampliação da janela de contexto

Para ampliar a análise dos modelos, a equipe implementou técnicas avançadas de paralelismo de contexto, como a atenção em anel (ring attention), que distribui o processamento da sequência de entrada entre GPUs, compartilhando pares chave-valor para acumular resultados de atenção. Isso permitiu aumentar o tamanho máximo de entrada de 640×640×1024 para 1536×1536×2048 voxels, um aumento de 4,5 vezes no volume analisado simultaneamente.

Impacto prático para os usuários

Redução drástica no tempo de treinamento: De seis meses para apenas cinco dias, possibilitando ciclos de iteração semanais.
Maior capacidade analítica: Modelos capazes de capturar informações geológicas locais e globais em uma única análise.
Infraestrutura resiliente e econômica: Uso eficiente de recursos com streaming direto do S3 e escalabilidade facilitada pelo SageMaker HyperPod.

Disponibilidade e como acessar

O Amazon SageMaker HyperPod está disponível para clientes AWS interessados em treinamento distribuído resiliente e escalável. Para começar, é necessário criar uma conta AWS (registro AWS) e explorar a documentação oficial do SageMaker HyperPod (documentação SageMaker HyperPod).

Além disso, a AWS oferece planos flexíveis de treinamento (SageMaker AI flexible training plans) que ajudam a otimizar custos conforme a demanda.