Cómo TGS aceleró modelos sísmicos con SageMaker HyperPod y ventanas de contexto más grandes

Introducción

TGS, proveedor de datos geocientíficos para el sector energético, abordó un reto crítico: cómo entrenar modelos foundation para datos sísmicos 3D extremadamente voluminosos de forma eficiente y con mayor contexto geológico. En conjunto con el AWS Generative AI Innovation Center (GenAIIC), TGS reestructuró su infraestructura en AWS para lograr escalado casi lineal en entrenamiento distribuido y ampliar la ventana de contexto de sus modelos Vision Transformer entrenados con Masked AutoEncoder (MAE). El resultado fue notable: un ciclo de entrenamiento que antes tardaba alrededor de seis meses se redujo a cinco días.

Desafíos del entrenamiento de modelos sísmicos

Los modelos sísmicos (SFM) enfrentan varias limitaciones técnicas y operacionales:

• Escala y complejidad de datos: los volúmenes sísmicos 3D contienen miles de millones de puntos almacenados en formatos específicos del dominio. TGS utiliza el formato MDIO, basado en Zarr, pensado para datos científicos a escala en la nube.
• Eficiencia de entrenamiento: procesar datos volumétricos 3D con arquitecturas pesadas como ViT+MAE exige recursos computacionales masivos y una estrategia de I/O que evite GPUs ociosas.
• Alcance analítico: la capacidad del modelo para identificar patrones geológicos depende de cuánto volumen 3D pueda analizar al mismo tiempo; ventanas de contexto mayores permiten capturar tanto detalles locales como estructuras regionales.

Frente a esto, la solución debía optimizar el pipeline de datos, distribuir el entrenamiento eficientemente y ampliar la ventana de contexto sin comprometer el rendimiento.

Visión general de la solución

La colaboración se centró en tres pilares: pipeline de datos eficiente, entrenamiento distribuido escalable y expansión de la ventana de contexto.

Componentes clave:

• SageMaker HyperPod: proporcionó una plataforma resistente y escalable para orquestar el cluster de entrenamiento, con monitoreo de salud y manejo de checkpoints.
• Seguridad y auditoría: roles de IAM con permisos mínimos, despliegue en VPC para aislamiento de red, y logs en CloudTrail y S3 access logging para trazabilidad.
• Datos en S3: terabytes de datos sísmicos se transmitieron directamente desde Amazon S3, evitando capas intermedias de almacenamiento.
• Paralelismo avanzado: técnicas de paralelización y paralelismo de contexto que permitieron al modelo procesar volúmenes 3D significativamente mayores.

La configuración final del cluster incluyó 16 instancias EC2 P5 (workers) integradas con los planes de entrenamiento flexible de SageMaker. Cada instancia tenía: 8 GPUs NVIDIA H200 con 141 GB de HBM3e por GPU, 192 vCPUs, 2048 GB de RAM y conectividad EFAv3 a 3200 Gbps.

Optimización del pipeline de datos

TGS evaluó dos alternativas para alimentar datos a las GPUs:

1. Amazon FSx for Lustre: copiar datos desde S3 a un filesystem distribuido de alta velocidad que los nodos leen. Ofrece latencias muy bajas pero requiere pre-carga y capacidad provisionada.
2. Streaming directo desde Amazon S3: utilizar las capacidades nativas de MDIO sobre Zarr con librerías multihilo para abrir múltiples conexiones concurrentes por nodo.

La diferencia arquitectónica clave fue cómo escala el throughput cuando crece el cluster. Con streaming directo, cada nodo abre conexiones independientes a S3, permitiendo que el throughput agregado escale de forma lineal. Con FSx, los nodos dependen de un volumen compartido cuya capacidad provisionada puede convertirse en cuello de botella.

Tras pruebas y análisis de costos, el streaming directo ganó por:

• Rendimiento: se logró un throughput sostenido de 4–5 GB/s por nodo usando múltiples procesos de data loader con prefetch sobre endpoints HTTPS (TLS 1.2), suficiente para mantener las GPUs ocupadas.
• Eficiencia de costos: al evitar la provisión de FSx, se redujo el costo de infraestructura en más del 90%. El modelo pay-per-use de S3 resultó más económico para este patrón de acceso.

En conjunto, el cluster alcanzó 64–80 GB/s de throughput agregado, lo que permitió explotar el rendimiento de los nodos P5 sin cuellos de botella de I/O.

Entrenamiento distribuido y escalado

SageMaker HyperPod facilitó el despliegue de un cluster resiliente con gestión automática de salud y checkpoints. La arquitectura aprovechó técnicas avanzadas de paralelización para maximizar la utilización de las GPUs H200 y la red EFAv3 de baja latencia.

Las prácticas clave incluyeron:

• Multiprocesamiento de loaders y prefetch para minimizar latencia de datos.
• Conexiones S3 concurrentes por nodo para escalar throughput por participante del cluster.
• Checkpointing continuo gestionado por HyperPod para recuperación rápida ante fallos.

Gracias a estas optimizaciones, el entrenamiento mostró escalado casi lineal conforme se agregaban instancias, reduciendo los tiempos de experimentación y permitiendo iteraciones de modelo más rápidas.

Ampliar la ventana de contexto

Para capturar estructuras geológicas a diferentes escalas, el equipo implementó paralelismo de contexto: estrategias que permiten al modelo procesar bloques más grandes del volumen 3D distribuyendo el trabajo entre nodos y GPUs sin sacrificar rendimiento. Esto posibilitó analizar volúmenes 3D mayores que los previamente procesables, mejorando la capacidad del modelo para identificar patrones regionales y locales de forma simultánea.

Resultados y lecciones prácticas

Resultados observados:

• Tiempo de entrenamiento reducido de aproximadamente 6 meses a 5 días en la configuración final.
• Escalado casi lineal gracias al streaming desde S3 y la configuración de HyperPod con instancias P5.
• Ahorro significativo en costos de almacenamiento al evitar FSx provisioning (más del 90% en infraestructura de almacenamiento para este caso).
• Capacidad para procesar volúmenes 3D más grandes mediante paralelismo de contexto.

Lecciones aplicables para equipos en Latinoamérica:

• Evaluar patrones de acceso a datos antes de provisionar sistemas de archivos de alto rendimiento: el streaming desde S3 puede ser más sencillo y económico cuando el formato de datos y la librería soportan múltiples conexiones concurrentes.
• Diseñar pipelines de datos que mantengan las GPUs ocupadas: loaders multihilo y prefetching son críticos para aprovechar GPUs de gran memoria como las H200.
• Considerar soluciones gestionadas como SageMaker HyperPod para orquestar clusters grandes, simplificar la recuperación ante fallos y acelerar la iteración de experimentos.

Para empresas energéticas y de recursos en la región, estas prácticas permiten incorporar modelos foundation que analicen mayores volúmenes sísmicos y ofrezcan resultados más rápidos y con mayor contexto geológico.

Conclusión

La colaboración entre TGS y el AWS GenAIIC demuestra que, combinando formatos de datos diseñados para la nube (MDIO/Zarr), streaming eficiente desde S3, y una infraestructura de entrenamiento escalable con SageMaker HyperPod, es posible transformar un proceso de entrenamiento que tomaba meses en uno de días. Más allá del ahorro de tiempo y costos, la capacidad de procesar volúmenes 3D mayores abre nuevas oportunidades analíticas para la exploración geocientífica. Para organizaciones latinoamericanas del sector energético, adoptar estas prácticas puede acelerar la adopción de modelos foundation y mejorar la toma de decisiones basada en datos sísmicos de alta resolución.