在 AWS 上扩展地震基础模型:使用 Amazon SageMaker HyperPod 进行分布式训练并扩展上下文窗口

本文介绍了 TGS 如何使用 Amazon SageMaker HyperPod 实现分布式训练的近线性扩展以及基于 Vision Transformer 的 SFM 的扩展上下文窗口。该联合解决方案将训练时间从 6 个月缩短至仅 5 天,同时能够分析比以前更大的地震体积。

来源:亚马逊云科技 _机器学习

这篇文章是与 TGS 的 Altay Sansal 和 Alejandro Valenciano 共同撰写的。

TGS 是能源行业的地球科学数据提供商,通过先进的地震基础模型 (SFM) 支持公司的勘探和生产工作流程。这些模型分析复杂的 3D 地震数据,以确定对能源勘探至关重要的地质结构。为了帮助增强下一代模型,作为 AWS 基础设施现代化的一部分,TGS 与 AWS 生成 AI 创新中心 (GenAIIC) 合作,优化其 SFM 培训基础设施。

本文介绍了 TGS 如何使用 Amazon SageMaker HyperPod 实现分布式训练的近线性扩展以及基于 Vision Transformer 的 SFM 的扩展上下文窗口。该联合解决方案将训练时间从 6 个月缩短至仅 5 天,同时能够分析比以前更大的地震体积。

解决地震基础模型训练挑战

TGS 的 SFM 使用由 TGS 团队设计的 Vision Transformer (ViT) 架构和 Masked AutoEncoder (MAE) 训练来分析 3D 地震数据。扩展此类模型面临一些挑战:

  • 数据规模和复杂性 - TGS 可处理以特定领域格式存储的大量专有 3D 地震数据。这些数据的庞大数量和结构需要高效的流策略来保持高吞吐量并帮助防止训练期间的 GPU 空闲时间。
  • 训练效率——在 3D 体积数据上训练大型 FM 需要大量计算。加快培训周期将使 TGS 能够更频繁地整合新数据并更快地迭代模型改进,从而为客户提供更多价值。
  • 扩展的分析能力 - 模型可以分析的地质背景取决于它可以一次处理多少 3D 体积。扩展此功能将使模型能够同时捕获局部细节和更广泛的地质模式。
  • 解决方案概述

  • 8 个 NVIDIA H200 GPU,每个 GPU 具有 141GB HBM3e 内存