Scaling seismic basis fashions on AWS: Distributed coaching with Amazon SageMaker HyperPod and increasing context home windows

This put up is cowritten with Altay Sansal and Alejandro Valenciano from TGS.

TGS, a geoscience information supplier for the vitality sector, helps firms’ exploration and manufacturing workflows with superior seismic basis fashions (SFMs). These fashions analyze advanced 3D seismic information to determine geological buildings very important for vitality exploration. To assist improve their next-generation fashions as a part of their AWS infrastructure modernization, TGS partnered with the AWS Generative AI Innovation Heart (GenAIIC) to optimize their SFM coaching infrastructure.

This put up describes how TGS achieved near-linear scaling for distributed coaching and expanded context home windows for his or her Imaginative and prescient Transformer-based SFM utilizing Amazon SageMaker HyperPod. This joint resolution lower coaching time from 6 months to only 5 days whereas enabling evaluation of seismic volumes bigger than beforehand doable.

Addressing seismic basis mannequin coaching challenges

TGS’s SFM makes use of a Imaginative and prescient Transformer (ViT) structure with Masked AutoEncoder (MAE) coaching designed by the TGS group to research 3D seismic information. Scaling such fashions presents a number of challenges:

• Information scale and complexity – TGS works with giant volumes of proprietary 3D seismic information saved in domain-specific codecs. The sheer quantity and construction of this information required environment friendly streaming methods to keep up excessive throughput and assist stop GPU idle time throughout coaching.
• Coaching effectivity – Coaching giant FMs on 3D volumetric information is computationally intensive. Accelerating coaching cycles would allow TGS to include new information extra regularly and iterate on mannequin enhancements sooner, delivering extra worth to their shoppers.
• Expanded analytical capabilities – The geological context a mannequin can analyze is determined by how a lot 3D quantity it may well course of without delay. Increasing this functionality would enable the fashions to seize each native particulars and broader geological patterns concurrently.

Understanding these challenges highlights the necessity for a complete method to distributed coaching and infrastructure optimization. The AWS GenAIIC partnered with TGS to develop a complete resolution addressing these challenges.

Resolution overview

The collaboration between TGS and the AWS GenAIIC centered on three key areas: establishing an environment friendly information pipeline, optimizing distributed coaching throughout a number of nodes, and increasing the mannequin’s context window to research bigger geological volumes. The next diagram illustrates the answer structure.

The answer makes use of SageMaker HyperPod to assist present a resilient, scalable coaching infrastructure with computerized well being monitoring and checkpoint administration. The SageMaker HyperPod cluster is configured with AWS Id and Entry Administration (IAM) execution roles scoped to the minimal permissions required for coaching operations, deployed inside a digital non-public cloud (VPC) with community isolation and safety teams limiting communication to licensed coaching nodes. Terabytes of coaching information streams straight from Amazon Easy Storage Service (Amazon S3), assuaging the necessity for intermediate storage layers whereas sustaining excessive throughput. AWS CloudTrail logs API calls to Amazon S3 and SageMaker companies, and Amazon S3 entry logging is enabled on coaching information buckets to offer an in depth audit path of information entry requests. The distributed coaching framework makes use of superior parallelization methods to effectively scale throughout a number of nodes, and context parallelism strategies allow the mannequin to course of considerably bigger 3D volumes than beforehand doable.

The ultimate cluster configuration consisted of 16 Amazon Elastic Compute Cloud (Amazon EC2) P5 cases for the employee nodes built-in by means of the SageMaker AI versatile coaching plans, every containing:

• 8 NVIDIA H200 GPUs with 141GB HBM3e reminiscence per GPU
• 192 vCPUs
• 2048 GB system RAM
• 3200 Gbps EFAv3 networking for ultra-low latency communication

Optimizing the coaching information pipeline

TGS’s coaching dataset consists of 3D seismic volumes saved within the TGS-developed MDIO format—an open supply format constructed on Zarr arrays designed for large-scale scientific information within the cloud. Such volumes can comprise billions of information factors representing underground geological buildings.

