AWS собрала отказоустойчивое обучение на EKS с NVRx: 99%+ эффективности и восстановление за секунды

AWS 16 сентября 2026 года опубликовала техническое руководство о том, как собрать отказоустойчивое распределённое обучение на Amazon EKS с NVIDIA Resiliency Extension (NVRx) в PyTorch Fully Sharded Data Parallel (FSDP). Авторы материала — Аравинд Нилакантан и Шрея Гупта, уровень Advanced (300).

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

Цена одного сбоя разобрана по шагам: падение одного GPU поднимает таймауты NVIDIA Collective Communication Library (NCCL) до здоровых воркеров, поды падают и перезапускаются несинхронно, а кластер сжигает дорогие GPU-часы, не делая прогресса в обучении.

Две поломки производительности: чекпойнты и восстановление

Вторая проблема — синхронные чекпойнты: каждая запись блокирует все ранги на ввод-вывод. На размерах кластера из публикации это съедало до 40% общего времени работы.

NVRx закрывает обе: асинхронная запись накладывает ввод-вывод на обучение, in-process restart восстанавливает работу за секунды без обращения к жизненному циклу контейнера, а in-job restart через `ft_launcher` автоматически поднимает воркеры после жёстких падений.

Библиотека ставится как обычный Python-пакет командой `pip install nvidia-resiliency-ext` — без собственных ядер, без форка PyTorch и без перекомпиляции; примитивы подключаются в существующий скрипт FSDP обычными импортами, а код модели и обучения остаётся нетронутым. Каждую возможность можно принять независимо от других.

Асинхронные чекпойнты

`TorchAsyncCheckpoint` заменяет `torch.save` на вызов `async_save()`, который передаёт state dict фоновому процессу и немедленно возвращает управление; парный вызов `finalize_async_save()` фиксирует предыдущую запись.

В связке с FSDP `LOCAL_STATE_DICT` каждый ранг пишет свой шард напрямую — без all-gather и без узкого места на нулевом ранге.

In-process restart

`inprocess.Wrapper` оборачивает функцию обучения и владеет циклом перезапуска: исключение внутри обёрнутой функции или зависание, которое обнаружил watchdog, приводит к повторному входу вместо падения процесса.

В конструкторе собираются четыре части: таймауты (soft, hard, barrier и completion), проверки здоровья `CudaHealthCheck` и `FaultCounter`, контроллер `RetryController`, ограничивающий общее число перезапусков и задающий минимальный живой размер мира, и назначение рангов `ActivateAllRanks` вместе с `ShiftRanks`, которое подтягивает выживших влево, чтобы мир оставался непрерывным.

Под этим слоем `ProgressWatchdog` через `Py_AddPendingCall` расставляет метки времени между инструкциями байт-кода, поэтому зависание внутри NCCL, будучи расширением на C, всё равно обнаруживается. `MonitorThread` опрашивает inprocess TCPStore и поднимает в главном потоке `RankShouldRestart`, после чего `AbortTorchDistributed` собирает трассы Flight Recorder, абортит бэкенды NCCL и уничтожает process group. Python-процесс остаётся живым — пересобирается только распределённая группа процессов.

In-job restart через ft_launcher

Запуск переключается с `torchrun` на `ft_launcher`: он понимает те же флаги рандеву и добавляет свои — `--max-restarts=20`, `--ft-restart-policy=any-failed`, `--ft-rank-heartbeat-timeout=900`, `--ft-initial-rank-heartbeat-timeout=1200`, `--monitor-interval=5`.

Внутри скрипта один раз после инициализации распределённого режима создаётся `RankMonitorClient`, который отправляет heartbeat на каждом шаге. Таймаут heartbeat следует держать выше самого длинного легитимного интервала между сигналами, а начальный таймаут в 1200 секунд (20 минут) рассчитан на первую загрузку модели.

На стороне лаунчера `RankMonitorServer` для каждого ранга считает интервалы между heartbeat. По таймауту выжившие получают SIGTERM, отстающие — SIGKILL, память GPU освобождается, происходит повторное рандеву и запуск новых воркеров, которые поднимаются с последнего чекпойнта. Частота чекпойнтов напрямую ограничивает объём потерянной работы.

