Что такое NCCL и зачем он для tensor parallelism?

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

NCCL (NVIDIA Collective Communications Library) — это библиотека от NVIDIA для эффективных коммуникаций между GPU (и между узлами). Она предоставляет оптимизированные коллективные операции (AllReduce, Broadcast, AllGather и др.), критически важные для tensor parallelism — стратегии распределённого обучения/инференса, при которой один тензор (например, вес слоя) разрезается на части и размещается на разных GPU. Без NCCL tensor parallelism был бы неэффективен из-за высоких накладных расходов на передачу данных.

1. Что такое NCCL

NCCL (NVIDIA Collective Communications Library) — низкоуровневая библиотека, реализующая коллективные операции для GPU. Она использует аппаратные возможности NVIDIA: NVLink (высокоскоростное соединение GPU-GPU внутри одного узла) и InfiniBand (межузловое соединение). NCCL автоматически выбирает оптимальный путь передачи данных (Ring, Tree, и т.д.) в зависимости от топологии.

Основные операции NCCL

• AllReduce — суммирование/усреднение тензоров со всех GPU и рассылка результата каждому.
• Broadcast — отправка тензора с одного GPU на все остальные.
• AllGather — сборка фрагментов тензора со всех GPU в полный тензор на каждом.
• ReduceScatter — суммирование фрагментов и распределение результата по GPU (обратная AllGather).

Пример инициализации NCCL в PyTorch

import torch
import torch.distributed as dist

dist.init_process_group(backend='nccl')
rank = dist.get_rank()
world_size = dist.get_world_size()

# Пример AllReduce
tensor = torch.tensor([rank]).cuda()
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
print(f"Rank {rank}: after all_reduce = {tensor.item()}")

Термины

• NVLink — высокоскоростное соединение GPU-GPU (до 900 ГБ/с на H100).
• InfiniBand — сетевая технология для соединения узлов (до 400 Гбит/с).
• Ring AllReduce — алгоритм, при котором данные передаются по кольцу GPU, минимизируя количество передач.

2. Tensor Parallelism (тензорный параллелизм)

Tensor Parallelism (TP) — способ распределения модели, при котором один слой (например, линейный слой или attention) разрезается на части по измерению тензора (обычно по hidden dimension). Каждая GPU хранит и вычисляет свою часть. Для выполнения forward pass требуется обмен данными между GPU.

Пример: Пусть есть линейный слой y = xW, где W — матрица весов размером (d_in, d_out). При TP с world_size=2:

• GPU 0 хранит W[:, :d_out/2]
• GPU 1 хранит W[:, d_out/2:]
• Каждая GPU получает полный вход x, вычисляет свою часть выхода, затем выполняется AllGather для получения полного y.

Другие варианты TP

• Column-wise TP (как выше) — разрезание по выходному измерению.
• Row-wise TP — разрезание по входному измерению (требует ReduceScatter).
• 2D TP — комбинация обоих.

Сравнение стратегий параллелизма

3. Зачем NCCL для tensor parallelism

Tensor parallelism требует интенсивного обмена данными между GPU на каждом шаге forward/backward. Без эффективной библиотеки накладные расходы на коммуникацию могут превысить выигрыш от параллелизма.

Роль NCCL в TP

• AllGather — после вычисления части выхода нужно собрать полный тензор для следующего слоя. NCCL реализует AllGather с пропускной способностью, близкой к пиковой скорости NVLink.
• ReduceScatter — при row-wise TP перед вычислением градиентов нужно суммировать частичные результаты и распределить их. NCCL оптимизирует эту операцию.
• AllReduce — может использоваться для синхронизации градиентов, если TP комбинируется с Data Parallelism.

Пример кода с PyTorch FSDP + TP (упрощённо):

# Инициализация NCCL
torch.distributed.init_process_group(backend='nccl')

