Что такое NCCL и почему он критичен для multi-GPU инференса?

Краткий тезис

NCCL (NVIDIA Collective Communications Library) — это высокопроизводительная библиотека для организации обмена данными между GPU. Для multi-GPU инференса, особенно при использовании tensor parallelism, каждый forward pass требует синхронизации результатов между GPU через операцию AllReduce. NCCL реализует такие коллективные операции с максимальной эффективностью, используя аппаратные возможности NVLink и PCIe. Без NCCL latency инференса может вырасти в 5–10 раз из-за неоптимальных программных решений.

1. Термин: NCCL (NVIDIA Collective Communications Library)

NCCL — это проприетарная библиотека от NVIDIA, предназначенная для выполнения коллективных коммуникаций между GPU в одном узле или между узлами через сеть (InfiniBand, RoCE). Она предоставляет оптимизированные реализации операций AllReduce, Broadcast, AllGather, ReduceScatter и других.

Исторически NCCL появилась для ускорения распределённого обучения (training|distributed training), но сегодня она критична и для инференса больших моделей, когда одна модель не помещается в память одного GPU и требуется параллелизм.

Ключевые особенности:

• Использует аппаратные пути NVLink (меж-GPU) и InfiniBand (меж-узлы) для минимизации задержек.
• Автоматически выбирает оптимальный алгоритм (ring, tree, direct) в зависимости от топологии.
• Работает на уровне CUDA, что позволяет избегать лишних копирований через CPU.

2. Multi-GPU инференс: зачем и какие виды параллелизма

При инференсе больших языковых моделей (LLM) или моделей компьютерного зрения часто не хватает памяти одного GPU или требуется снизить latency за счёт распараллеливания. Основные стратегии:

Для инференса parallelism|tensor parallelism даёт наименьшую задержку, но требует самой интенсивной синхронизации — именно здесь NCCL незаменим.

3. Коллективные операции NCCL

Библиотека реализует несколько базовых операций. Рассмотрим ключевые для multi-GPU инференса:

• AllReduce — каждый GPU имеет свой тензор; после операции все GPU получают одинаковый тензор, равный сумме (или другой редукции) исходных. Используется в parallelism|tensor parallelism для суммирования частичных результатов.
• Broadcast — один GPU рассылает свои данные всем остальным. Полезен при загрузке весов модели на все GPU.
• AllGather — каждый GPU собирает данные со всех GPU в один большой тензор. Применяется, например, для сбора скрытых состояний.
• ReduceScatter — комбинация Reduce и Scatter: сначала редукция, затем результат распределяется по GPU. Используется в некоторых реализациях tensor parallelism для экономии памяти.

Пример кода на Python с использованием PyTorch (который внутри использует NCCL):

import torch
import torch.distributed as dist

# Инициализация группы процессов (NCCL backend)
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)

# Каждый GPU имеет свой тензор
tensor = torch.tensor([rank], device=f'cuda:{rank}')

# AllReduce: суммируем все тензоры
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
print(f'Rank {rank}: after all_reduce = {tensor.item()}')
# Результат: на всех GPU будет сумма рангов (0+1+2+...)

4. Почему NCCL критичен для multi-GPU инференса

Основная причина — производительность. Рассмотрим альтернативы:

Для tensor parallelism типичный размер сообщения — несколько мегабайт (скрытые состояния). NCCL использует ring-алгоритм для AllReduce, который разбивает данные на чанки и передаёт их по кольцу, минимизируя пиковую нагрузку на каналы. Без NCCL пришлось бы использовать CPU как промежуточное звено или писать неэффективные копирования, что резко увеличивает latency.

Кроме того, NCCL умеет перекрывать коммуникации с вычислениями (overlap), используя CUDA streams. Это позволяет GPU одновременно считать и передавать данные, что критично для инференса в реальном времени.

5. Tensor Parallelism и AllReduce

Рассмотрим, как работает tensor parallelism на примере линейного слоя y = xW. Если веса W разделены между двумя GPU (по столбцам), то:

1. Каждый GPU получает полный вход x.
2. Каждый GPU вычисляет свою часть y_i = x @ W_i.
3. Чтобы получить полный y, нужно выполнить AllReduce по y_i (суммирование).

