PyTorch 2.14: NVGEMM, nccl2 и живая отказоустойчивость

PyTorch 2.14: NVGEMM, nccl2 и живая отказоустойчивость

Вышел PyTorch 2.14, и это один из тех релизов, где за сухим списком изменений прячется смена философии. 2995 коммитов от 487 контрибьюторов со времён 2.13, но главное не количество. Fault tolerance впервые стал концепцией уровня c10d, а не деталью бэкенда: тренировочный джоб теперь может пережить падение ноды без перезапуска всего кластера. Для команд, которые гоняют обучение неделями на сотнях GPU, это меняет экономику процесса.

Релиз продолжает линию, которую команда ведёт всю серию 2.x: PyTorch из research-first фреймворка превращается в единую hardware-agnostic платформу для продакшн-обучения и инференса. Версия 2.12 добавила device-agnostic torch.accelerator.Graph API, 2.13 принесла FlexAttention на Apple Silicon и torchcomms для больших кластеров. В 2.14 эти нити доведены до зрелости: CuTeDSL-путь вырос в полноценный GEMM-backend NVGEMM, torchcomms переехал в дерево как backend nccl2, а динамические размерности тензоров стали декларативными.

NVGEMM: CUTLASS-ядра прямо в Inductor

Самое заметное изменение для производительности это NVGEMM, новый backend для матричных умножений в TorchInductor, построенный на официальном API cutlass.operators от NVIDIA. В версии 2.13 CuTeDSL мог сгенерировать standalone-ядро для GEMM, но всё, что шло следом (bias add, активация, rescale), оставалось отдельным ядром, которое заново читало результат из памяти. Теперь эпилоги сливаются в одно ядро так же, как это делают шаблоны Triton: bias add в addmm, цепочки pointwise-операций и даже редукции по результату GEMM, включая случаи, когда ядро возвращает и редуцированное значение, и полную выходную матрицу.

Fusion добрался и до низкоточных путей. Pointwise-работа после scaled GEMM теперь складывается внутрь ядра, а runtime global scale для формата NVFP4 применяется в эпилоге, а не отдельным умножением. Сгенерированные кандидаты NVGEMM конкурируют с Triton и ATen за mm, addmm и scaled_mm в рамках autotune, а сами слитые ядра кешируются на диск, так что повторная компиляция не начинается с нуля. Для включения достаточно добавить NVGEMM в max_autotune_gemm_backends. Потребуется nvidia-cutlass-dsl 4.6.0, а NVFP4-пути работают только на Blackwell. Эпилоги, которые backend выразить не может, откатываются на Triton без потери существующего fusion.

nccl2 и отказоустойчивость как первоклассная концепция

Распределённая часть релиза, пожалуй, важнее всего для больших тренировок. torchcomms, появившийся в 2.13, теперь в дереве исходников как новый c10d backend nccl2 за флагом USE_C10D_NCCL. Он реализует полный контракт Work поверх переиспользуемой абстракции NcclApi и приносит с собой односторонние окна, отказоустойчивость и suspend/resume offload памяти. Для тех, кому нужно старое ленивое поведение, есть совместимая обёртка nccl-lazy, которая строит per-peer P2P коммуникаторы по требованию.

Ключевой сдвиг в том, как обрабатываются сбои. Раньше при падении ранка в большом джобе стандартным восстановлением был снос process group и рестарт, что сбрасывало прогретое состояние по всему кластеру. Теперь Backend и ProcessGroup получили интерфейсы реконфигурации: группу можно пересобрать на месте, с abort-хуками и pre/post-хуками коллективов, проходящими через тот же путь. Gloo получил поддержку отказоустойчивости наравне с nccl2.

Рядом появились односторонние RMA-окна: ранк может читать или писать память соседа без встречного вызова с его стороны. Это ложится на нерегулярные паттерны доступа вроде lookup по эмбеддингам, передачи весов и роутинга экспертов в mixture-of-experts. Через nccl2 это раскрывает новые API ncclGet и ncclPut. Довершает картину backend-agnostic Flight Recorder: трассировка инцидентов теперь работает для любого backend, а не только для NCCL, что упрощает расследование зависаний в смешанных конфигурациях.

