Оглавление

Команда PyTorch представила torchcomms — экспериментальный легковесный API для распределенных вычислений, предназначенный для работы с PyTorch Distributed. Вместе с базовым API открывается исходный код NCCLX — нового бэкенда, способного масштабироваться на более чем 100 000 графических процессоров.

Новая эра распределенных вычислений

Первая версия torchcomms предлагает фундаментальные API и бэкенды, необходимые для обучения крупномасштабных моделей в PyTorch. Основное внимание уделено базовым коммуникационным примитивам, обеспечивающим надежное и производительное распределенное обучение в промышленных масштабах.

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

Ключевые цели проекта

Разработчики ставят перед собой амбициозные задачи:

  • Быстрое прототипирование — исследователи смогут экспериментировать с новыми коммуникационными парадигмами независимо от основных численных примитивов PyTorch
  • Масштабирование до 100K+ GPU — введение агрессивной инициализации и оптимизация распределения ресурсов для сверхмасштабных задач
  • Поддержка гетерогенного железа — работа в смешанных средах с разными поколениями и производителями оборудования
  • Отказоустойчивость — создание надежного бэкенда для алгоритмов вроде fault-tolerant HSDP и Streaming DiLoCo
  • Односторонняя коммуникация — поддержка RDMA-подобных семантик для асинхронных workflow
  • Устройственно-ориентированные коллективные операции — тесная интеграция вычислений и коммуникаций для сверхнизких задержек

Почему новый API?

Существующие API c10d в PyTorch Distributed несут значительное техническое бремя, что затрудняет их расширение и модернизацию. Torchcomms предоставляет свободу для экспериментов и эволюции дизайна без ограничений устаревших интерфейсов.

Это тот случай, когда технический долг становится настолько велик, что проще построить новый дом, чем ремонтировать старый. Существующие ограничения NCCL и c10d давно тормозили развитие распределенных вычислений, особенно в эпоху моделей с триллионами параметров. Meta* фактически признала, что текущая архитектура не справляется с масштабами Llama 3 и 4, и создала инструмент, который позволит не просто догнать, а перегнать текущие потребности.

Быстрый старт

Базовое использование torchcomms выглядит интуитивно понятно:

import torchcomms

# Инициализация коммуникатора
comm = torchcomms.new_comm("ncclx", torch.device("cuda"), name="my_comm")
print(f"I am rank {comm.get_rank()} of {comm.get_size()}!")

t = torch.full((10, 20), value=comm.rank, dtype=torch.float)

# Синхронный all_reduce
comm.allreduce(t, torchcomms.ReduceOp.SUM, async_op=False)

# Асинхронный all_reduce
work = comm.allreduce(t, torchcomms.ReduceOp.SUM, async_op=True)
work.wait()

# Разделение коммуникатора на группы
split_groups = torch.arange(comm.get_size()).view(-1, 8).tolist()
tp_comm = comm.split(split_groups)

Совместимость с DeviceMesh

Torchcomms поддерживает совместимость с DeviceMesh для работы с библиотеками параллелизма PyTorch, такими как FSDP2:

import torchcomms
from torchcomms.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard

comm = torchcomms.new_comm("ncclx", torch.device("cuda:0"), name="global")

mesh = init_device_mesh(
 mesh_dim_comms=(comm,),
 mesh_dim_names=("global",),
)
fully_shard(model, device_mesh=mesh)

Вместе с новыми API выпущено несколько бэкендов для различных аппаратных платформ. NCCLX содержит расширения Meta* для популярной библиотеки NCCL. Этот бэкенд уже прошел проверку в продакшене — он используется для крупномасштабного обучения и инференса больших языковых моделей, таких как Llama3 и Llama4.

Сегодня все генеративные модели Meta* работают на этой технологии, что подтверждает ее готовность к промышленному использованию.

По материалам PyTorch.

* Meta (признана экстремистской и запрещена в РФ)