Распределенное обучение PyTorch 2026: стратегии DP, FSDP, TP, PP с кодом | AiManual
AiManual Logo Ai / Manual.
12 Апр 2026 • Инструмент

Как работает распределённое обучение в PyTorch: разбор стратегий DP, FSDP, TP и PP на примерах кода

Глубокий разбор распределенного обучения в PyTorch. Data Parallelism, Fully Sharded DP, Tensor и Pipeline Parallelism на актуальных примерах кода для больших яз

Когда одной видеокарты мало: зачем нужны все эти сложности

Ваша модель Llama 4 80B не помещается в память даже самой навороченной H100. Вы пробовали уменьшить batch size до одного, но тогда обучение растянется на годы. Знакомая история? В 2026 году модели уже перестали быть просто «большими» – они стали чудовищно огромными. К счастью, в PyTorch есть целый арсенал стратегий, чтобы раскидать вычисления по нескольким GPU. И нет, это не просто «включи DDP и забудь».

💡
Если вы только начинаете, сначала прочитайте про то, как GPU общаются между собой. Без понимания NVLink и сети все дальнейшие оптимизации – стрельба из пушки по воробьям.

Data Parallelism (DP): призрак из прошлого

Помните старый добрый DataParallel? В 2026 году он уже не просто устарел – он опасен. PyTorch до сих пор поддерживает его для обратной совместимости, но использовать в продакшене – все равно что ездить на работу на паровозе.

Как он работает? Один GPU (главный) хранит полную копию модели. Во время forward pass батч данных разбивается на части и рассылается на все доступные GPU. Каждый GPU вычисляет градиенты для своей части, после чего они собираются на главном GPU, усредняются, и обновление применяется к главной модели. Затем новая версия модели рассылается обратно на все GPU.

Основная проблема – узкое горлышко на главном GPU. Вся коммуникация идет через него, и при большом количестве карт он превращается в пробку. Скорость обучения упрется в пропускную способность PCIe и память главной карты. На практике линейное ускорение заканчивается на 4-8 GPU.

# КАК НЕ НАДО ДЕЛАТЬ в 2026 году
import torch.nn as nn
import torch

# Модель автоматически реплицируется на все видимые GPU
model = nn.DataParallel(MyModel().cuda())

# Обучение выглядит привычно, но под капотом – кошмар
for data, target in dataloader:
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

Вы все еще видите этот код в старых туториалах. Удалите его и никогда не используйте для реальных задач.

Distributed Data Parallel (DDP): рабочий лошадь

DDP – это то, что используют все, кто хоть раз запускал обучение на кластере. В отличие от DP, здесь нет главного GPU. Каждая карта хранит свою полную копию модели и работает с частью данных. Градиенты усредняются между всеми процессами с помощью коллективных операций all-reduce.

Звучит просто? На бумаге – да. На практике вас ждут сетевые лаги, падения узлов и магические ошибки синхронизации.

1Базовый шаблон DDP в PyTorch 2.4

# main.py
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

def main(rank, world_size):
    # Инициализация процесса
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # Модель, данные, оптимизатор – каждый процесс создает свои
    model = MyModel().cuda(rank)
    ddp_model = DDP(model, device_ids=[rank])
    
    optimizer = torch.optim.Adam(ddp_model.parameters())
    dataloader = create_distributed_dataloader(world_size, rank)
    
    for epoch in range(epochs):
        for batch in dataloader:
            optimizer.zero_grad()
            output = ddp_model(batch)
            loss = compute_loss(output)
            loss.backward()
            # DDP автоматически синхронизирует градиенты!
            optimizer.step()
    
    dist.destroy_process_group()

if __name__ == "__main__":
    # Запуск: torchrun --nproc_per_node=4 main.py
    import sys
    rank = int(sys.argv[1])
    world_size = int(sys.argv[2])
    main(rank, world_size)