Choosing the proper storage method

The group evaluated two approaches for delivering information to coaching GPUs:

• Amazon FSx for Lustre – Copy information from Amazon S3 to a high-speed distributed file system that the nodes learn from. This method gives sub-millisecond latency however requires pre-loading and provisioned storage capability.
• Streaming straight from Amazon S3 – Stream information straight from Amazon S3 utilizing MDIO’s native capabilities with multi-threaded libraries, opening a number of concurrent connections per node.

Deciding on streaming straight from Amazon S3

The important thing architectural distinction lies in how throughput scales with the cluster. With streaming straight from Amazon S3, every coaching node creates unbiased Amazon S3 connections, so combination throughput can scale linearly. With Amazon FSx for Lustre, the nodes share a single file system whose throughput is tied to provisioned storage capability. Utilizing Amazon FSx along with Amazon S3 requires solely a small Amazon FSx storage quantity, which limits the whole cluster to that quantity’s throughput, making a bottleneck because the cluster grows.

Complete testing and value evaluation revealed streaming from Amazon S3 straight because the optimum alternative for this configuration:

• Efficiency – Achieved 4–5 GBps sustained throughput per node utilizing a number of information loader processes with pre-fetching over HTTPS endpoints (TLS 1.2)—adequate to completely make the most of the GPUs.
• Value effectivity – Streaming from Amazon S3 alleviated the necessity for Amazon FSx provisioning, decreasing storage infrastructure prices by over 90% whereas serving to ship 64-80 GBps cluster-wide throughput. The Amazon S3 pay-per-use mannequin was extra economical than provisioning high-throughput Amazon FSx capability.
• Higher scaling – Streaming from Amazon S3 straight scales naturally—every node brings its personal connection bandwidth, avoiding the necessity for advanced capability planning.
• Operational simplicity – No intermediate storage to provision, handle, or synchronize.

The group optimized Amazon S3 connection pooling and carried out parallel information loading to maintain excessive throughput throughout the 16 nodes.

Deciding on the distributed coaching framework

When coaching giant fashions throughout a number of GPUs, the mannequin’s parameters, gradients, and optimizer states should be distributed throughout units. The group evaluated completely different distributed coaching approaches to search out the optimum steadiness between reminiscence effectivity and coaching throughput:

• ZeRO-2 (Zero Redundancy Optimizer Stage 2) – This method partitions gradients and optimizer states throughout GPUs whereas conserving a full copy of mannequin parameters on every GPU. This helps cut back reminiscence utilization whereas sustaining quick communication, as a result of every GPU can straight entry the parameters throughout the ahead cross with out ready for information from different GPUs.
• ZeRO-3 – This method goes additional by additionally partitioning mannequin parameters throughout GPUs. Though this helps maximize reminiscence effectivity (enabling bigger fashions), it requires extra frequent communication between GPUs to assemble parameters throughout computation, which might cut back throughput.
• FSDP2 (Absolutely Sharded Information Parallel v2) – PyTorch’s native method equally shards parameters, gradients, and optimizer states. It presents tight integration with PyTorch however entails related communication trade-offs as ZeRO-3.

Complete testing revealed DeepSpeed ZeRO-2 because the optimum framework for this configuration, delivering sturdy efficiency whereas effectively managing reminiscence:

• ZeRO-2 – 1,974 samples per second (carried out)
• FSDP2 – 1,833 samples per second
• ZeRO-3 – 869 samples per second

This framework alternative offered the inspiration for attaining near-linear scaling throughout a number of nodes. The mixture of those three key optimizations helped ship the dramatic coaching acceleration:

• Environment friendly distributed coaching – DeepSpeed ZeRO-2 enabled near-linear scaling throughout 128 GPUs (16 nodes × 8 GPUs)
• Excessive-throughput information pipeline – Streaming from Amazon S3 straight sustained 64–80 GBps combination throughput throughout the cluster

Collectively, these enhancements helped cut back coaching time from 6 months to five days—enabling TGS to iterate on mannequin enhancements weekly relatively than semi-annually.

