AWS показала, как подключить NVIDIA Resiliency Extension, или NVRx, к распределённому обучению на PyTorch FSDP без форка PyTorch, собственных ядер и изменений модели. Это не очередная попытка ускорить вычисления на несколько процентов. NVRx сокращает время, когда дорогие GPU ничего не обучают: ждут записи чекпойнта или восстановления после падения одного из воркеров.
В конфигурациях AWS из 2–8 узлов p5.48xlarge, по восемь H100 80 GB в каждом, синхронное сохранение чекпойнтов занимало до 40% общего времени запуска. Причина проста: при каждом сохранении все процессы останавливаются и ждут завершения I/O.
NVRx заменяет torch.save на асинхронный вызов. Состояние передаётся фоновому процессу, а обучение продолжается. Каждый процесс сохраняет собственную часть состояния через FSDP LOCAL_STATE_DICT, поэтому не нужно собирать всю модель на одном процессе и создавать узкое место на rank 0. Перед следующим сохранением система завершает предыдущую запись.
Отказы разделили по тяжести
Вторая часть NVRx отвечает за восстановление. Для необработанных исключений и зависаний NCCL используется перезапуск внутри процесса. Расширение останавливает текущую группу процессов, проверяет GPU, NVLink и сетевые интерфейсы, заново объединяет исправные процессы и продолжает работу с последнего чекпойнта. Интерпретатор Python и CUDA allocator при этом сохраняются.
Для SIGKILL, остановки из-за нехватки памяти и зависаний на уровне ОС предусмотрен ft_launcher. Он следит за heartbeat каждого процесса, при сбое завершает оставшиеся процессы, освобождает память GPU и запускает воркеры заново внутри того же задания. Потерю целого узла по-прежнему должен обрабатывать оркестратор кластера.
Слои независимы: команда может внедрить только асинхронные чекпойнты или добавить восстановление для своего класса сбоев. NVRx устанавливается как pip-пакет, а его компоненты подключаются к FSDP-скрипту обычными импортами. AWS также опубликовала код для воспроизведения реализации.
Для бизнеса вывод довольно приземлённый: стоимость обучения определяется не только скоростью GPU, но и долей оплаченного простоя. Если многодневный запуск регулярно тормозит на сохранениях или целиком откатывается из-за одного воркера, NVRx можно проверить как отдельный инфраструктурный слой, не переписывая модель и основной код обучения.
Но цифру в 40% нельзя автоматически переносить на любой кластер. AWS тестировала конкретную связку EKS, H100, EFA и FSx for Lustre, причём загрузка чекпойнта на больших масштабах всё равно доминировала во времени восстановления. Сначала имеет смысл измерить собственные паузы на сохранение и потери от каскадных падений. Иначе отказоустойчивость легко превращается ещё в один компонент, который команда обслуживает ради красивой архитектурной схемы.