Приятный бонус, который не требует никаких изменений кода: Inductor теперь по умолчанию включает simple_overlap, перекрытие коммуникаций с независимыми вычислениями. Раньше это была opt-in опция, теперь распределённые тренировки, скомпилированные через Inductor, автоматически получают лучшую утилизацию GPU просто за счёт обновления версии.

Apple Silicon: нативная линейная алгебра

Линейная алгебра на MPS исторически опиралась на примитивы Apple MPSGraph или вовсе падала на CPU для всего, что сложнее базовых операций. Смешанные CPU/MPS переходы были постоянным источником тормозов в численном коде. В 2.14 эти дыры закрыты нативными Metal-ядрами. SVD, eigh и lstsq теперь работают через Jacobi-ядра для float32 и complex64, с откатом на CPU для float64 (у Metal нет double) и для маленьких матриц, где накладные расходы на запуск GPU не окупаются. Заодно зажглись зависящие от них matrix_rank, pinv, cond и norm.

Cholesky получил более быстрый алгоритм панельной факторизации с trailing update на базе matmul2d, что даёт примерно 1.2–2.8x в зависимости от размера, плюс исправление корректности для complex dtypes, которые раньше могли молча выдавать неверный результат. lu_factor и lu_solve переехали с MPSMatrixDecompositionLU на рукописные Metal-ядра, и здесь цифры самые впечатляющие: авторы замерили больше 100x на маленьких батчах матриц и 2–9x на больших одиночных. Добавлены geqrf и linalg.polar с backward-проходом, а matrix_exp доступен на MPS, хотя обгоняет CPU только примерно от 512×512.

Миграция операторов с MPSGraph на Metal продолжилась и вне линейной алгебры: index_add, index_select, argmin, argmax, conv3d, median, linspace, arange и ещё с десяток операторов. Редукции получили отдельный переписанный путь из пяти частей, включая split-K для argmax и argmin.

Компилятор: torch.switch и декларативные dynamic shapes

torch.cond умел выразить только двухстороннюю ветку, так что n-сторонний dispatch приходилось писать вложенными условиями. Это раздувало трасcируемый граф и затемняло намерение. Новый higher-order op torch.switch делает multi-way ветвление по индексу первоклассной операцией, с дедупликацией lifted-аргументов в Dynamo, чтобы общие операнды не поднимались заново для каждой ветки. Практический бенефициар очевиден: mixture-of-experts архитектуры, где вложенный torch.cond был реальным барьером. А torch.while_loop теперь можно захватить в CUDA graph.

Вторая история это @dynamic_spec из torch.fx.experimental.dynamic_spec. До сих пор указать, какие размерности входов меняются, приходилось по-разному для каждой точки входа: словарь dynamic_shapes для torch.export, грубый флаг dynamic для torch.compile, глобальный режим трейсинга для make_fx. И в каждом случае объявление жило на call site, далеко от модели, которую оно описывает. Теперь размерность именуется один раз (ShapeVar("batch", min=2, max=128)), переиспользуется между входами, поддерживает производные вроде batch * 2 и допущения вроде batch % 2 == 0. Все три точки входа принимают спецификацию под единым ключом dynamic_shapes, а декоратор @dynamic_spec прикрепляет её прямо к forward модуля.

У решения есть осознанная цена: объявленные так размерности становятся unbacked symbols, и компилятор не может тихо специализироваться на том batch size, который случайно попался при трейсе. Ветвления, зависящие от формы, теперь всплывают как data-dependent ошибка вместо guard с перекомпиляцией. API экспериментальный, поддержка make_fx ограничена tracing_mode="fake".

Производительность под капотом: что получают бесплатно

Отдельного разговора заслуживают изменения, которые не требуют от пользователя вообще ничего, кроме обновления версии. Inductor научился агрессивнее батчировать мелкие GPU-ядра через combo kernels: вместо запуска десятка крошечных ядер с собственными накладными расходами на каждое, они собираются в комбинированные запуски. Снижен и per-call overhead в Dynamo, той самой постоянной цене, которую каждый вызов скомпилированной функции платит сверх чистого исполнения. На моделях с короткими forward-проходами, где время итерации измеряется миллисекундами, именно эта постоянная составляющая часто съедала весь выигрыш от компиляции.