Increasing analytical capabilities

One of the important achievements was increasing the mannequin’s discipline of view—how a lot 3D geological quantity it may well analyze concurrently. A bigger context window permits the mannequin to seize each effective particulars (small fractures) and broad patterns (basin-wide fault techniques) in a single cross, serving to present insights that had been beforehand undetectable throughout the constraints of smaller evaluation home windows for TGS’s shoppers. The implementation by the TGS and AWS groups concerned adapting the next superior methods to allow ViTs to course of considerably bigger 3D seismic volumes:

• Ring consideration implementation – Every GPU processes a portion of the enter sequence whereas circulating key-value pairs to neighboring GPUs, regularly accumulating consideration outcomes throughout the distributed system. PyTorch gives an API that makes this easy:

from torch.distributed.tensor.parallel import context_parallel

# Wrap consideration computation with context parallelism
with context_parallel(
buffers=[query, key, value], # Tensors to shard
buffer_seq_dims=[1, 1, 1] # Dimension to shard alongside (sequence dimension)
):
# Normal scaled dot-product consideration – mechanically turns into Ring Consideration
attention_output = torch.nn.practical.scaled_dot_product_attention(
question, key, worth, attn_mask=None
)

• Dynamic masks ratio adjustment – The MAE coaching method required ensuring unmasked patches plus classification tokens are evenly divisible throughout units, necessitating adaptive masking methods.
• Decoder sequence administration – The decoder reconstructs the complete picture by processing each the unmasked patches from the encoder and the masked patches. This creates a special sequence size that additionally must be divisible by the variety of GPUs.

The previous implementation enabled processing of considerably bigger 3D seismic volumes as illustrated within the following desk.

Metric
Earlier (Baseline)
With Context Parallelism

Most enter dimension
640 × 640 × 1,024 voxels
1,536 × 1,536 × 2,048 voxels

Context size
102,400 tokens
1,170,000 tokens

Quantity enhance
1×
4.5×

The next determine gives an instance of 2D mannequin context dimension.

This growth permits TGS’s fashions to seize geological options throughout broader spatial contexts, serving to improve the analytical capabilities they’ll supply to shoppers.

Outcomes and affect

The collaboration between TGS and the AWS GenAIIC delivered substantial enhancements throughout a number of dimensions:

• Vital coaching acceleration – The optimized distributed coaching structure lowered coaching time from 6 months to five days—an approximate 36-fold speedup, enabling TGS to iterate sooner and incorporate new geological information extra regularly into their fashions.
• Close to-linear scaling – The answer demonstrated sturdy scaling effectivity from single-node to 16-node configurations, attaining roughly 90–95% parallel effectivity with minimal efficiency degradation because the cluster dimension elevated.
• Expanded analytical capabilities – The context parallelism implementation permits coaching on bigger 3D volumes, permitting fashions to seize geological options throughout broader spatial contexts.
• Manufacturing-ready, cost-efficient infrastructure – The SageMaker HyperPod primarily based resolution with streaming from Amazon S3 helps present a cheap basis that scales effectively as coaching necessities develop, whereas serving to ship the resilience, flexibility, and operational effectivity wanted for manufacturing AI workflows.

These enhancements set up a robust basis for TGS’s AI-powered analytics system, delivering sooner mannequin iteration cycles and broader geological context per evaluation to shoppers whereas serving to defend TGS’s priceless information property.

Classes realized and finest practices

A number of key classes emerged from this collaboration which may profit different organizations working with large-scale 3D information and distributed coaching:

