Команда 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 (признана экстремистской и запрещена в РФ)