PyTorch Distributed и NCCL: как заставить 8 GPU работать как одна

Почему все так боятся torch.distributed?

Видел десятки инженеров, которые берут DataParallel как костыль и молятся, чтобы он работал. Потому что torch.distributed выглядит как чёрная магия: процессы, ранки, группы, блокирующие операции. Но если вы хотите реальной скорости на 4+ GPU, без него никуда.

DataParallel - это как детский велосипед с боковыми колёсами. Работает, но медленно. Он копирует модель на каждый GPU, разбивает батч, собирает градиенты на GPU:0. Всё через один узел. А torch.distributed с NCCL - это гоночный болид, где все GPU общаются напрямую через высокоскоростные линки.

NCCL против RCCL: битва бэкендов

NVIDIA Collective Communications Library (NCCL) - де-факто стандарт для Nvidia GPU. Но с 2024 года у неё появился серьёзный конкурент - RCCL от AMD для их ускорителей. PyTorch поддерживает оба, но есть нюансы.

Вот что многие упускают: PyTorch не всегда выбирает оптимальный бэкенд автоматически. На системе с 8x NVIDIA H100 он может взять GLOO вместо NCCL, если что-то настроено неправильно. Проверяйте всегда:

import torch.distributed as dist
print("Используемый бэкенд:", dist.get_backend())

1Базовая настройка: не повторяйте эти ошибки

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

Правильный паттерн на 2026 год:

import os
import torch
import torch.distributed as dist

def setup(rank, world_size):
    # Важно: одна и та же мастер-адрес для всех процессов
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    
    # Инициализация с таймаутом (новое в PyTorch 2.3)
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        rank=rank,
        world_size=world_size,
        timeout=datetime.timedelta(seconds=30)  # Защита от зависаний
    )
    
    # Привязка процесса к конкретному GPU
    torch.cuda.set_device(rank)


def cleanup():
    dist.destroy_process_group()

Почему 'env://' лучше файлов? Потому что файловая синхронизация ломается на NFS и медленных дисках. А сетевая инициализация через переменные окружения работает стабильнее.

2Коллективные операции: что они реально делают

Всего 6 основных операций, но большинство используют только all_reduce. Ошибка.

• broadcast - отправить тензор с одного процесса всем остальным
• all_reduce - все процессы получают сумму всех тензоров (для градиентов)
• reduce - собрать тензоры на одном процессе
• all_gather - собрать тензоры со всех процессов на всех процессах
• gather - собрать тензоры на одном процессе
• scatter - разослать тензоры со одного процесса всем

Самая важная для обучения - all_reduce. Но вот что мало кто знает: в PyTorch 2.4+ появилась gradient_accumulation_with_all_reduce. Она делает all_reduce во время аккумуляции градиентов, а не после. Экономит 15-20% времени на коммуникации.

# Старый способ (медленный)
loss.backward()
if (i + 1) % accumulation_steps == 0:
    for param in model.parameters():
        dist.all_reduce(param.grad, op=dist.ReduceOp.SUM)
        param.grad /= world_size
    optimizer.step()
    optimizer.zero_grad()

# Новый способ в PyTorch 2.4+
from torch.distributed.algorithms import gradient_accumulation_with_all_reduce

model = gradient_accumulation_with_all_reduce(model, accumulation_steps)
# Теперь all_reduce происходит внутри backward()

Почему ваш NCCL тормозит: 3 скрытые проблемы

1. Неоптимальная топология. NCCL не всегда правильно определяет, какие GPU соединены NVLink. На системах с 8 GPU часто бывает, что GPU 0-3 связаны NVLink, 4-7 тоже, но между этими группами только PCIe. Если разместить процессы без учёта этого - получите узкое горло.

Решение:

# Перед запуском проверьте топологию
nvidia-smi topo -m

# Или из Python
import pynvml
pynvml.nvmlInit()
for i in range(torch.cuda.device_count()):
    handle = pynvml.nvmlDeviceGetHandleByIndex(i)
    print(f"GPU {i}: {pynvml.nvmlDeviceGetName(handle)}")

2. Буферы неправильного размера. NCCL аллоцирует буферы при инициализации. Если вы потом пытаетесь передать тензор большего размера - будет переаллокация. Это дорого.

3. Конфликт с CUDA graphs. С 2025 года многие используют CUDA graphs для ускорения. Но NCCL операции внутри графа ведут себя иначе. Если график содержит коллективную операцию, она должна быть одинаковой во всех процессах (одинаковые размеры тензоров, одинаковые типы).

Группы процессов: зачем они нужны

По умолчанию все процессы в одной группе. Но что если вам нужно:

• Сделать all_reduce только на половине GPU (например, для pipeline параллелизма)
• Разные операции на разных подмножествах
• Иерархическая коммуникация (сначала внутри узла, потом между узлами)

Создание группы:

# Создаём группу из первых 4 GPU
ranks = [0, 1, 2, 3]
group = dist.new_group(ranks)

