Cómo TGS y AWS redujeron el entrenamiento de un modelo sísmico de IA de seis meses a cinco días

TGS, junto con AWS, ha logrado un resultado poco habitual en los grandes proyectos de IA: el entrenamiento de un modelo fundacional sísmico, que antes llevaba unos seis meses, se ha reducido a cinco días. Al mismo tiempo, el equipo amplió el tamaño de la ventana de contexto, es decir, el volumen de datos geológicos tridimensionales que el modelo puede analizar de una sola pasada. Para las empresas que trabajan en exploración del subsuelo, esto se traduce en iteraciones más rápidas y una imagen más completa de las estructuras subterráneas.

TGS es un proveedor de datos geocientíficos para el sector energético. La compañía utiliza modelos fundacionales sísmicos para analizar matrices 3D complejas y localizar estructuras geológicas relevantes para la exploración y la producción. La arquitectura base del modelo se apoya en Vision Transformer y en el esquema de entrenamiento Masked AutoEncoder. El principal problema no es solo el tamaño del modelo, sino la naturaleza de los datos: los volúmenes sísmicos contienen miles de millones de puntos, se almacenan en formatos especializados y hay que alimentar las GPU de forma continua, sin tiempos muertos.

Junto con el AWS Generative AI Innovation Center, la compañía trasladó 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 a 3200 Gbit/s. En total, 128 GPU. Según AWS, esta configuración logró un escalado casi lineal: la eficiencia del trabajo en paralelo se mantuvo en torno al 90–95 % al pasar de uno a 16 nodos.

El pipeline de datos recibió atención específica. En lugar del esquema clásico con Amazon FSx for Lustre, el equipo optó por el streaming directo desde Amazon S3. El motivo es sencillo: al crecer el clúster, S3 permite que cada nodo añada su propio ancho de banda, mientras que un sistema de archivos compartido se convierte antes en cuello de botella.

El dataset de entrenamiento de TGS se almacena en MDIO, un formato desarrollado por la propia empresa sobre Zarr y optimizado para grandes datos científicos en la nube. Tras ajustar la carga multihilo y el prefetching, cada nodo recibía de forma estable 4–5 GB/s y todo el clúster, 64–80 GB/s. Además, esto redujo más de un 90 % el coste de la infraestructura de almacenamiento.

Para el entrenamiento distribuido, el equipo comparó varios enfoques: DeepSpeed ZeRO-2, ZeRO-3 y FSDP2. ZeRO-2 ofreció el mejor equilibrio entre velocidad y ahorro de memoria: 1974 muestras por segundo, frente a 1833 de FSDP2 y 869 de ZeRO-3. Es relevante porque, en este tipo de tareas, el máximo ahorro de memoria no siempre da el mejor resultado final: si el intercambio entre GPU resulta demasiado costoso, el throughput cae en picado. En este proyecto no se apostó por la opción teóricamente más compacta, sino por la configuración que ejecuta el entrenamiento más rápido en un escenario de producción real.

El segundo logro clave tiene que ver con la ampliación de la ventana de contexto. En los modelos sísmicos influye directamente en la calidad del análisis: cuanto mayor es el volumen de roca que el modelo ve de una pasada, mejor capta tanto los detalles locales, como pequeñas fracturas, como los grandes patrones, por ejemplo los sistemas de fallas a escala de toda una cuenca. Gracias al context parallelism y a la adaptación de ring attention a la arquitectura ViT para datos 3D, el tamaño máximo de entrada pasó de 640 x 640 x 1024 a 1536 x 1536 x 2048 vóxeles.

La longitud del contexto aumentó de 102 400 a 1,17 millones de tokens, y el volumen analizado, aproximadamente 4,5 veces.

La utilidad práctica del caso es que TGS ya puede actualizar sus modelos prácticamente cada semana, y no una vez cada seis meses, incorporar antes nuevos datos geológicos y ofrecer a los clientes un contexto de análisis más amplio. Para AWS es otro ejemplo ilustrativo de que los modelos fundacionales especializados en ciencia e industria dependen no solo de la arquitectura del modelo, sino también de organizar bien los datos, la red y el entrenamiento distribuido. Para el mercado en general, la señal es clara: la IA muy especializada empieza a imponerse donde antes dominaban los largos ciclos de cómputo y una infraestructura demasiado cara.