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

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

TGS, a geoscience knowledge supplier for the vitality sector, helps corporations’ exploration and manufacturing workflows with superior seismic basis fashions (SFMs). These fashions analyze complicated 3D seismic knowledge 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 Middle (GenAIIC) to optimize their SFM coaching infrastructure.

This publish 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 answer lower coaching time from 6 months to simply 5 days whereas enabling evaluation of seismic volumes bigger than beforehand attainable.

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 investigate 3D seismic knowledge. Scaling such fashions presents a number of challenges:

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

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

Answer overview

The collaboration between TGS and the AWS GenAIIC targeted on three key areas: establishing an environment friendly knowledge pipeline, optimizing distributed coaching throughout a number of nodes, and increasing the mannequin’s context window to investigate 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 automated 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 approved coaching nodes. Terabytes of coaching knowledge streams instantly 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 knowledge buckets to offer an in depth audit path of knowledge 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 attainable.

The ultimate cluster configuration consisted of 16 Amazon Elastic Compute Cloud (Amazon EC2) P5 cases for the employee nodes built-in by 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 knowledge 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 knowledge within the cloud. Such volumes can include billions of knowledge factors representing underground geological buildings.

Selecting the best storage strategy

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

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

Selecting streaming instantly from Amazon S3

The important thing architectural distinction lies in how throughput scales with the cluster. With streaming instantly 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 instantly because the optimum selection for this configuration:

• Efficiency – Achieved 4–5 GBps sustained throughput per node utilizing a number of knowledge loader processes with pre-fetching over HTTPS endpoints (TLS 1.2)—adequate to completely make the most of the GPUs.
• Price effectivity – Streaming from Amazon S3 alleviated the necessity for Amazon FSx provisioning, lowering 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 instantly scales naturally—every node brings its personal connection bandwidth, avoiding the necessity for complicated capability planning.
• Operational simplicity – No intermediate storage to provision, handle, or synchronize.

The group optimized Amazon S3 connection pooling and applied parallel knowledge loading to maintain excessive throughput throughout the 16 nodes.

Choosing 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 totally different distributed coaching approaches to search out the optimum steadiness between reminiscence effectivity and coaching throughput:

• ZeRO-2 (Zero Redundancy Optimizer Stage 2) – This strategy partitions gradients and optimizer states throughout GPUs whereas holding a full copy of mannequin parameters on every GPU. This helps scale back reminiscence utilization whereas sustaining quick communication, as a result of every GPU can instantly entry the parameters in the course of the ahead move with out ready for knowledge from different GPUs.
• ZeRO-3 – This strategy 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 collect parameters throughout computation, which might scale back throughput.
• FSDP2 (Absolutely Sharded Information Parallel v2) – PyTorch’s native strategy equally shards parameters, gradients, and optimizer states. It gives 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 robust efficiency whereas effectively managing reminiscence:

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

This framework selection supplied the muse for reaching near-linear scaling throughout a number of nodes. The mix 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 knowledge pipeline – Streaming from Amazon S3 instantly sustained 64–80 GBps combination throughput throughout the cluster

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

Increasing analytical capabilities

Some 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 positive particulars (small fractures) and broad patterns (basin-wide fault techniques) in a single move, serving to present insights that had been beforehand undetectable throughout the constraints of smaller evaluation home windows for TGS’s purchasers. 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, step by step accumulating consideration outcomes throughout the distributed system. PyTorch supplies 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 - routinely turns into Ring Consideration
    attention_output = torch.nn.purposeful.scaled_dot_product_attention(
        question, key, worth, attn_mask=None
    )

• Dynamic masks ratio adjustment – The MAE coaching strategy 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 distinct 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.

The next determine supplies an instance of 2D mannequin context measurement.

This growth permits TGS’s fashions to seize geological options throughout broader spatial contexts, serving to improve the analytical capabilities they will provide to purchasers.

Outcomes and affect

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

• Important coaching acceleration – The optimized distributed coaching structure diminished coaching time from 6 months to five days—an approximate 36-fold speedup, enabling TGS to iterate sooner and incorporate new geological knowledge extra ceaselessly into their fashions.
• Close to-linear scaling – The answer demonstrated robust scaling effectivity from single-node to 16-node configurations, reaching roughly 90–95% parallel effectivity with minimal efficiency degradation because the cluster measurement elevated.
• Expanded analytical capabilities – The context parallelism implementation allows 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 answer 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 purchasers whereas serving to defend TGS’s precious knowledge property.

Classes realized and greatest practices

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

• Systematic scaling strategy – 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 knowledge pipeline design can present robust efficiency. Direct streaming from object storage with applicable parallelization and prefetching delivered the throughput wanted with out complicated intermediate storage layers.
• Batch measurement tuning is nuanced – Growing batch measurement doesn’t all the time enhance throughput. The group discovered excessively giant batch measurement can create bottlenecks in getting ready and transferring knowledge to GPUs. Via systematic testing at totally different scales, the group recognized the purpose the place throughput plateaued, indicating the info loading pipeline had grow to be the limiting issue reasonably than GPU computation. This optimum steadiness maximized coaching effectivity with out over-provisioning assets.
• Framework choice depends upon your particular necessities – Completely different distributed coaching frameworks contain trade-offs between reminiscence effectivity and communication overhead. The optimum selection depends upon mannequin measurement, {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 in the course of the improvement part.

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 purchasers within the vitality sector. The technical improvements developed throughout this collaboration—notably the variation of context parallelism to ViT architectures for 3D volumetric knowledge—display 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 help future enhancements equivalent to multi-modal integration and temporal evaluation.

To be taught extra about scaling your personal FM coaching workloads, discover SageMaker HyperPod for resilient distributed coaching infrastructure, or evaluate the distributed coaching greatest practices within the SageMaker documentation. For organizations concerned about related collaborations, the AWS Generative AI Innovation Middle 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