Distributed Training of 3D Seismic Foundation Models

This technical report presents a systematic approach to scaling the training of Seismic Foundation Models (SFMs) for 3D volumetric data using a distributed training framework. The study addresses challenges in data management, computational efficiency, and the need for expanded spatial context. The authors demonstrate the effectiveness of the Vision Transformer-Masked AutoEncoder (ViT-MAE) architecture, achieving near-linear scaling from one to sixteen nodes with 128 GPUs. Key findings include the identification of optimal data streaming methods that enhance throughput significantly compared to shared file systems. The report also discusses the implementation of context parallelism to manage memory constraints and expand the model's analytical capacity. Additionally, it evaluates various distributed training frameworks, concluding that the DeepSpeed ZeRO Stage 2 framework is optimal for their specific memory profile. The results indicate substantial improvements in processing large 3D seismic volumes, which are critical for geological interpretation tasks.