В PyTorch это выглядит так:

# Псевдокод для tensor parallelism с NCCL
class LinearTP(torch.nn.Module):
    def __init__(self, in_features, out_features, rank, world_size):
        super().__init__()
        self.rank = rank
        self.world_size = world_size
        # Каждый GPU хранит только часть весов
        self.weight = torch.nn.Parameter(torch.randn(in_features, out_features // world_size))
    
    def forward(self, x):
        # Локальное вычисление
        local_out = x @ self.weight
        # AllReduce для суммирования результатов со всех GPU
        dist.all_reduce(local_out, op=dist.ReduceOp.SUM)
        return local_out

Без NCCL пришлось бы вручную копировать local_out на CPU, суммировать и рассылать обратно — это добавило бы десятки миллисекунд к каждому слою.

6. Pipeline Parallelism и P2P

В pipeline parallelism данные передаются от GPU к GPU последовательно. Для этого используются операции send/recv (point-to-point). NCCL поддерживает их через dist.send() и dist.recv(), но часто для pipeline применяют более лёгкие механизмы (например, P2P через NVLink напрямую). Однако NCCL всё равно может быть задействован для синхронизации границ микробатчей.

7. Data Parallelism (для справки)

Хотя в чистом инференсе data parallelism не требует коммуникаций, в распределённом обучении (training) AllReduce градиентов — основная нагрузка. NCCL здесь также критичен: без него обучение на многих GPU было бы невозможно из-за узкого места на CPU.

8. Что происходит без NCCL: latency 5–10x

Если попытаться реализовать AllReduce вручную через cudaMemcpy на CPU, каждый GPU должен:

• Скопировать свой тензор в память CPU (через PCIe, latency ~10 мкс).
• CPU выполняет суммирование (медленно).
• Скопировать результат обратно на каждый GPU.

Для 8 GPU это 16 копирований через PCIe (8 туда, 8 обратно) плюс вычисления на CPU. NVLink обеспечивает прямые соединения GPU-GPU с latency ~1 мкс и пропускной способностью до 600 ГБ/с. NCCL использует эти соединения, минуя CPU. В результате latency падает в 5–10 раз для типичных размеров сообщений.

9. Практические аспекты настройки NCCL

При развёртывании multi-GPU инференса полезно знать переменные окружения NCCL:

• NCCL_DEBUG=INFO — выводит детальную информацию о выбранных алгоритмах и топологии.
• NCCL_IB_DISABLE=1 — отключить InfiniBand (если используется только NVLink).
• NCCL_NET_GDR_LEVEL — управление GPU Direct RDMA.
• NCCL_SOCKET_IFNAME — указать сетевой интерфейс для меж-узловых коммуникаций.

Также можно профилировать коммуникации с помощью nsys (NVIDIA Nsight Systems) и видеть, сколько времени занимают коллективные операции.

Пет-проект для закрепления

Задача Написать скрипт, который запускает инференс маленькой модели (например, два линейных слоя) на 2 GPU с tensor parallelism, используя NCCL для AllReduce. Измерить время одного forward pass с NCCL и без него (имитация через CPU).

Инструменты Python, PyTorch, две GPU (можно в облаке или локально), torch.distributed.

Шаги:

1. Установить PyTorch с поддержкой CUDA и NCCL.
2. Написать скрипт tp_inference.py, который инициализирует процессную группу с backend='nccl'.
3. Реализовать класс LinearTP как в разделе 5.
4. Создать модель из двух таких слоёв.
5. Запустить инференс на случайном батче (batch_size=1, seq_len=128, hidden=1024).
6. Замерить время с помощью torch.cuda.synchronize() и time.time().
7. Для сравнения написать версию без NCCL: каждый GPU копирует свой результат на CPU, суммирует вручную, копирует обратно.
8. Сравнить latency.

Ожидаемый результат Скрипт покажет, что NCCL в 5–10 раз быстрее. Вы также увидите в логах NCCL выбранный алгоритм (например, ring).

Связь с другими вопросами

Навигация

• Предыдущий: 401
• Следующий: 403
• Индекс: 00. Индекс разборов