AI в мульти‑GPU: точечные и коллективные операции в PyTorch

Почему распределённые вычисления становятся обязательными

Современные модели искусственного интеллекта требуют гигабайтов памяти и терафлопов вычислительной мощности. Одна видеокарта уже не может обеспечить необходимый объём ресурсов, поэтому разработчики переходят к распределённым системам, где несколько GPU работают совместно. Основным инструментом для организации такой параллельной работы в экосистеме PyTorch является пакет torch.distributed. Он позволяет реализовать два фундаментальных типа коммуникаций: точечные (point‑to‑point) и коллективные (collective) операции.

Точечные операции: send / recv и их применение

Точечные операции — это прямой обмен данными между двумя процессами (или двумя GPU). В PyTorch они реализованы через функции torch.distributed.send и torch.distributed.recv.

# процесс 0 отправляет тензор t
torch.distributed.send(tensor=t, dst=1)

# процесс 1 получает тензор в буфер recv_t
torch.distributed.recv(tensor=recv_t, src=0)

Когда использовать

• Алгоритмы с асимметричной нагрузкой – когда один узел отвечает за подготовку данных, а остальные только их потребляют.
• Пайплайнинг – последовательный проход модели, где каждый слой размещён на отдельном GPU, и результаты передаются дальше по цепочке.
• Контрольные сообщения – небольшие сигналы синхронизации, которые не требуют глобального барьера.

Точечные операции дают гибкость, но требуют от разработчика явного управления порядком отправки и приёма, иначе возможны взаимные блокировки (deadlock). Для избежания такой ситуации часто используют неблокирующие версии isend/irecv, а затем вызывают torch.distributed.wait.

Коллективные операции: синхронизация и агрегация

Коллективные операции позволяют одновременно взаимодействовать со всеми процессами в группе. Ключевые функции:

Применение в обучении моделей

• Синхронный SGD – каждый GPU вычисляет градиенты локально, а all_reduce суммирует их, после чего каждый процесс обновляет свои параметры. Это обеспечивает идентичные весовые коэффициенты на всех устройствах.
• Data Parallelism – broadcast используется для распространения начального состояния модели от главного процессора к остальным.
• Model Parallelism – scatter и gather позволяют разбивать большие весовые матрицы между GPU и собирать их для последующего вычисления.

Коллективные операции оптимизированы под конкретные бекенды и часто реализуются с помощью высокопроизводительных библиотек, таких как NCCL (для NVIDIA GPU) и Gloo (универсальный, поддерживает CPU и GPU).

Выбор бекенда и инициализация процесса

Для корректной работы torch.distributed необходимо задать бекенд и инициализировать процессную группу. Наиболее распространённые варианты:

import torch.distributed as dist

dist.init_process_group(
    backend='nccl',          # 'nccl' для NVIDIA, 'gloo' для CPU/AMD
    init_method='env://',   # читаем переменные окружения (MASTER_ADDR, MASTER_PORT)
    world_size=4,           # количество процессов (GPU)
    rank=local_rank         # уникальный идентификатор текущего процесса
)

• NCCL обеспечивает наилучшую пропускную способность и низкую латентность при работе с несколькими NVIDIA GPU.
• Gloo более гибок, поддерживает как CPU, так и GPU, но обычно медленнее NCCL в чисто GPU‑сценариях.

Важно, чтобы каждый процесс запускался с правильным local_rank, часто передаваемым через переменную окружения CUDA_VISIBLE_DEVICES. В сценариях с torchrun (ранее torch.distributed.launch) это делается автоматически.

Практические рекомендации по производительности

1. Минимизировать синхронные барьеры – каждый barrier останавливает весь кластер, поэтому используйте их только в случае реальной необходимости.
2. Пакетировать мелкие сообщения – вместо частых небольших send/recv лучше собрать данные в один буфер и отправить за один раз.
3. Привязывать тензоры к нужным устройствам – перед вызовом коллективных функций убедитесь, что все тензоры находятся на тех же GPU, иначе будет лишняя копия через PCIe.
4. Профилировать с torch.profiler – он позволяет увидеть, какие операции занимают большую часть времени и где происходит узкое место коммуникации.
5. Учесть топологию узла – в многосокетных серверах связь между GPU, находящимися в разных сокетах, может иметь большую латентность; распределяйте процессные группы так, чтобы часто общающиеся GPU находились в одном сокете.

Пример полного цикла обучения с DistributedDataParallel

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def train(rank, world_size):
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

    model = nn.Linear(1024, 10).cuda()
    ddp_model = DDP(model, device_ids=[rank])

    optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01)
    loss_fn = nn.CrossEntropyLoss()

    dataset = torch.utils.data.TensorDataset(
        torch.randn(10000, 1024), torch.randint(0, 10, (10000,))
    )
    sampler = torch.utils.data.distributed.DistributedSampler(
        dataset, num_replicas=world_size, rank=rank
    )
    loader = torch.utils.data.DataLoader(dataset, batch_size=32, sampler=sampler)

    for epoch in range(5):
        sampler.set_epoch(epoch)  # обеспечивает перемешивание в каждом эпо
        for data, target in loader:
            data, target = data.cuda(rank), target.cuda(rank)
            optimizer.zero_grad()
            output = ddp_model(data)
            loss = loss_fn(output, target)
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()

if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size)

В этом примере DistributedDataParallel автоматически использует all_reduce для синхронизации градиентов, а DistributedSampler гарантирует, что каждый процесс получает уникальную часть датасета. Такой шаблон покрывает большинство сценариев обучения на нескольких GPU и служит надёжной отправной точкой для дальнейшего экспериментирования с точечными и коллективными операциями.