NVIDIA NVRx и Amazon EKS: как асинхронное чекпоинтинг сокращает простои при обучении ИИ
AWS Machine Learning · оригинал
Материал подготовлен автоматизированной редакционной системой. Факты можно сверить по указанному первоисточнику.
Команда AWS продемонстрировала интеграцию NVIDIA Resiliency Extension (NVRx) в PyTorch FSDP на Amazon EKS, что позволяет решать проблемы блокировок ввода-вывода и сбоев оборудования. Тесты на GPU H100 показали, что асинхронное сохранение чекпоинтов повышает эффективность обучения до 99%, а восстановление после сбоев занимает секунды вместо минут.
Ценность: Масштабное обучение моделей ИИ уязвимо к сбоям оборудования и медленному сохранению данных. Представленное решение устраняет два ключевых источника потерь времени: блокирующие операции I/O и длительные перезапуски контейнеров, что критически важно для экономии дорогостоящих GPU-часов.
Масштабные задачи обучения нейросетей, длящиеся часы или дни на десятках узлов, статистически неизбежно сталкиваются с сбоями: от сетевых разрывов до ошибок памяти. Даже единичный отказ GPU может вызвать каскадный эффект: тайм-ауты библиотеки NCCL распространяются на здоровые узлы, поды перезапускаются вне синхронизации, и кластер сжигает дорогие GPU-часы без прогресса. Дополнительной проблемой является синхронное сохранение чекпоинтов, которое блокирует все ранги во время записи на диск. По данным AWS, в крупных кластерах это съедало до 40% общего времени обучения.
Для решения этих проблем команда AWS интегрировала NVIDIA Resiliency Extension (NVRx) в PyTorch Fully Sharded Data Parallel (FSDP) на базе Amazon EKS. NVRx представляет собой устанавливаемый через pip Python-слой, который добавляет механизмы отказоустойчивости без необходимости перекомпиляции или изменения ядра PyTorch. Решение опирается на инфраструктуру EKS с инстансами p5.48xlarge (8 GPU H100 каждый), высокоскоростной сетью EFA и общим хранилищем Amazon FSx for Lustre для персистентности данных.
Ключевым элементом является асинхронное чекпоинтинг через класс TorchAsyncCheckpoint. Вместо блокирующего torch.save состояние модели передается фоновому процессу, позволяя основному потоку продолжать обучение. Каждый ранг записывает свой шард независимо, что исключает узкое место на ранге 0. Бенчмарки на LLaMA-3.1-8B показали, что асинхронный подход сохраняет эффективность обучения выше 99% при масштабировании от 2 до 8 узлов (16–64 GPU). В то же время синхронное сохранение стагнирует на уровне 57–61% эффективности, так как время записи на FSx Lustre (~275 секунд) остается постоянным независимо от количества GPU.
Вторая часть решения — многоуровневое восстановление после сбоев. Механизм in-process restart перехватывает «мягкие» сбои (исключения или зависания NCCL), не убивая Python-процесс. Он выполняет проверки здоровья GPU и сети, пересобирает группу процессов и возобновляет обучение из последнего чекпоинта. Для «жестких» сбоев, таких как SIGKILL или OOM, используется загрузчик ft_launcher, который отслеживает heartbeat-сигналы и при потере связи перезапускает рабочих процессы внутри того же задания Kubernetes.
Эксперименты по введению детерминированных сбоев на 16 GPU H100 выявили значительные различия в эффективности. NVRx in-process restart обеспечил 31% полезного времени обучения (training goodput) и восстановление за ~10 секунд. Загрузчик ft_launcher показал 25.5% goodput со временем восстановления 17 секунд. В качестве базового варианта использовался стандартный перезапуск Kubernetes, который дал лишь 11.5% goodput и требовал 270 секунд на восстановление из-за каскадных тайм-аутов и несогласованности подов.
Результаты демонстрируют, что при увеличении масштаба кластера, где сбои происходят каждые несколько часов, традиционные методы восстановления приводят к тому, что система тратит больше времени на перезапуск, чем на обучение. NVRx сокращает это окно до секунд, позволяя поддерживать высокую утилизацию GPU даже при высокой частоте сбоев. Асинхронное чекпоинтинг дополнительно позволяет сохранять состояние модели чаще, минимизируя потери прогресса при сбоях, без существенного влияния на скорость обучения.