Distributed training of 3D seismic foundation models at scale

Paper submitted to IMAGE 2026, by Manoj Alwani, Haotian An, Andy Lapastora, Prashanth Ramaswamy, Bingchen Liu, Rohit Thekkanal, Jared Kramer, Arun Ramanathan (Amazon Web Services), Ben Lasscock, Alejandro Valenciano, Altay Sansal, Sathiya Namasivayam (TGS).

Abstract

Developing Seismic Foundation Models (SFMs) for 3D volumetric data requires addressing challenges in data management, computational efficiency, and the need for expanded spatial context. This paper presents a systematic approach to scaling 3D SFM training based on the Vision Transformer-Masked AutoEncoder (ViT-MAE) (He et al., 2021) to 16 nodes (128 GPUs). We show that our approach achieves near-linear scaling from 1 to 16 nodes. We demonstrate that, across multiple nodes, streaming data directly from object storage to each node in the distributed training cluster achieved 5.7× cluster-wide throughput and identified the DeepSpeed ZeRO Stage2 (Rajbhandari et al., 2020) training framework as optimal for our memory profile (127% faster than ZeRO Stage3). Through context parallelism using ring attention, we enable processing of 1536×1536×2048 voxel seismic volumes, a 4.5x volumetric expansion that was previously infeasible due to memory constraints, allowing the model to see a larger geological context critical for interpretation tasks.

Read the full article here.