NVRx (NVIDIA Resiliency Extension) - это pip-устанавливаемый Python-пакет, который дополняет PyTorch FSDP тремя примитивами отказоустойчивости: асинхронным сохранением чекпоинтов через TorchAsyncCheckpoint, перезапуском обучения внутри того же процесса при мягких сбоях и обёрткой ft_launcher для жёстких падений вроде SIGKILL или OOM. На кластере Amazon EKS из инстансов p5.48xlarge с GPU NVIDIA H100, сетью EFA и хранилищем FSx for Lustre это даёт две измеримые величины: эффективность обучения держится выше 99% вместо 57-61% у синхронных чекпоинтов, а восстановление после мягкого сбоя занимает около 10 секунд без пересоздания контейнеров.
Боль знакома каждому, кто гонял многоузловое обучение: запись чекпоинта останавливает вычисления на всех рангах, а после падения процесса кластер заново проходит путь от планирования пода до инициализации NCCL, пока восемь H100 на узле простаивают. NVRx не заменяет PyTorch и не трогает математику обучения, он добавляет слой управления записью состояния и сбоями.
Дальше разберу, из чего состоит пакет, как выглядит кластер на EKS под такую нагрузку, какие цифры дают замеры на H100, где лежат ограничения и как проверить отказоустойчивость на своём стенде.
Что такое NVRx и какие проблемы он решает
FSDP шардит параметры модели, состояния оптимизатора и градиенты по всем GPU кластера. Чекпоинт для модели на десятки миллиардов параметров - это не один файл, а набор шардов общим объёмом в сотни гигабайт, которые нужно согласованно собрать и записать.
Стандартный путь (torch.save или torch.distributed.checkpoint) делает это синхронно: пока данные уходят на диск, ни один ранг не считает шаги. Второе узкое место - восстановление. Процесс упал, torchrun завершился, Kubernetes пересоздал под, кластер заново поднял NCCL-коммуникации и загрузил последнюю точку возобновления. Простой измеряется минутами, а на восьмиузловой конфигурации это десятки GPU-часов в месяц.
NVRx закрывает обе болевые точки на прикладном уровне. Пакет ставится через pip под именем nvidia-resiliency-ext, не требует патчей ядра PyTorch и не диктует, как писать модель. Это набор утилит вокруг вашего тренировочного скрипта, а не отдельный фреймворк обучения.
Три примитива NVRx: TorchAsyncCheckpoint, in-process restart, ft_launcher
TorchAsyncCheckpoint сохраняет состояние модели и оптимизатора асинхронно. Ранги складывают тензоры в буферы, а запись на диск идёт в фоне, пока цикл обучения продолжает считать шаги. Класс работает с шардированными состояниями FSDP и с распределёнными чекпоинтами PyTorch, поэтому собирать полный state dict на одном ранге не нужно.
In-process restart (в материалах NVIDIA встречается как in-job restart) перехватывает мягкие сбои: исключение в Python, зависание NCCL по таймауту, деградацию одного из рангов. Процесс-супервизор внутри контейнера останавливает обучение, поднимает его заново и продолжает с последнего чекпоинта. Под не пересоздаётся, образ не перезаливается, GPU остаются зарезервированными за вами.
ft_launcher - обёртка запуска в стиле torchrun с той разницей, что она следит за воркерами и перезапускает их при жёстких сбоях: SIGKILL, убийство по OOM, аварийное завершение процесса. У лаунчера есть параметр --max-restarts, который ограничивает число попыток и не даёт получить бесконечный цикл падений.
Разделение ролей простое: мягкий сбой ловится изнутри процесса и стоит секунд, жёсткий убивает процесс целиком, и поднимать его должен внешний супервизор. Асинхронные чекпоинты отвечают за то, чтобы точка возобновления оставалась свежей для обоих механизмов.
Почему стандартные чекпоинты и восстановление тормозят обучение
Эффективность обучения в таких замерах считают как долю времени, когда GPU реально выполняют вычисления, от общего времени работы. Синхронная запись роняет её до 57-61%: все ранги ждут, пока шарды уйдут в хранилище, и только потом возобновляют шаги. Чем чаще вы сохраняетесь, тем сильнее просадка.
Асинхронный режим удерживает показатель выше 99%, потому что запись перекрывается вычислениями. На коротких интервалах между чекпоинтами разница становится драматичной: при сохранении каждые 10 минут синхронная схема отдаёт диску заметную долю бюджетного времени кластера.
Восстановление устроено так же невыгодно. Пересоздание пода добавляет к простою планирование, скачивание образа, подъём сетевых интерфейсов и повторную инициализацию NCCL. In-process restart пропускает все эти этапы: контейнер уже запущен, EFA-интерфейсы настроены, остаётся перечитать состояние с Lustre. Отсюда примерно 10 секунд вместо минут.
Архитектура кластера Amazon EKS для обучения с NVRx
Целевая конфигурация выглядит так: Amazon EKS как оркестратор, GPU-инстансы p5.48xlarge, сеть EFA для межузлового обмена и FSx for Lustre под чекпоинты и данные. EKS отвечает за размещение подов и перезапуск на уровне кластера, NVRx добавляет отказоустойчивость приложения.
Выбор инстансов и сети: p5.48xlarge и EFA
p5.48xlarge несёт 8 GPU NVIDIA H100 80 ГБ SXM, около 192 vCPU, порядка 2 ТБ оперативной памяти и 8 локальных NVMe-дисков по 3,84 ТБ. Для FSDP это удобная единица масштабирования: 8 GPU внутри одного узла общаются по NVLink, а межузловой трафик уходит в сеть.
Именно здесь нужна EFA (Elastic Fabric Adapter). У инстанса 32 сетевых карты и агрегированная пропускная способность до 3200 Гбит/с, а EFA даёт обход ядра и низкую задержку для коллективных операций. NCCL использует её через провайдер libfabric, поэтому all-gather и reduce-scatter в FSDP не упираются в TCP-стек.
В EKS EFA подключается через aws-efa-k8s-device-plugin. В манифесте пода запрашивается ресурс vpc.amazonaws.com/efa, а в окружении выставляются переменные вроде FI_PROVIDER=efa. Без этого плагина многоузловое обучение на 4-8 узлах начинает ждать сеть вместо GPU, и выигрыш от H100 частично теряется.
Хранилище для чекпоинтов: FSx for Lustre
Асинхронная запись снимает блокировку вычислений, но не отменяет физику: сотни гигабайт на чекпоинт должны куда-то уходить с достаточной скоростью. Пропускная способность FSx for Lustre масштабируется вместе с объёмом файловой системы, у Persistent SSD это порядка 1000 МБ/с на каждый терабайт, чего хватает, чтобы фоновая запись не догоняла обучение.
Файловая система подключается к подам как PVC через CSI-драйвер и поддерживает режим ReadWriteMany, то есть все воркеры пишут в одну точку монтирования. Привязка репозитория к бакету S3 даёт долговременное хранение: горячие чекпоинты лежат в Lustre, архивные уезжают в S3 задачами репозитория.
Альтернативы проигрывают по разным причинам. Отдельный том EBS упирается в потолок пропускной способности одного тома, и на сотнях гигабайт запись становится узким местом. S3 напрямую неудобен из-за задержки на объект и стоит денег за каждую операцию, что плохо ложится на частую запись шардов.
Установка и настройка NVRx в контейнерах EKS
Пакет добавляется в образ одной строкой, лучше отдельным слоем после установки PyTorch, чтобы кэш Docker не инвалидировался при каждом обновлении кода:
FROM nvcr.io/nvidia/pytorch:xx.xx-py3
RUN pip install --no-cache-dir nvidia-resiliency-ext
COPY train_fsdp.py /workspace/train_fsdp.py
ENTRYPOINT ["ft_launcher", "--nnodes=4", "--nproc-per-node=8", \
"--rdzv-backend=c10d", "--rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT", \
"--max-restarts=3", "/workspace/train_fsdp.py"]Точка входа пода - ft_launcher, а не python напрямую: тогда жёсткие сбои переживает сам контейнер. Асинхронные чекпоинты требуют свежей версии PyTorch с поддержкой распределённых сохранений, поэтому версии torch и nvidia-resiliency-ext стоит фиксировать в образе, а не подтягивать последние при каждом билде.
Асинхронные чекпоинты: как удержать эффективность обучения выше 99%
Механика простая: ранги копируют состояние в промежуточные буферы, после чего фоновый поток отправляет данные в хранилище. Пока идёт запись, обучение продолжается, а следующий барьер по чекпоинту ставится только тогда, когда предыдущая запись действительно завершилась.
Сравнение с синхронными чекпоинтами: цифры и графики
Замеры проводились на кластере от 2 до 8 узлов p5.48xlarge с H100, сетью EFA и хранилищем FSx for Lustre, обучение шло на PyTorch FSDP.
| Параметр | Синхронные чекпоинты | Асинхронные (NVRx) |
|---|---|---|
| Эффективность обучения, 2-8 узлов | 57-61% | выше 99% |
| Блокировка цикла обучения | на всё время записи | нет, запись идёт в фоне |
| Восстановление после мягкого сбоя | пересоздание пода, минуты | in-process restart, около 10 секунд |
Разброс 57-61% связан с числом узлов: чем больше рангов участвует в синхронизации перед записью, тем длиннее пауза. Асинхронная схема держится выше 99% во всём диапазоне, потому что пауза на барьер не растягивается на время дискового ввода-вывода.
Цифры бенчмарков всегда зависят от конфигурации, и это стоит проверять на своём стенде. Хороший пример того, как одна и та же модель даёт разные результаты в разных условиях, - замеры DeepSeek V4.1 Flash на восьми A40: разница между квантизациями и настройками рантайма там больше, чем между железками.
Как подключить TorchAsyncCheckpoint к PyTorch FSDP
Замена синхронного сохранения сводится к нескольким строкам. Вместо torch.save используется объект TorchAsyncCheckpoint, которому передаются модель, оптимизатор и каталог в Lustre:
from nvidia_resiliency_ext.checkpointing.async_ckpt.torch_ckpt import TorchAsyncCheckpoint
ckpt = TorchAsyncCheckpoint("torch", checkpoint_dir="/mnt/fsx/ckpt")
# внутри цикла обучения, после optimizer.step()
if step % ckpt_interval == 0:
ckpt.save({"model": model, "optim": optim, "step": step},
checkpoint_id=f"step_{step}")
# перед завершением обучения дождитесь, пока фоновые записи закроются
# и только потом выходите из процесса
checkpoint = ckpt.load({"model": model, "optim": optim}, checkpoint_id="step_1000")Точные имена методов и сигнатуры сверяйте с документацией пакета: API молодое и между версиями меняется. Практические нюансы важнее синтаксиса. Асинхронный режим держит копию состояния в буферах, поэтому растёт потребление оперативной памяти узла и появляется нагрузка на CPU-потоки, которые пишут данные. Если хранилище не успевает, очередь записей копится и память заканчивается раньше, чем диск. Интервал между чекпоинтами стоит подбирать так, чтобы запись гарантированно завершалась до следующего вызова.
Восстановление после сбоев: in-process restart и ft_launcher
Сбои в распределённом обучении делятся на два класса. Мягкие: исключение в коде, зависание коллективной операции, деградация ранга, из-за которой шаг не завершается. Жёсткие: процесс убит по SIGKILL или OOM, узел отвалился, GPU выдала ошибку, несовместимую с продолжением работы. Под каждый класс в NVRx есть свой механизм.
In-process restart: восстановление за ~10 секунд без перезапуска контейнеров
Супервизор внутри контейнера отслеживает состояние воркеров. При исключении или срабатывании таймаута NCCL он останавливает обучение, разрывает process group и запускает тренировочный цикл заново в том же контейнере, подтягивая последний чекпоинт из Lustre. Планировщик Kubernetes в этом процессе не участвует, образ повторно не скачивается, GPU не освобождаются.
Отсюда примерно 10 секунд в тестовой конфигурации. Основное время уходит на загрузку состояния и повторную инициализацию коммуникаций, а не на организационные накладные. Для моделей с большим числом параметров загрузка состояния добавит времени, поэтому цифру стоит воспринимать как порядок величины, а не как константу.
Ft_launcher: обработка жёстких сбоев (SIGKILL, OOM)
Когда процесс убит системным OOM-killer или сигналом, изнутри ловить нечего: обучение прекращается мгновенно. Здесь работает ft_launcher. Он выступает родителем для воркеров, замечает их смерть и перезапускает группу заново через rendezvous, используя последний чекпоинт. Параметр --max-restarts ограничивает число таких попыток, что защищает кластер от бесконечного цикла падений.
Если падает сам контейнер или узел, дальше подключается EKS: под перепланируется на здоровом инстансе, и ft_launcher внутри нового контейнера снова собирает группу. Такое восстановление занимает уже минуты, потому что включает планирование и старт контейнера.
Как слои восстановления сочетаются с асинхронными чекпоинтами
| Класс сбоя | Примеры | Механизм | Что происходит с подом |
|---|---|---|---|
| Мягкий | исключение в Python, зависание NCCL | in-process restart | под живёт, пауза около 10 с |
| Жёсткий на уровне процесса | SIGKILL, OOM | ft_launcher | под живёт, воркеры перезапускаются |
| Жёсткий на уровне узла | отказ инстанса, вытеснение пода | EKS и ft_launcher | новый под, минуты |
Асинхронные чекпоинты связывают эти слои воедино: любой из механизмов восстанавливается из последней сохранённой точки. Чем свежее чекпоинт, тем меньше потерянных шагов, и именно поэтому асинхронность важна не только для эффективности обучения, но и для цены сбоя.
Бенчмарки и тестирование отказоустойчивости
Методика: детерминированная инъекция сбоев
Чтобы сравнить механизмы восстановления, сбой нужно вызывать в предсказуемый момент. Схема такая: в sidecar-контейнере или отдельном скрипте задаётся номер шага, на котором обучение ломается. Дальше выбирается тип отказа.
- Жёсткое убийство: kill -9 по PID воркера. Работает, если в поде включён shareProcessNamespace, тогда sidecar видит процессы обучения.
- OOM: контейнер-«пожиратель» памяти в том же поде или cgroup-лимит, который срабатывает на заданном шаге.
- Зависание NCCL: рангу блокируется сетевой обмен или процесс ставится на паузу, коллективная операция не завершается и срабатывает таймаут.
Замеряются четыре отметки времени: момент сбоя, момент обнаружения, момент перезапуска обучения и первый успешный шаг после восстановления. Разница между третьей и четвёртой отметками даёт время простоя, а разница между четвёртой и первой, умноженная на длительность шага, показывает потерянный прогресс. Каждый сценарий повторяется несколько раз, чтобы получить медиану, а не случайную удачную попытку. Такой подход к измерениям похож на любые честные сравнения: важны одинаковые условия и повторяемость, как в разборе сравнения nInfer, vLLM и llama.cpp на одной RTX 5090.
Результаты на H100: эффективность и время восстановления
Сводка по тестовой конфигурации от 2 до 8 узлов p5.48xlarge:
| Метрика | Значение |
|---|---|
| Эффективность обучения с асинхронными чекпоинтами | выше 99% |
| Эффективность обучения с синхронными чекпоинтами | 57-61% |
| Восстановление in-process restart | около 10 секунд |
| Восстановление после жёсткого сбоя через ft_launcher | дольше, зависит от типа отказа |
Результаты зависят от размера модели, интервала чекпоинтов, скорости Lustre и стабильности сети. На двух узлах выигрыш от асинхронности меньше, чем на восьми, потому что синхронизация дешевле. Оценивайте свою конфигурацию отдельно: смысл NVRx раскрывается там, где чекпоинты частые, а кластер большой.
Ограничения и кому подходит NVRx
Когда NVRx не нужен или может навредить
Пакет заточен под PyTorch FSDP. Если вы обучаете модель на DDP, в TensorFlow или JAX, он не поможет. Если обучение идёт на одном узле и чекпоинты делаются раз в несколько часов, синхронная запись почти не влияет на общую эффективность, а вы получите дополнительный слой абстракций и новые точки отказа.
Без быстрого хранилища асинхронность превращается в проблему: очередь записей растёт, буферы занимают оперативную память, и узел может упасть по OOM раньше, чем закончится обучение. На слабом хранилище честнее оставить редкие синхронные чекпоинты.
Наконец, часть возможностей пакета молодая. Механизм перезапуска внутри процесса NVIDIA развивает активно, поэтому на продакшене его стоит включать после собственных тестов с инъекцией сбоев, а не по умолчанию. Для локальных сценариев с двумя потребительскими картами, как в схеме запуска Qwen 3.8 27B на двух RTX 3090, инструмент попросту не про эту задачу: там нет ни многоузлового обучения, ни EFA.
Сравнение с альтернативами: torch.distributed.checkpoint, AWS SageMaker
| Решение | Что даёт | Чего не хватает |
|---|---|---|
| torch.distributed.checkpoint | шардированное сохранение, в свежих версиях PyTorch есть асинхронный вариант | нет слоя перезапуска и контроля зависаний |
| NVRx | асинхронные чекпоинты плюс восстановление на уровне процесса | работает с FSDP, требует быстрого хранилища и памяти под буферы |
| AWS SageMaker | управляемое обучение с автоматическим перезапуском задач | меньше контроля над стеком и другой уровень абстракции |
Штатный DCP закрывает часть задачи и остаётся разумным выбором, если сбои редкие. NVRx выигрывает там, где нужно сократить часы простоя, а не только ускорить запись. Управляемые сервисы экономят силы команды, но забирают контроль над конфигурацией сети, хранилища и версий библиотек.
Практические шаги: как добавить NVRx в свой пайплайн
- Проверьте инфраструктуру: EKS с GPU-узлами p5.48xlarge, установленный aws-efa-k8s-device-plugin, PVC на FSx for Lustre.
- Соберите образ с фиксированными версиями torch и nvidia-resiliency-ext, добавьте пакет отдельным слоем.
- Переведите код на PyTorch FSDP, если он ещё на DDP: NVRx требует шардирования.
- Замените синхронное сохранение на TorchAsyncCheckpoint и подберите интервал так, чтобы запись успевала завершиться.
- Сделайте точкой входа ft_launcher с параметром --max-restarts и настройте rendezvous.
- Прогоните сценарии с детерминированной инъекцией сбоев: SIGKILL, OOM, зависание NCCL.
- Настройте мониторинг времени восстановления и эффективности обучения, чтобы видеть эффект на длинной дистанции.
Чек-лист для запуска NVRx на EKS
- Кластер EKS с GPU-инстансами p5.48xlarge.
- Установлен aws-efa-k8s-device-plugin, в поде запрошен ресурс vpc.amazonaws.com/efa.
- FSx for Lustre подключён как PVC с режимом ReadWriteMany, настроена привязка к S3 для архива.
- В Docker-образ добавлен nvidia-resiliency-ext.
- Обучение идёт на PyTorch FSDP.
- Сохранение чекпоинтов обёрнуто в TorchAsyncCheckpoint.
- Точка входа контейнера - ft_launcher с ограничением на число перезапусков.
- Мониторинг собирает время восстановления, эффективность и объём потерянных шагов.
- Проведены тесты с инъекцией сбоев на двух узлах перед масштабированием.
Начните с двух узлов и одного типа сбоя: так вы за пару дней поймёте, сколько времени реально экономит асинхронный чекпоинт на вашей модели и хранилище. Если восстановление укладывается в десятки секунд, а эффективность держится у верхней границы, масштабируйте конфигурацию на нужное число узлов и добавляйте сценарии с отказами железа.