Слои закрывают разные классы отказов

Публикация разводит слои по масштабу аварии: in-process ловит то, что укладывается в один Python-процесс — транзиентные исключения и видимые watchdog зависания NCCL.

`ft_launcher` ловит то, что убивает процесс или вешает его на уровне операционной системы: SIGKILL, OOM и дедлоки ниже Python. Потерю узла обрабатывает планировщик кластера, а асинхронные чекпойнты ортогональны восстановлению — они сочетаются с любым слоем и сокращают объём потерянной работы.

Результаты: 99%+ эффективности и секунды на восстановление

Замеры выполнены на обучении LLaMA-3.1-8B с FSDP. Асинхронные чекпойнты держат 99%+ эффективности обучения на всех масштабах — 99.2% на 2 узлах и 99.8% на 8 узлах, то есть от 16 до 64 GPU H100, тогда как синхронный вариант ровно держится на 57–61%.

Причина в том, что оверхед чекпойнта упирается в хранилище: запись в FSx Lustre занимает около 275 секунд независимо от того, 16 GPU или 64, а асинхронный режим прячет эту задержку, накладывая её на следующий отрезок обучения.

Свип по частоте чекпойнтов на 8 узлах усиливает разрыв: при интервале в 1000 шагов — 99.8% против 60.3%, а при 100 шагах синхронный вариант падает до 14.7%, тогда как асинхронный снижается мягче, до 29.6%, оставаясь вдвое эффективнее. Точка перехода — там, где разрыв между чекпойнтами (около 280 секунд на 100 шагов) подходит к времени записи в FSx (около 275 секунд).

Практический вывод для команд в публикации такой: чекпойнтить редко, чтобы не платить за блокировку, и терять часы прогресса при сбое больше не обязательно — асинхронная запись позволяет писать часто с почти нулевым оверхедом.

Сравнение механизмов восстановления

Эксперименты по восстановлению AWS провела на двух узлах p5.48xlarge, то есть на 16 GPU H100, обучая LLaMA-3.1-8B с FSDP, с чекпойнтом каждые 500 шагов на прогоне в 2000 шагов и пятью детерминированными сбоями на каждый прогон.

In-process restart даёт 31% goodput обучения и 87% инфраструктурного goodput, восстанавливаясь примерно за 10 секунд без единого перезапуска контейнера. `ft_launcher` показывает 25.5% и 85.9% при 17 секундах на сбой.

Базовый перезапуск средствами Kubernetes выглядит катастрофически: 11.5% training goodput, 35.8% инфраструктурного и 270 секунд на сбой. Причина — каскад: один ранг падает, выживший упирается в 60-секундный таймаут NCCL, оба пода перезапускаются несинхронно, что запускает CrashLoopBackOff и дополнительные циклы таймаутов до возобновления обучения.

Ставка на масштабе объяснена так: при сотнях и тысячах GPU аппаратные сбои случаются каждые несколько часов, и восстановление в 4+ минуты на сбой означает, что кластер с 3–4 сбоями в час тратит на восстановление больше времени, чем на само обучение.

Воспроизводимость и уборка

Чтобы сравнивать механизмы в одинаковых условиях, в публикации используется детерминированная инъекция сбоев с seeded-RNG: точное число сбоев на фиксированных шагах и рангах — `--fault_count=5 --fault_seed=42 --fault_types=exception,hang --fault_type_weights=0.6,0.4`.

Задания разворачиваются тонкой обёрткой над `kubectl`, а тип инстанса, число GPU и количество устройств Elastic Fabric Adapter (EFA) задаются переменными окружения, поэтому один и тот же код и манифесты работают на разных GPU (p5, p4de) сменой одного конфигурационного файла.

Полный воспроизводимый пример с тестовым сценарием лежит в репозитории `awslabs/awsome-distributed-ai`.

Отдельно AWS напоминает про порядок уборки: объекты Kubernetes, которые монтируют FSx, удаляются до самой файловой системы, иначе удаление блокируется финализаторами тома, а файловая система SCRATCH_2 удаляется без финальной резервной копии.