PyTorch 2.15: NCCL по умолчанию и FP8 на Apple Silicon

PyTorch 2.15: NCCL по умолчанию и FP8 на Apple Silicon

Самое неприятное в обучении на сотнях GPU не падение процесса, а тихое зависание. Коллектив all-reduce не возвращает ошибку, лосс стоит на месте, карты заняты, и остаётся гадать, какой из рангов перестал отвечать. Раньше это лечилось воспроизведением сбоя руками, теперь в PyTorch 2.15 есть flight recorder: как только один ранг ловит таймаут, все остальные сбрасывают на диск отчёт о том, что делали в момент остановки.

PyTorch 2.15 собрали из 3117 коммитов от 583 контрибьюторов. Разница с 2.14 по объёму примерно такая же, но суть изменений другая. Предыдущие версии добавляли возможности, эта доводит до состояния по умолчанию то, что раньше было опцией: переписанный бэкенд NCCL теперь обслуживает каждый процесс-группу, а не только тех, кто явно попросил.

Что такое PyTorch 2.15

Это очередной релиз открытого фреймворка машинного обучения, вышедший через два месяца после 2.14. Внутри: новый NCCL-бэкенд по умолчанию, one-sided транспорты с бэкендом NIXL, точная настройка памяти в FSDP2, FP8 и быстрый decode на Apple Silicon, поддержка NVFP4 для инференса на Blackwell, ROCm 10.0 и упрощённое написание SYCL-ядер для видеокарт Intel.

Общая линия серии 2.x видна без микроскопа. Фреймворк перестал быть инструментом исследователя и стал платформой, на которой одна и та же команда запускает обучение на NVIDIA, AMD, Apple и Intel. 2.13 принесла FlexAttention на Apple Silicon и CuTeDSL в Inductor, 2.14 превратила CuTeDSL в полноценный GEMM-бэкенд NVGEMM. В 2.15 эти ветки доходят до продакшн-состояния, а самая заметная работа ушла в распределённую коммуникацию.

Переписанный NCCL стал дефолтом

В 2.14 новый бэкенд, портированный из torchcomms, существовал как опция под именем nccl2. Теперь он используется по умолчанию для любой группы процессов с бэкендом nccl. Откат остался: переменная окружения TORCH_DIST_USE_NCCL2=0 или имя бэкенда nccl-legacy возвращают старый ProcessGroupNCCL.

Что поменялось по существу. Коммуникаторы создаются сразу вместе с группой процессов, а не лениво при первом коллективе. Реконфигурация получила быстрый путь shrink/grow на NCCL 2.29 и новее: выжившие ранги пропускают полный bootstrap и не тратят на него минуты после падения одной ноды. На ROCm то же самое работает с RCCL 2.30.7.

Flight recorder оказался важнее, чем звучит. В обычном режиме отчёт о зависании пишется на диск, а состояние watchdog уходит на debug-сервер. Таймаут на любом ранге запускает дампы со всех рангов одновременно, то есть в отчёте оказывается не только жертва, но и соседи, которые её ждали. На кластере из тысячи GPU один застрявший узел останавливает весь прогон, и раньше поиск причины занимал часы ручного перебора.

One-sided транспорты и NIXL

В torch.distributed появился Python API для односторонней передачи данных, и первым бэкендом стал NIXL. Транспорт подключается к пиру через store, без группы процессов и без коллективов. Методы Transport.read_stream и Transport.write_stream ставят передачу в очередь на CUDA stream, поэтому её можно захватить в torch.cuda.graph.

Смысл в том, чтобы развязать передачу данных и синхронизацию. Коллектив по определению заставляет всех участников ждать самого медленного, а one-sided операция позволяет одному рангу забрать нужный кусок памяти у другого, не собирая барьер. Для инференса с разделением модели по узлам, для стриминга весов и для асинхронных чекпоинтов это меняет схему обмена. API пока помечен как нестабильный, что честно: интерфейс будет меняться.

FSDP2 и пайплайны: контроль точности и памяти

Главное изменение для тех, кто обучает MoE. Чувствительные к точности параметры, такие как роутерные гейты и веса нормализации, можно держать в fp32, пока остальная модель учится в bf16. Делается через MixedPrecisionPolicy(param_dtype_override_fn=...), и all-gather на каждый fully_shard остаётся один. Раньше выбор был бинарным: либо весь шард в bf16, либо потеря скорости.

Причина, по которой это важно, в арифметике роутера. Гейт выбирает эксперта по логитам, которые отличаются в третьем знаке, и bf16 с его восемью битами мантиссы легко переворачивает выбор. Один неверный эксперт на токен, и обучение уезжает. Точечный fp32 для таких тензоров дешевле, чем полный fp32 для всего блока.