Eager-режим тоже не забыт: переработаны диспетчеризация и CPU-ядра. На практике это означает, что даже код, который никогда не увидит torch.compile, в 2.14 работает немного быстрее. Суммарно команда описывает эффект как сокращение end-to-end времени исполнения модели без единой правки пользовательского кода, и это редкий случай, когда маркетинговая формулировка подкреплена конкретными механизмами.

Платформы: ROCm 7.14, XPU, Rubin и Python 3.15

Колёсики ROCm 7.14 теперь собираются из TheRock pip SDK, поддержка 7.1 сброшена. Для AMD существеннее другое: модели mixture-of-experts раньше не получали Triton-compiled grouped GEMM в Inductor, потому что он был ограничен NVIDIA SM90+, и ROCm откатывался на медленный цикл по вызовам hipBLASLt. Теперь Triton lowering работает и на ROCm, включая scaled FP8-вариант. А Origami, аналитический селектор tile size от AMD, включён по умолчанию для max-autotune: Inductor выбирает близкие к оптимальным конфигурации GEMM по модели латентности вместо полного прогона автотюнинга.

Intel XPU получил нативный захват графов со сниженными накладными расходами на capture и replay для Arc B-Series и новее, а scaled_mm обзавёлся поддержкой MXFP8 и MXFP4: ранняя программная готовность к следующему поколению Intel GPU. На стороне NVIDIA Inductor уже целится в архитектуру Rubin (sm_107).

Из инфраструктурного: появились бинарные сборки под Python 3.15, включая free-threaded вариант 3.15t без GIL, для всех платформ и бэкендов. Есть нюанс: torch.compile на Python 3.15 пока не поддерживается. И важная новость для мейнтейнеров: torchvision 0.29 теперь ABI-стабилен относительно torch 2.14, то есть останется совместим с 2.15, 2.16 и дальше. Обновлять torchvision при апгрейде torch больше не обязательно, и синхронный релизный цикл двух проектов может закончиться.

Что сломается: депрекации

TorchScript теперь выдаёт видимые предупреждения о депрекации и убран из путей импорта. Профилировщик почистили: устаревшая опция use_cuda удалена, Python function events исключены из key_averages() по умолчанию, что заметно изменит вывод профиля. Из Dynamo вырезан Relay-путь TVM backend, взамен предлагается relax frontend. В distributed _set_pg_timeout уступил место torch.distributed.set_timeout, а реализация control collectives удалена. Разрежённые тензоры при загрузке с weights_only теперь валидируются на консистентность.

Часто задаваемые вопросы

Стоит ли обновляться на PyTorch 2.14?

Да, если вы используете torch.compile с распределённой тренировкой: перекрытие коммуникаций и вычислений теперь включено по умолчанию и даёт прирост утилизации GPU без изменений кода. Обладатели Apple Silicon получат нативную линейную алгебру, а большие кластеры, отказоустойчивые коллективы.

Что такое NVGEMM в PyTorch 2.14?

NVGEMM это backend матричных умножений для TorchInductor на базе NVIDIA CUTLASS DSL. Он генерирует ядра, которые конкурируют с Triton и ATen в autotune, сливает эпилоги (bias, активации, редукции) в одно ядро и поддерживает низкоточный формат NVFP4 на Blackwell.

Как работает отказоустойчивость в nccl2?

Backend и ProcessGroup получили API реконфигурации: при падении ранка группу можно пересобрать на месте, не снося весь process group и не теряя прогретое состояние кластера. Поддерживаются abort-хуки и pre/post-хуки коллективов, Gloo получил такую же поддержку.

Итог

PyTorch 2.14 не про одну громкую фичу, а про доведение до ума всего, что накопилось в серии 2.x. NVGEMM закрывает вопрос зрелых GEMM-ядер в Inductor, nccl2 и fault-tolerant коллективы превращают большие тренировки из хрупкого процесса в управляемый, а @dynamic_spec наконец даёт единый язык для динамических размерностей во всех точках входа. Если вы откладывали переход на torch.compile для распределённых нагрузок, этот релиз лучший повод попробовать: половина улучшений достаётся бесплатно, просто за счёт обновления версии.

← Все записи