Оптимизируйте обучение моделей с помощью Amazon SageMaker HyperPod

Компания TGS, занимающаяся геонаучными данными для энергетического сектора, разрабатывает продвинутые модели сейсмических оснований (SFM), которые анализируют сложные 3D сейсмические данные для выявления геологических структур, необходимых для энергетической разведки. В рамках модернизации своей инфраструктуры на AWS, TGS сотрудничает с AWS Generative AI Innovation Center (GenAIIC) для оптимизации своей инфраструктуры обучения SFM. В данной статье описывается, как TGS достигла почти линейного масштабирования для распределенного обучения и расширила контекстные окна для своей модели SFM на основе Vision Transformer с помощью Amazon SageMaker HyperPod. Это совместное решение сократило время обучения с 6 месяцев до всего 5 дней, что позволило анализировать сейсмические объемы, превышающие ранее возможные.

Модель SFM использует архитектуру Vision Transformer (ViT) с обучением Masked AutoEncoder (MAE), разработанную командой TGS. Масштабирование таких моделей сопряжено с несколькими проблемами: объем и сложность данных, эффективность обучения и расширенные аналитические возможности. TGS работает с большими объемами закрытых 3D сейсмических данных, хранящимися в специфичных для области форматах. Огромный объем и структура этих данных требуют эффективных стратегий потоковой передачи для поддержания высокой пропускной способности и предотвращения простоя GPU во время обучения.

Обучение больших моделей на 3D объемных данных требует значительных вычислительных ресурсов. Ускорение циклов обучения позволит TGS чаще интегрировать новые данные и быстрее вносить улучшения в модели, что принесет больше ценности их клиентам. Расширенные аналитические возможности модели зависят от того, сколько 3D объема она может обрабатывать одновременно. Расширение этой способности позволит моделям захватывать как локальные детали, так и более широкие геологические паттерны одновременно.

Понимание этих вызовов подчеркивает необходимость комплексного подхода к распределенному обучению и оптимизации инфраструктуры. AWS GenAIIC сотрудничала с TGS для разработки комплексного решения, которое решает эти проблемы. Основное внимание в сотрудничестве было уделено трем ключевым областям: установлению эффективного канала передачи данных, оптимизации распределенного обучения на нескольких узлах и расширению контекстного окна модели для анализа больших геологических объемов.

Решение использует SageMaker HyperPod для обеспечения устойчивой и масштабируемой инфраструктуры обучения с автоматическим мониторингом состояния и управлением контрольными точками. Кластер SageMaker HyperPod настроен с ролями выполнения AWS Identity and Access Management (IAM), ограниченными минимально необходимыми разрешениями для операций обучения, развернутыми в виртуальной частной облачной сети (VPC) с сетевой изоляцией и группами безопасности, ограничивающими связь только с авторизованными узлами обучения. Терабайты обучающих данных передаются непосредственно из Amazon Simple Storage Service (Amazon S3), что исключает необходимость в промежуточных слоях хранения, сохраняя при этом высокую пропускную способность.