# Разделение весов линейного слоя
class LinearTP(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.world_size = dist.get_world_size()
        self.rank = dist.get_rank()
        # Каждая GPU хранит часть весов
        self.weight = torch.nn.Parameter(
            torch.randn(in_features, out_features // self.world_size)
        )
    
    def forward(self, x):
        # x: (batch, in_features) — полный вход
        # Локальное вычисление
        local_out = torch.mm(x, self.weight)
        # AllGather для получения полного выхода
        full_out = [torch.zeros_like(local_out) for _ in range(self.world_size)]
        dist.all_gather(full_out, local_out)
        return torch.cat(full_out, dim=1)

Почему именно NCCL, а не MPI или собственные реализации:

• NCCL оптимизирован для GPU: использует CUDA kernels, NVLink, GPUDirect RDMA (прямой доступ к памяти GPU через InfiniBand).
• Автоматический выбор алгоритма (Ring, Tree, NVLink-only) в зависимости от размера сообщения и топологии.
• Низкая задержка (latency) и высокая пропускная способность (bandwidth) для типичных размеров тензоров в TP (от нескольких МБ до ГБ).

Сравнение NCCL с другими библиотеками

4. Практические аспекты использования NCCL для TP

Настройка

• Убедиться, что NCCL установлен (обычно входит в CUDA Toolkit).
• Установить переменные окружения для диагностики: NCCL_DEBUG=INFO, NCCL_DEBUG_SUBSYS=ALL.
• Для TP внутри одного узла достаточно NVLink; для межузлового TP нужен InfiniBand и NCCL_SOCKET_IFNAME.

Проблемы и решения

• Out of Memory при AllGather — можно использовать sequence parallelism (разрезать по длине последовательности).
• Deadlock — все GPU должны участвовать в коллективной операции; асинхронные вызовы требуют синхронизации.
• Пропускная способность — для маленьких тензоров (менее 1 МБ) latency доминирует; для больших — bandwidth. NCCL автоматически переключает алгоритм.

Мониторинг:

• nsys (NVIDIA Nsight Systems) для профилирования коммуникаций.
• nvidia-smi topo -m для просмотра топологии NVLink.

5. Связь NCCL и tensor parallelism с Agentic RAG

В контексте Agentic RAG (системы, где агенты используют LLM для принятия решений) большие модели (например, Llama 3 70B, GPT-4) требуют распределённого инференса. Tensor parallelism с NCCL позволяет:

• Уместить модель на нескольких GPU (если одна GPU не вмещает).
• Уменьшить latency инференса за счёт параллельных вычислений.
• Масштабировать агентов на кластере GPU.

Без NCCL tensor parallelism был бы неэффективен, и агенты не могли бы работать с большими моделями в реальном времени.

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

Задача Реализовать простой tensor parallelism для линейного слоя на двух GPU с использованием NCCL и измерить ускорение.

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

Шаги:

1. Написать скрипт, который инициализирует torch.distributed с backend='nccl'.
2. Создать класс LinearTP, который разрезает веса по выходному измерению.
3. В forward выполнить локальное умножение, затем all_gather для сборки полного выхода.
4. Сравнить время выполнения с обычным линейным слоем на одной GPU (без распределения).
5. Измерить пропускную способность NCCL с помощью torch.cuda.Event.

Ожидаемый результат

• Для больших матриц (например, in_features=4096, out_features=8192) TP на двух GPU должен быть быстрее, чем на одной, за счёт параллелизма.
• Для маленьких матриц накладные расходы на коммуникацию могут перевесить выигрыш.

Пример кода (фрагмент):

import torch
import torch.distributed as dist
import time

def benchmark_tp():
    dist.init_process_group(backend='nccl')
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    torch.cuda.set_device(rank)

    in_features = 4096
    out_features = 8192
    batch = 32

    # TP слой
    weight = torch.randn(in_features, out_features // world_size).cuda()
    x = torch.randn(batch, in_features).cuda()

    # Тепловой запуск
    for _ in range(10):
        local_out = torch.mm(x, weight)
        out_list = [torch.zeros_like(local_out) for _ in range(world_size)]
        dist.all_gather(out_list, local_out)
        full_out = torch.cat(out_list, dim=1)

    # Измерение
    torch.cuda.synchronize()
    start = time.time()
    for _ in range(100):
        local_out = torch.mm(x, weight)
        dist.all_gather(out_list, local_out)
        full_out = torch.cat(out_list, dim=1)
    torch.cuda.synchronize()
    elapsed = time.time() - start

    if rank == 0:
        print(f"TP with {world_size} GPUs: {elapsed/100:.6f} sec per forward")

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

8. Навигация

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

Навигация

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