Полный шаг накопления градиента теперь захватывается в один CUDA graph, а точность градиентов задаётся отдельно для каждого параметра: они аккумулируются прямо в .grad и следуют за grad_dtype. Пайплайн-параллелизм перестал выделять буферы приёма на каждый микробатч заранее, вместо этого буфер берётся прямо перед приёмом, а выходы forward освобождаются после завершения отправки. Расписания принимают бюджет max_outstanding_sends, расшардирование FSDP размазывается по стартовому пузырю ранга, и каждая уникальная пара рангов в пайплайне получает собственный P2P-коммуникатор. Reduce градиентов FSDP теперь перекрывается с последующими вычислениями.

В сумме пайплайнинг стал пригодным для захвата в CUDA graph и детерминированным, а не просто быстрым. Для команд, которые гоняют недельные прогоны, детерминизм важнее процента скорости: без него невозможно отладить ни один сбой.

NVFP4 decode на Blackwell

NVGEMM, появившийся в 2.14, научился втягивать скалярный коэффициент деквантования слоёв ModelOpt NVFP4 прямо в GEMM. Раньше на каждый такой слой приходился отдельный pointwise-запуск, который читал результат из памяти и умножал его на масштаб. Во время decode, где на один слой уходят миллисекунды, лишний запуск ядра это заметная доля времени.

Кроме этого добавилась политика выбора кандидатов, учитывающая форму тензоров, для случаев NVFP4-decode, которые пропускает эвристика nvMatmulHeuristics. Ядра NVGEMM умеют использовать programmatic dependent launch, поэтому Triton-потребитель стартует, пока хвост производящего GEMM ещё считается. Бэкенд переехал на CUTLASS Operators.

Apple Silicon: FP8, потоки и QR за 1,5 миллисекунды

Самая интересная часть релиза для тех, у кого нет кластера. На MPS появилась поддержка torch.float8_e4m3fn вместе с _scaled_mm в двух режимах масштабирования, tensorwise и rowwise, с быстрым путём gemv для декодерных форм. Плюс публичный API потоков, смоделированный по образцу CUDA: torch.mps.Stream, current_stream, set_stream, контекстный менеджер torch.mps.stream и Tensor.record_stream.

Цифры по ядрам выглядят несопоставимо с прошлыми релизами. Одиночный softmax по словарю в 128 256 токенов ускорился в 8 раз. Обратный проход upsample стал быстрее до 21,7 раза. Редукции перевели на общий набор ядер, и полные редукции для норм ускорились более чем в 100 раз.

Но главная цифра другая. QR-разложение матрицы 512 на 512 упало примерно с 7 секунд до 1,5 миллисекунды. Это не опечатка и не результат агрессивного округления: старый путь шёл через MPSGraph с постоянными переключениями и копированиями, новый написан как нативное Metal-ядро. Треугольные решения тоже переписали на блочный алгоритм. Для тех, кто использует Mac как рабочую станцию для прототипов, разница между семью секундами и полутора миллисекундами означает, что численные эксперименты перестают быть поводом идти обедать.

Для декодера F.linear перевели на существующие gemv-ядра: ускорение 1,23 раза по геометрическому среднему и до 5 раз на маленьких проекциях. Новый grouped_mm получил gemv-путь и даёт прирост пропускной способности MoE-decode от 1,32 до 1,34 раза на M3 Ultra.

AMD: FlyDSL и MI450

FlyDSL встал в один ряд с CuTeDSL, Triton и Helion как DSL, который PyTorch умеет использовать. На gfx950 Inductor автотюнит шаблоны FlyDSL для GEMM, ragged F.grouped_mm и MXFP4/MXFP8 scaled_mm наравне с Triton и ATen, а в eager-режиме FlyDSL может подменять RMSNorm и topk. Опция необязательная и требует FlyDSL 0.3.x.

На gfx1250, это MI450, Inductor задействует Tensor Data Mover для плотных mm и addmm, pointwise и reduction-кодогенерации, а также для forward и decode в FlexAttention. Архитектуру добавили в список ночных сборок, то есть поддержка появилась до того, как карты доехали до покупателей.

Отдельно про колёса. Появились сборки под ROCm 10.0 рядом с ROCm 7.14 из 2.14, и все они теперь собираются на пакетах TheRock. ROCm ставится как pip-зависимость, а не вшивается в wheel. Для тех, кто держит несколько окружений на одной машине, это снимает главную боль прошлых версий, когда обновление PyTorch тянуло за собой замену системного ROCm.

Intel XPU: SYCL без шаблонных функторов

Для XPU-бэкенда добавили поддержку свободных функций в SYCL, и большинство существующих ядер уже перевели на новую модель. Вместо многословных функторов теперь пишутся обычные функции, что снижает порог входа для тех, кто портирует CUDA-ядра на Intel.