Ключевая магия DDP – в хуке, который автоматически вызывает all-reduce для градиентов после backward(). Вам не нужно делать это вручную. Но за эту магию приходится платить: каждая карта должна хранить полную копию модели, оптимизатора и его состояния. Для модели на 70 миллиардов параметров это сразу отсекает всех, у кого нет батареи из дорогих GPU с огромной памятью.

Fully Sharded Data Parallel (FSDP): когда память кончилась

FSDP – это ответ PyTorch на тренировку моделей, которые не помещаются в память ни на одной карте. Идея проста до гениальности: давайте разобьем параметры модели, градиенты и состояние оптимизатора на части (шарды) и распределим их по всем GPU. В каждый момент времени на конкретной карте находятся только те шарды, которые нужны для текущего вычисления.

В PyTorch 2.4 API FSDP значительно упростился, но документация все еще вызывает легкую панику.

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
import torch.nn as nn

def main(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # Политика автоматического обертывания: оборачиваем слои с >100M параметров
    auto_wrap_policy = size_based_auto_wrap_policy(min_num_params=100_000_000)
    
    model = MyGiantModel()
    # Модель должна быть на CPU перед обертыванием!
    fsdp_model = FSDP(
        model,
        auto_wrap_policy=auto_wrap_policy,
        sharding_strategy=ShardingStrategy.FULL_SHARD,  # Шардим все: параметры, градиенты, состояние оптимизатора
        device_id=rank
    )
    
    optimizer = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-4)
    
    for batch in dataloader:
        optimizer.zero_grad()
        output = fsdp_model(batch.cuda(rank))
        loss = output.sum()
        loss.backward()
        # Градиенты собираются автоматически
        optimizer.step()
        # Освобождаем непоследовательные блоки памяти (важно для эффективности)
        torch.cuda.empty_cache()

Самая частая ошибка – пытаться переместить модель на GPU до обертывания в FSDP. Делать это категорически нельзя. Сначала создаете модель на CPU, потом оборачиваете в FSDP, и только тогда она попадет на GPU.

FSDP не бесплатен. За экономию памяти вы платите увеличением коммуникации. Вместо одного all-reduce для градиентов, как в DDP, FSDP выполняет множество all-gather и reduce-scatter операций для шардов. На медленной сети это может съесть всю выгоду.

Tensor Parallelism (TP): режем слои пополам

Представьте линейный слой с матрицей весов размером 8192x8192. Что если разрезать ее по столбцам и раздать две половинки разным GPU? Каждая карта будет вычислять свою часть выхода, а результаты потом соединятся. Это и есть Tensor Parallelism.

Раньше для этого нужны были монструозные фреймворки вроде Megatron-LM. Сейчас в PyTorch 2.4 есть torch.distributed.tensor.parallel.

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
)
from torch.distributed.device_mesh import init_device_mesh

dist.init_process_group("nccl")
# Создаем сетку устройств 1D: разбиваем по одному измерению
device_mesh = init_device_mesh("cuda", (2,))  # 2 GPU

class SimpleMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1024, 2048)
        self.fc2 = nn.Linear(2048, 512)
    
    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = SimpleMLP().cuda()
# Параллелизуем модуль по плану:
# fc1 режем по столбцам (ColwiseParallel), fc2 по строкам (RowwiseParallel)
parallel_plan = {
    "fc1": ColwiseParallel(),  # Веса [1024, 2048] -> на 2 GPU: [1024, 1024] на каждом
    "fc2": RowwiseParallel(),  # Веса [2048, 512] -> на 2 GPU: [1024, 512] на каждом
}
parallel_model = parallelize_module(
    model,
    device_mesh,
    parallel_plan,
)

# Теперь forward pass автоматически распределяется
input = torch.randn(32, 1024).cuda()
output = parallel_model(input)  # Вычисления и коммуникация скрыты

Tensor Parallelism требует, чтобы карты были соединены высокоскоростными линиями (NVLink). Каждый слой порождает коммуникацию между GPU. Если у вас 8 карт в режиме TP, то для одного forward pass может потребоваться 7 операций обмена данными. На медленной сети это убьет производительность.

Pipeline Parallelism (PP): конвейер как на заводе