# Теперь операции только внутри группы
if dist.get_rank() in ranks:
    dist.all_reduce(tensor, group=group)

# Иерархическая коммуникация (для мульти-нод)
intra_node_group = dist.new_group([0, 1, 2, 3])  # GPU на одной ноде
inter_node_group = dist.new_group([0, 4])        # GPU0 на разных нодах

Иерархическая коммуникация может ускорить обмен градиентами в 2-3 раза на кластерах. Сначала собираем градиенты внутри ноды, потом обмениваемся между нодами.

Асинхронные операции: рисковано, но быстро

Все коллективные операции по умолчанию блокирующие. Процесс ждёт завершения. Но можно сделать асинхронно:

handle = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True)
# Делаем что-то ещё...
handle.wait()  # Ждём завершения

Проблема в том, что если начать следующую операцию до wait(), NCCL может упасть. Или хуже - повиснуть. На практике асинхронные операции стоит использовать только когда точно знаете, что делаете.

Отладка: когда всё падает в 3 часа ночи

Включите логирование NCCL:

export NCCL_DEBUG=INFO
export NCCL_DEBUG_SUBSYS=INIT,COLL
export NCCL_DEBUG_FILE=/tmp/nccl_debug_%h_%p.log

Самые частые ошибки:

Если NCCL падает с загадочной ошибкой, попробуйте уменьшить количество потоков:

export NCCL_MAX_NCHANNELS=1
export NCCL_BUFFSIZE=16777216

Интеграция с DDP: что происходит под капотом

DistributedDataParallel (DDP) использует коллективные операции, но скрывает сложность. На 2026 год в PyTorch 2.4+ DDP стал умнее:

• Автоматически определяет, нужно ли использовать bucketizing (группировка градиентов)
• Поддерживает gradient accumulation внутри
• Может работать со статическими графами (torch.compile)

Но DDP не панацея. Если у вас custom backward или сложная модель, возможно, придётся писать свой all_reduce. Например, для тензорного параллелизма как в Llama.cpp.

Гетерогенные системы: когда есть и NVIDIA, и AMD

С 2025 года это стало реальностью благодаря совместной работе разных архитектур. Но NCCL и RCCL между собой не общаются. Решение:

# Разные группы для разных архитектур
nvidia_ranks = [0, 1]  # NVIDIA GPU
amd_ranks = [2, 3]     # AMD GPU

nvidia_group = dist.new_group(nvidia_ranks, backend='nccl')
amd_group = dist.new_group(amd_ranks, backend='rccl')

# Обмен через CPU (медленно, но работает)
cpu_tensor = tensor.cpu()
dist.all_reduce(cpu_tensor, group=...)

Лучше так не делать. Если нужно гетерогенное окружение, посмотрите на оптимизацию AI-станций.

Производительность: цифры, которые имеют значение

На 8x NVIDIA H100 с NVLink 4.0:

• all_reduce для 1GB тензора: ~15 мс
• Пропускная способность: до 600 GB/s
• Задержка: 2-5 мс в зависимости от топологии

Но это в идеальных условиях. На практике из-за неправильной настройки часто получается 100 GB/s и 50 мс. Проверяйте с torch.benchmark:

import torch.distributed as dist
import time

def benchmark_all_reduce(size_mb=100):
    tensor = torch.randn(size_mb * 1024 * 1024 // 4, device='cuda')  # float32
    
    # Синхронизируем все процессы
    dist.barrier()
    
    start = time.time()
    dist.all_reduce(tensor)
    torch.cuda.synchronize()
    end = time.time()
    
    # Только rank 0 выводит результат
    if dist.get_rank() == 0:
        bandwidth = (size_mb * 2) / (end - start)  * 2  * 2
        print(f"Bandwidth: {bandwidth:.2f} MB/s")
        print(f"Time: {(end-start)*1000:.2f} ms")

Куда двигаться дальше

Когда освоите базовые коллективные операции, смотрите в сторону:

1. Pipeline параллелизм - для моделей, которые не влезают в память одной GPU
2. Tensor параллелизм - разделение матриц умножения между GPU
3. Zero Redundancy Optimizer (ZeRO) - экономия памяти за счёт разделения состояний оптимизатора
4. Fully Sharded Data Parallel (FSDP) - следующее поколение после DDP

Но помните: каждая абстракция добавляет накладные расходы. Иногда простой DDP с правильно настроенным NCCL даст больше, чем сложная система с FSDP.

Если столкнётесь с проблемами масштабирования на кластерах, посмотрите как исправить проблемы DGX Spark. А для продакшн-развёртывания пригодится опыт с Kubernetes и KServe.

Коллективные операции - это фундамент распределённого обучения. Потратьте день на их изучение, сэкономите недели на отладке. И не бойтесь заглядывать в исходники PyTorch. Там нет магии, только код. Иногда слишком много кода.