Рантайм получил IPC для torch.xpu.Event, так что события можно передавать между процессами и использовать с torch.multiprocessing. torch.xpu.MemPool научился политикам use_on_oom и non_split, что полезно при работе с распределённой симметричной памятью и при обработке нехватки памяти.

Профилирование стало глубже. Появилась видимость ядер при replay XPU Graph, кросс-ранговая корреляция коллективов XCCL и сбор аппаратных счётчиков по каждому ядру сразу на всех видеокартах в многогпу-задаче.

Платформа: CUDA 13.2, Python 3.15 и Cuspy

Дефолтное колесо переехало на CUDA 13.2, рядом публикуются сборки под CUDA 13.4 GA для Linux x86_64, Linux aarch64 и Windows. В 13.4 добавлена экспериментальная поддержка GPU NVIDIA Rubin.

Python поддерживается от 3.11 до 3.15, включая сборки со свободными потоками 3.14t и 3.15t. Dynamo научился работать на 3.15, а load_inline теперь умеет генерировать pybind11-модули, которые безопасно импортируются в free-threaded Python.

Стабильный ABI расширили функцией torch::stable::sort, а в Python-обвязке появились конверсии dtype и device плюс проверка is_tensor. Часть c10/util уехала в torch/headeronly, старые пути остались как перенаправляющие заголовки. Расширения теперь получают ошибку компиляции, если TORCH_TARGET_VERSION новее их заголовков libtorch: защита от ситуации, когда плагин собран против другой версии и падает в рантайме.

Профилировщик CUPTI переименовали в Cuspy, и он научился главному для тех, кто ловит просадки в графах: ядра, воспроизведённые из CUDA graph, теперь привязываются к коду, который их захватил. Аннотации пишутся в export_chrome_trace инлайном, а с record_py_stacks граф сохраняет и Python-стек запуска для каждого узла.

Что сломается при обновлении

Python 3.10 больше не поддерживается: колёс под него нет, в CI он не тестируется, сборка из исходников не гарантируется. macOS 14 тоже выпал, для MPS нужен macOS 15 или новее, обходные пути под старую систему удалили. Если у вас CI на macOS 14, обновление фреймворка потянет за собой обновление раннера.

Два изменения поведут себя тихо. Первое: бэкенд nccl теперь выбирает переписанную реализацию, и код, который лез во внутренности ProcessGroupNCCL, может сломаться без внятной ошибки. Второе: FSDP2 разделил dtype шардированных и нешардированных градиентов, и при reduce_dtype=None поведение по умолчанию изменилось. Если в проекте есть свои хуки на градиенты, их стоит перечитать. Предупреждения C++ из torch.ops и кастомных операторов теперь доходят до Python, так что шумные расширения начнут сыпать предупреждениями в логах.

FAQ

PyTorch 2.15 стоит ставить сразу? Если вы обучаете на нескольких GPU через nccl, стоит: flight recorder и быстрый shrink/grow окупаются на первом же зависании. Если у вас inference на Apple Silicon, обновление даёт FP8 и заметно более быстрый decode. Для проектов на Python 3.10 обновление заблокировано до переезда на 3.11.

Как откатиться на старый NCCL, если что-то пошло не так? Поставить TORCH_DIST_USE_NCCL2=0 или указать бэкенд nccl-legacy при создании группы процессов. Старая реализация остаётся в дереве, так что откат не требует пересборки.

Зачем FP8 на ноутбуке, если там нет тензорных ядер Blackwell? FP8 на MPS нужен не для обучения больших моделей, а для инференса квантизованных весов и для экспериментов с масштабированием. Половина памяти под веса и в 1,3 раза быстрее decode MoE на M3 Ultra это ощутимая разница, когда модель не влезает в unified memory целиком.

Почему QR ускорился в тысячи раз, а матричное умножение только на проценты? Потому что до этого QR шёл через MPSGraph, где каждая операция сопровождалась переключением контекста и копированием, а GEMM уже был оптимизирован. Ускорение в 4600 раз это не магия нового алгоритма, а устранение накладных расходов, которые раньше были в сотни раз больше полезной работы.

Итог

PyTorch 2.15 это релиз про надёжность распределённых прогонов и про то, чтобы железо кроме NVIDIA перестало быть вторым сортом. Переписанный NCCL с flight recorder уменьшает время от зависания до диагноза, one-sided транспорты открывают схему обмена без барьеров, FSDP2 даёт точечный fp32 там, где bf16 ломает роутеры MoE, а Apple Silicon получает FP8, публичные потоки и QR за полторы миллисекунды вместо семи секунд.

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

← Все записи