Самая сложная для отладки стратегия. Идея: разбиваем модель на последовательные этапы (например, первые 10 слоев на GPU 0, следующие 10 на GPU 1 и т.д.). Данные проходят через этот конвейер как детали по сборочной линии.

Проблема в том, что в наивной реализации большая часть GPU простаивает: пока GPU 1 обрабатывает первый микро-батч, GPU 0 уже ждет. Решение – Pipeline Parallelism с чередованием микро-батчей (interleaved), но код для этого выглядит так, что хочется плакать.

В PyTorch 2.4 есть экспериментальная поддержка через torch.distributed.pipeline.sync.Pipe.

# Внимание: код упрощен, в реальности нужно учитывать чередование батчей и балансировку
import torch
import torch.nn as nn
from torch.distributed.pipeline.sync import Pipe

dist.init_process_group("nccl")

# Разбиваем модель на 4 последовательных блока
class LayerChunk1(nn.Module):
    def forward(self, x):
        return x * 2

class LayerChunk2(nn.Module):
    def forward(self, x):
        return x + 1

# Создаем модель-конвейер
model = nn.Sequential(LayerChunk1(), LayerChunk2()).cuda()
# Обертываем в Pipe. Каждый модуль из Sequential попадет на отдельный GPU.
# Требуется, чтобы world_size был равен числу chunks (2 в данном случае).
pipeline_model = Pipe(model, chunks=4)  # chunks - число микро-батчей

input = torch.randn(16, 10).cuda(0)  # Тензор должен быть на первом устройстве
# Forward pass автоматически проходит по конвейеру
output = pipeline_model(input)

Pipeline Parallelism – это про балансировку. Если один этап конвейера выполняется в 2 раза дольше остальных, все остальные GPU будут ждать. Подбор оптимального разбиения модели – это отдельное искусство, часто требующее профайлинга и угадывания.

Что выбрать? Сравнительная таблица

СтратегияКогда использоватьГлавная больАктуальность в 2026
DPНикогда. Серьезно, забудьте.Узкое горлышко на главном GPUУстарело
DDPМодель помещается в память одной карты, нужно ускорить обучение.Синхронизация и отладка на кластереПромышленный стандарт
FSDPМодель НЕ помещается в память одной карты (например, Llama 70B).Сложная настройка, оверхед коммуникацииКритически важно для больших моделей
TPОтдельные слои модели слишком велики (большие матрицы в MLP).Требует NVLink, сложность реализацииИспользуется в гибридных схемах
PPМодель очень глубокая (сотни слоев), а GPU мало.Балансировка конвейера, простои устройствНишевое применение

В реальности никто не использует одну стратегию. Гибридные схемы – это норма. Например, FSDP поверх TP для модели с гигантскими слоями. Или DDP внутри узла с 8 GPU и Pipeline Parallelism между узлами. В гайде по тонкой настройке больших моделей мы как раз разбираем подобные комбинации.

Совет, который сэкономит вам месяц отладки

Не пытайтесь реализовать распределенное обучение с нуля, если только вы не работаете в команде, разрабатывающей новый фреймворк. Возьмите готовое решение: Hugging Face Accelerate, DeepSpeed (который, к слову, прекрасно интегрируется с PyTorch) или хотя бы изучите продакшен-пайплайн на DDP.

Но даже с готовыми инструментами подготовьтесь к тому, что 80% времени уйдет не на обучение, а на настройку коммуникации и отладку распределенного кода. И да, аренда мощного кластера через провайдеров вроде CloudProvider может быть дешевле, чем покупка собственных карт, но учтите – их инфраструктура не всегда оптимизирована для low-latency коммуникации между узлами.

Мой прогноз? К 2027 году все эти стратегии окончательно абстрагируются за высокоуровневыми API. Вы будете просто указывать, сколько GPU у вас есть и какого размера модель, а фреймворк сам выберет оптимальную схему распределения. Но пока что – придется разбираться в деталях. Удачи, и пусть ваши градиенты всегда синхронизируются.

Подписаться на канал

←
На главную