Cómo TGS y AWS Redujeron el Entrenamiento del Modelo de IA Sísmica de Seis Meses a Cinco Días

TGS junto con AWS demostró un resultado raro para grandes proyectos de IA: el entrenamiento de un modelo foundation sísmico, que anteriormente requería aproximadamente seis meses, se redujo a cinco días. Al mismo tiempo, el equipo aumentó el tamaño de la ventana de contexto, es decir, el volumen de datos geológicos tridimensionales que el modelo puede analizar en un único paso. Para las empresas que trabajan en exploración de subsuelo, esto significa iteraciones más rápidas y una visión más completa de las estructuras subterráneas.

TGS es un proveedor de datos geocientíficos para el sector energético. La empresa utiliza modelos foundation sísmicos para analizar matrices 3D complejas y encontrar estructuras geológicas importantes para la exploración y producción. La arquitectura básica del modelo se construye sobre Vision Transformer y el esquema de entrenamiento Masked AutoEncoder.

El problema principal aquí no es solo el tamaño del modelo en sí, sino también la naturaleza de los datos: los volúmenes sísmicos constan de miles de millones de puntos, se almacenan en formatos especializados y requieren entrega continua a GPU sin tiempo de inactividad. Junto con AWS Generative AI Innovation Center, la empresa migró el entrenamiento a Amazon SageMaker HyperPod y montó un clúster de 16 nodos Amazon EC2 P5. Cada nodo incluía 8 GPU NVIDIA H200 con 141 GB de memoria HBM3e, 192 vCPU, 2 TB de RAM y red EFAv3 con una velocidad de 3200 Gbps.

En total, son 128 GPU. Según AWS, tal configuración proporcionó un escalado casi lineal: la eficiencia de la operación paralela al escalar de uno a 16 nodos se mantuvo en aproximadamente 90–95%. Un enfoque separado fue en el pipeline de datos.

En lugar del esquema clásico con Amazon FSx for Lustre, el equipo eligió streaming directo desde Amazon S3. La razón es simple: a medida que crece el clúster, S3 permite que cada nodo agregue su propio ancho de banda, mientras que un sistema de archivos compartido se convierte en un cuello de botella más rápidamente. El conjunto de datos de entrenamiento de TGS se almacena en el formato MDIO desarrollado por la empresa basado en Zarr, optimizado para grandes datos científicos en la nube.

Después de configurar la carga multihilo y el prefetching, cada nodo logró consistentemente 4–5 GB/s, y todo el clúster — 64–80 GB/s. Además, esto redujo los costos de infraestructura de almacenamiento en más del 90%. Para el entrenamiento distribuido, el equipo comparó varios enfoques: DeepSpeed ZeRO-2, ZeRO-3 y FSDP2.

ZeRO-2 mostró el mejor equilibrio entre velocidad y ahorro de memoria: 1.974 muestras por segundo en comparación con 1.833 para FSDP2 y 869 para ZeRO-3.

Esto es importante porque en tales tareas, el máximo ahorro de memoria no siempre produce el mejor resultado general: si la comunicación entre GPU se vuelve demasiado cara, el rendimiento cae drásticamente. En este proyecto, el enfoque no fue en la opción teóricamente más densa, sino en la configuración que realmente procesa el entrenamiento más rápido en un escenario de producción. El segundo logro clave se relaciona con la expansión de la ventana de contexto.

Para modelos sísmicos, esto afecta directamente la calidad del análisis: cuanto mayor sea el volumen de roca que el modelo ve en un único paso, mejor captura tanto los detalles locales, como pequeñas fracturas, como los patrones grandes — por ejemplo, sistemas de fallas a la escala de toda una cuenca. Gracias al context parallelism y la adaptación de ring attention para la arquitectura Vision Transformer para datos 3D, el tamaño máximo de entrada aumentó de 640 x 640 x 1.024 a 1.

536 x 1.536 x 2.048 voxels.

La longitud del contexto aumentó de 102.400 a 1,17 millones de tokens, y el volumen analizado — aproximadamente 4,5 veces. El significado práctico del caso es que TGS ahora puede actualizar modelos no una vez cada seis meses, sino efectivamente semanalmente, conectar más rápidamente nuevos datos geológicos y proporcionar a los clientes un contexto de análisis más amplio.

Para AWS, este es otro ejemplo demostrativo de que los modelos foundation especializados en ciencia e industria están limitados no solo por la arquitectura del modelo, sino también por la organización adecuada de datos, redes y entrenamiento distribuido. Para el mercado en general, la señal es clara: la IA especializada está comenzando a ganar donde antes dominaban ciclos computacionales largos e infraestructura demasiado cara.