• Systematic scaling method – Beginning with a single-node baseline institution earlier than progressively increasing to bigger clusters enabled systematic optimization at every stage whereas managing prices successfully.
• Information pipeline optimization is crucial – For data-intensive workloads, considerate information pipeline design can present sturdy efficiency. Direct streaming from object storage with applicable parallelization and prefetching delivered the throughput wanted with out advanced intermediate storage layers.
• Batch dimension tuning is nuanced – Rising batch dimension doesn’t at all times enhance throughput. The group discovered excessively giant batch dimension can create bottlenecks in making ready and transferring information to GPUs. By means of systematic testing at completely different scales, the group recognized the purpose the place throughput plateaued, indicating the information loading pipeline had change into the limiting issue relatively than GPU computation. This optimum steadiness maximized coaching effectivity with out over-provisioning sources.
• Framework choice is determined by your particular necessities – Completely different distributed coaching frameworks contain trade-offs between reminiscence effectivity and communication overhead. The optimum alternative is determined by mannequin dimension, {hardware} traits, and scaling necessities.
• Incremental validation – Testing configurations at smaller scales earlier than increasing to full manufacturing clusters helped determine optimum settings whereas controlling prices throughout the growth section.

Conclusion

By partnering with the AWS GenAIIC, TGS has established an optimized, scalable infrastructure for coaching SFMs on AWS. The answer helps speed up coaching cycles whereas increasing the fashions’ analytical capabilities, serving to TGS ship enhanced subsurface analytics to shoppers within the vitality sector. The technical improvements developed throughout this collaboration—notably the difference of context parallelism to ViT architectures for 3D volumetric information—show the potential for making use of superior AI methods to specialised scientific domains. As TGS continues to broaden its subsurface AI system and broader AI capabilities, this basis can assist future enhancements resembling multi-modal integration and temporal evaluation.

To study extra about scaling your individual FM coaching workloads, discover SageMaker HyperPod for resilient distributed coaching infrastructure, or evaluation the distributed coaching finest practices within the SageMaker documentation. For organizations thinking about related collaborations, the AWS Generative AI Innovation Heart companions with prospects to assist speed up their AI initiatives.

Acknowledgement

Particular due to Andy Lapastora, Bingchen Liu, Prashanth Ramaswamy, Rohit Thekkanal, Jared Kramer, Arun Ramanathan and Roy Allela for his or her contribution.

In regards to the authors

Haotian An

Haotian An is a Machine Studying Engineer on the AWS Generative AI Innovation Heart, the place he makes a speciality of customizing basis fashions and distributed coaching at scale. He works intently with prospects to adapt generative AI to their particular use instances, serving to them unlock new capabilities and drive measurable enterprise outcomes.

Manoj Alwani

Manoj Alwani is a Senior Utilized Scientist on the Generative AI Innovation Heart at AWS, the place he helps organizations unlock the potential of cutting-edge AI expertise. With deep experience throughout the whole generative AI analysis stack, Manoj works intently with prospects from numerous industries to speed up their GenAI adoption and drive significant enterprise outcomes. He brings over 13 years of hands-on expertise in growing and deploying machine studying options at scale.

Debby Wehner

Debby Wehner is a Machine Studying Engineer on the AWS Generative AI Innovation Heart, specializing in giant language mannequin customization and optimization. Beforehand, as a full-stack software program engineer at Amazon, she constructed AI-powered procuring functions reaching over 100 million month-to-month customers. She holds a PhD in Computational Geophysics from the College of Cambridge, in addition to a BSc and MSc from Freie Universität Berlin.

Altay Sansal

Altay Sansal is a Senior Information Science Lead at TGS in Houston, Texas, specializing in AI/ML functions for geophysics and seismic information, together with basis fashions, large-scale coaching, and open-source instruments just like the MDIO format. He holds an M.S. in Geophysics from the College of Houston and has authored key publications resembling “Scaling Seismic Basis Fashions” and “MDIO: Open-source format for multidimensional vitality information”, whereas actively contributing to geoscience ML by means of GitHub and trade occasions.

Alejandro Valenciano

Alejandro Valenciano is the Director of Information Science at TGS, the place he leads superior analytics and information science initiatives that unlock insights from subsurface and energy-related information, driving innovation throughout seismic, properly, and machine studying workflows. He has developed and utilized machine studying fashions for duties resembling basin-scale log prediction, superior seismic processing, and Basis Fashions. He regularly contributes to trade conferences and technical publications. His work spans information administration, ML/AI functions in geoscience, and the combination of scalable information platforms to assist exploration and vitality options.