Cómo TGS y AWS Redujeron el Entrenamiento del Modelo de IA Sísmica de Seis Meses a Cinco Días

TGS вместе с AWS показала редкий для крупных ИИ-проектов результат: обучение сейсмической foundation-модели, которое раньше занимало около шести месяцев, удалось сократить до пяти дней. Одновременно команда увеличила размер контекстного окна, то есть объем трехмерных геологических данных, который модель может анализировать за один проход. Для компаний, работающих с разведкой недр, это означает более быстрые итерации и более полную картину подземных структур.

TGS — поставщик геонаучных данных для энергетического сектора. Компания использует сейсмические foundation-модели для анализа сложных 3D-массивов, чтобы находить геологические структуры, важные для разведки и добычи. Базовая архитектура модели построена на Vision Transformer и схеме обучения Masked AutoEncoder.

Главная проблема здесь не только в размере самой модели, но и в характере данных: сейсмические объемы состоят из миллиардов точек, хранятся в специализированных форматах и требуют постоянной подачи на GPU без простоев. Совместно с AWS Generative AI Innovation Center компания перенесла обучение на Amazon SageMaker HyperPod и собрала кластер из 16 узлов Amazon EC2 P5. Каждый узел включал 8 GPU NVIDIA H200 с 141 ГБ памяти HBM3e, 192 vCPU, 2 ТБ оперативной памяти и сеть EFAv3 со скоростью 3200 Гбит/с.

В сумме это 128 GPU. По данным AWS, такая конфигурация дала почти линейное масштабирование: эффективность параллельной работы при росте от одного узла до 16 держалась на уровне примерно 90–95%. Отдельный фокус был на пайплайне данных.

Вместо классической схемы с Amazon FSx for Lustre команда выбрала прямой стриминг из Amazon S3. Причина проста: при росте кластера S3 позволяет каждому узлу добавлять собственную пропускную способность, тогда как общая файловая система быстрее становится узким местом. Обучающий датасет TGS хранится в разработанном компанией формате MDIO на базе Zarr, оптимизированном для больших научных данных в облаке.

После настройки многопоточной загрузки и prefetching каждый узел стабильно получал 4–5 ГБ/с, а весь кластер — 64–80 ГБ/с. Дополнительно это снизило затраты на storage-инфраструктуру более чем на 90%. Для распределенного обучения команда сравнила несколько подходов: DeepSpeed ZeRO-2, ZeRO-3 и FSDP2.

Лучший баланс между скоростью и экономией памяти показал ZeRO-2: 1 974 сэмпла в секунду против 1 833 у FSDP2 и 869 у ZeRO-3. Это важно, потому что в таких задачах максимальная экономия памяти не всегда дает лучший итоговый результат: если обмен между GPU становится слишком дорогим, throughput резко падает. В этом проекте ставку сделали не на теоретически самый плотный вариант, а на ту конфигурацию, которая реально быстрее прогоняет обучение в продакшн-сценарии.

Второе ключевое достижение связано с расширением контекстного окна. Для сейсмических моделей это напрямую влияет на качество анализа: чем больший объем породы модель видит за один проход, тем лучше она улавливает и локальные детали, вроде небольших трещин, и крупные паттерны — например, разломные системы на масштабе всего бассейна. За счет context parallelism и адаптации ring attention под ViT-архитектуру для 3D-данных максимальный размер входа вырос с 640 x 640 x 1 024 до 1 536 x 1 536 x 2 048 вокселей.

Длина контекста увеличилась со 102 400 до 1,17 млн токенов, а анализируемый объем — примерно в 4,5 раза. Практический смысл кейса в том, что TGS теперь может обновлять модели не раз в полгода, а фактически еженедельно, быстрее подключать новые геологические данные и давать клиентам более широкий контекст анализа. Для AWS это еще один показательный пример того, что специализированные foundation-модели в науке и индустрии упираются не только в архитектуру модели, но и в правильную организацию данных, сети и распределенного обучения.

Для рынка в целом сигнал понятный: узкоспециализированный ИИ начинает выигрывать там, где раньше доминировали долгие вычислительные циклы и слишком дорогая инфраструктура.