d2l 12.5–12.6: обучение на нескольких GPU

О переводе

Полный перевод двух разделов главы 12 «Производительность вычислений» из репозитория d2l-ai/d2l-zh (Apache-2.0): multiple-gpus.md и multiple-gpus-concise.md. В оригинале у каждого фрагмента кода есть варианты для MXNet, PyTorch и PaddlePaddle — мы оставили только PyTorch (как в видеокурсе). Ссылки вида «см. раздел …» сохранены как номера разделов книги. Рисунки не копируем — описываем и перерисовываем схемой. Вывод программ (графики, время эпохи) в исходниках книги не хранится — поэтому чисел мы не приводим. Наши пояснения — «Примечание переводчика». Видео Ли Му к этим разделам: «Обучение на нескольких GPU» и «Несколько GPU: реализация» (см. путеводитель).

12.5 Обучение на нескольких GPU

До сих пор мы обсуждали, как эффективно обучать модели на CPU и GPU, в разделе 12.3 показали, как фреймворки глубокого обучения автоматически распараллеливают вычисления и обмен данными между CPU и GPU, а в разделе 5.6 — как командой nvidia-smi вывести список всех доступных GPU компьютера. Но мы не обсуждали, как на самом деле распараллелить обучение. Можно ли как-то разделить данные между несколькими устройствами так, чтобы всё правильно работало? В этом разделе подробно показано, как с нуля обучать сеть параллельно; для этого понадобится стохастический градиентный спуск по мини-батчам (см. раздел 11.5). Дальше мы покажем, как обучать сеть параллельно с помощью высокоуровневых API (см. раздел 12.6).

Как разделить задачу

Начнём с простой задачи компьютерного зрения и чуть устаревшей сети. В сети несколько свёрточных слоёв и слоёв пулинга, а в конце, возможно, несколько полносвязных слоёв — она очень похожа на LeNet (LeCun и др., 1998) или AlexNet (Krizhevsky и др., 2012). Предположим, у нас несколько GPU (2 у настольного сервера, 4 на AWS g4dn.12xlarge, 8 на p3.16xlarge, 16 на p2.16xlarge). Мы хотим разделить обучение так, чтобы получить хорошее ускорение и при этом обойтись простыми и воспроизводимыми проектными решениями. В конце концов, несколько GPU дают больше и памяти, и вычислительной мощности. Коротко: для мини-батча обучающих данных, которые нужно классифицировать, у нас есть следующие варианты.

Первый способ — разделить сеть между GPU. То есть каждый GPU получает данные, входящие в определённый слой, обрабатывает их несколькими последующими слоями и отправляет на следующий GPU. Так можно обрабатывать данные сетью большего размера, чем поместилась бы на один GPU. К тому же объём занимаемой памяти (memory footprint) каждого GPU хорошо контролируется — это лишь часть памяти всей сети.

Однако плотная синхронизация между интерфейсами GPU может стать проблемой, особенно когда вычислительная нагрузка слоёв плохо сбалансирована, а также когда на стыках слоёв нужно передавать много данных (например, активации и градиенты; объём может превысить пропускную способность шины GPU). Кроме того, для разбиения очень важен порядок вычислительно тяжёлых операций — лучшее исследование на эту тему: Mirhoseini и др. (2017); по сути это остаётся трудной задачей, и неясно, можно ли на конкретной задаче добиться хорошего линейного масштабирования. В итоге этот способ не рекомендуется, если только фреймворк или операционная система сами не умеют связывать несколько GPU вместе.

Второй способ — разделить работу внутри слоя. Например, вместо того чтобы считать 64 канала на одном GPU, распределить задачу на 4 GPU, каждый из которых вычисляет 16 каналов. Для полносвязных слоёв так же можно разделить число выходных нейронов. На рис. 12.5.1 показан такой дизайн: эта стратегия применялась для GPU с очень маленькой памятью (тогда — 2 ГБ). Когда каналов или нейронов не слишком мало, это хорошо повышает производительность. Кроме того, поскольку доступная память растёт линейно, несколько GPU могут обрабатывать всё более крупные сети.

Рис. 12.5.1 в оригинале (img/alexnet-original.svg): параллелизм модели в исходном дизайне AlexNet, вызванный нехваткой памяти GPU, — сеть разрезана на две половины по каналам, каждая на своём GPU, с обменом данными между ними на некоторых слоях.

Однако здесь нужно очень много синхронизации, или барьерных операций (barrier operation), поскольку каждый слой зависит от результатов всех остальных. К тому же передавать нужно, возможно, даже больше данных, чем при разбиении слоёв между GPU. Поэтому из-за стоимости пропускной способности и сложности мы этот способ тоже не рекомендуем.

Последний способ — разделить данные между несколькими GPU. В этом случае все GPU выполняют одну и ту же работу, но над разными наблюдениями. После обработки каждого мини-батча градиенты агрегируются по GPU. Этот способ самый простой и применим в любой ситуации; синхронизация нужна только после обработки каждого мини-батча. Причём уже посчитанные градиенты параметров можно начинать передавать, пока другие ещё вычисляются. И чем больше GPU, тем больше данных в мини-батче — тем выше эффективность обучения. Но добавление GPU не позволяет обучать более крупные модели.

flowchart LR
subgraph A["Сеть по GPU<br/>(конвейер)"]
  A1["GPU 1: слои 1–2"] --> A2["GPU 2: слои 3–4"]
end
subgraph B["Внутри слоя<br/>(тензорный параллелизм)"]
  B1["GPU 1: каналы 1–32"]
  B2["GPU 2: каналы 33–64"]
end
subgraph C["По данным"]
  C1["GPU 1: вся сеть,<br/>половина батча"]
  C2["GPU 2: вся сеть,<br/>половина батча"]
  C1 <-. "сумма градиентов" .-> C2
end

По рис. 12.5.2 «Параллелизм на нескольких GPU»: слева направо — разбиение сети, разбиение слоёв, параллелизм по данным (схема переводчика)

На рис. 12.5.2 сравниваются разные способы параллелизма на нескольких GPU. В целом, если памяти GPU хватает, параллелизм по данным — самый удобный. Подробное описание разбиения при распределённом обучении — в Li и др. (2014). В ранние годы глубокого обучения память GPU была острой проблемой, но сегодня, кроме совсем особых случаев, она решена. Дальше мы сосредоточимся на параллелизме по данным.

Примечание переводчика: «память решена» — это про 2021 год

Фраза о том, что проблема памяти «решена», верна для CNN на картинках из этой главы. Для больших языковых моделей всё наоборот: модель на десятки миллиардов параметров не помещается на одну карту, и два «нерекомендованных» способа стали основными: разбиение по слоям — это конвейерный параллелизм (pipeline parallelism), разбиение внутри слоя — тензорный параллелизм (Megatron-LM, --tensor-parallel-size в vLLM). Плюс ZeRO/FSDP — разновидность параллелизма по данным, где параметры и состояния оптимизатора тоже делятся между GPU. Всё это — урок 5.4.3.

Параллелизм по данным

Пусть в машине $k$ GPU. Для обучаемой модели каждый GPU независимо хранит полный набор её параметров, хотя значения параметров на всех GPU одинаковы и синхронизированы. Например, рис. 12.5.3 показывает обучение с параллелизмом по данным при $k=2$.

Рис. 12.5.3 в оригинале (img/data-parallel.svg): мини-батч делится пополам; каждый GPU считает свою половину градиента, градиенты складываются, сумма рассылается обоим GPU, и каждый обновляет свою копию параметров.

В общем случае параллельное обучение на $k$ GPU идёт так:

• на каждой итерации случайный мини-батч делится на $k$ частей и поровну распределяется по GPU;
• каждый GPU по своей части мини-батча вычисляет потерю и градиенты параметров модели;
• локальные градиенты всех $k$ GPU агрегируются — получается стохастический градиент текущего мини-батча;
• агрегированный градиент рассылается обратно каждому GPU;
• каждый GPU по этому стохастическому градиенту обновляет полный набор параметров, который он хранит.

На практике учтите: при обучении на $k$ GPU размер мини-батча нужно увеличить в $k$ раз, чтобы у каждого GPU была та же нагрузка, что и при обучении на одном GPU. Поэтому на 16-GPU сервере можно заметно увеличить мини-батч, и, возможно, придётся соответственно увеличить скорость обучения. Учтите также, что пакетную нормализацию из раздела 7.5 нужно адаптировать — например, хранить для каждого GPU отдельные параметры пакетной нормализации.

Примечание переводчика: почему это математически то же самое

Градиент потери по мини-батчу — это сумма градиентов по примерам. Если батч из $B$ примеров разрезать на две половины, то $\nabla L = \sum_{i=1}^{B} \nabla \ell_i = \sum_{i \in \text{GPU}_1} \nabla \ell_i + \sum_{i \in \text{GPU}_2} \nabla \ell_i$. Поэтому 2 GPU с батчем 128 на каждом дают тот же шаг, что 1 GPU с батчем 256 (если нет пакетной нормализации, которая смотрит на статистику батча, — о ней и предупреждение выше).

Дальше мы покажем обучение на нескольких GPU на простой сети.

%matplotlib inline
from d2l import torch as d2l
import torch
from torch import nn
from torch.nn import functional as F

Простая сеть

Возьмём (немного изменённую) LeNet из раздела 6.6 и определим её с нуля, чтобы подробно показать обмен параметрами и синхронизацию.

# Инициализируем параметры модели
scale = 0.01
W1 = torch.randn(size=(20, 1, 3, 3)) * scale
b1 = torch.zeros(20)
W2 = torch.randn(size=(50, 20, 5, 5)) * scale
b2 = torch.zeros(50)
W3 = torch.randn(size=(800, 128)) * scale
b3 = torch.zeros(128)
W4 = torch.randn(size=(128, 10)) * scale
b4 = torch.zeros(10)
params = [W1, b1, W2, b2, W3, b3, W4, b4]

# Определяем модель
def lenet(X, params):
    h1_conv = F.conv2d(input=X, weight=params[0], bias=params[1])
    h1_activation = F.relu(h1_conv)
    h1 = F.avg_pool2d(input=h1_activation, kernel_size=(2, 2), stride=(2, 2))
    h2_conv = F.conv2d(input=h1, weight=params[2], bias=params[3])
    h2_activation = F.relu(h2_conv)
    h2 = F.avg_pool2d(input=h2_activation, kernel_size=(2, 2), stride=(2, 2))
    h2 = h2.reshape(h2.shape[0], -1)
    h3_linear = torch.mm(h2, params[4]) + params[5]
    h3 = F.relu(h3_linear)
    y_hat = torch.mm(h3, params[6]) + params[7]
    return y_hat

# Функция потерь — перекрёстная энтропия
loss = nn.CrossEntropyLoss(reduction='none')

Примечание переводчика: откуда 800

Вход Fashion-MNIST — 1×28×28. Свёртка 3×3 → 20×26×26, пулинг 2×2 → 20×13×13, свёртка 5×5 → 50×9×9, пулинг → 50×4×4 (округление вниз). $50 \cdot 4 \cdot 4 = 800$ — это и есть входная размерность W3. Как считать размеры после свёртки — урок 3.5.1.

Синхронизация данных

Для эффективного обучения на нескольких GPU нужны две базовые операции. Во-первых, нужно раздать параметры нескольким устройствам и подключить к ним градиенты (get_params). Без параметров сеть на GPU не посчитать. Во-вторых, нужно складывать параметры между устройствами — то есть нужна функция allreduce.

def get_params(params, device):
    new_params = [p.to(device) for p in params]
    for p in new_params:
        p.requires_grad_()
    return new_params

Скопируем параметры модели на один GPU.

new_params = get_params(params, d2l.try_gpu(0))
print('b1 веса:', new_params[1])
print('b1 градиент:', new_params[1].grad)

Поскольку вычислений ещё не было, градиенты параметров пока нулевые. Теперь пусть у нас есть вектор, распределённый по нескольким GPU; функция allreduce ниже складывает все векторы и рассылает результат всем GPU. Обратите внимание: чтобы функция работала, данные нужно копировать на устройство, где накапливается результат.

def allreduce(data):
    for i in range(1, len(data)):
        data[0][:] += data[i].to(data[0].device)
    for i in range(1, len(data)):
        data[i][:] = data[0].to(data[i].device)

Примечание переводчика

Строго говоря, до первого backward() градиент в PyTorch не «нулевой», а равен None — именно это и напечатает вторая строка. А allreduce здесь — самая наивная схема: всё собираем на GPU 0, затем рассылаем оттуда. При $k$ картах GPU 0 принимает и отправляет по $k-1$ копии — он становится «бутылочным горлышком». Библиотека NCCL, которую использует PyTorch, делает кольцевой all-reduce: каждая карта передаёт лишь около $2(k-1)/k$ объёма градиентов, независимо от числа карт. Об этом третье упражнение в конце раздела.

Проверим: создадим на разных устройствах векторы с разными значениями и агрегируем их.

data = [torch.ones((1, 2), device=d2l.try_gpu(i)) * (i + 1) for i in range(2)]
print('до allreduce:\n', data[0], '\n', data[1])
allreduce(data)
print('после allreduce:\n', data[0], '\n', data[1])

Примечание переводчика: что напечатается

До: на GPU 0 — [[1., 1.]], на GPU 1 — [[2., 2.]]. После: на обоих — [[3., 3.]]. Если у вас одна видеокарта, d2l.try_gpu(1) вернёт CPU, и пример всё равно сработает — просто «вторым устройством» будет процессор.

Распределение данных

Нужна простая вспомогательная функция, которая поровну распределяет мини-батч по нескольким GPU. Например, при двух GPU мы хотим, чтобы каждый получил половину данных. Встроенной функцией фреймворка это писать удобнее и короче, поэтому попробуем её на матрице $4 \times 5$.

data = torch.arange(20).reshape(4, 5)
devices = [torch.device('cuda:0'), torch.device('cuda:1')]
split = nn.parallel.scatter(data, devices)
print('input :', data)
print('load into', devices)
print('output:', split)

Для повторного использования определим функцию split_batch, которая разделяет сразу и данные, и метки.

#@save
def split_batch(X, y, devices):
    """Разделить X и y по нескольким устройствам"""
    assert X.shape[0] == y.shape[0]
    return (nn.parallel.scatter(X, devices),
            nn.parallel.scatter(y, devices))

Примечание переводчика

nn.parallel.scatter режет тензор по первой оси (по примерам) на столько частей, сколько устройств: матрица 4×5 на двух GPU превратится в два куска 2×5 — строки 0–1 на cuda:0, строки 2–3 на cuda:1. Этот пример требует двух видеокарт; #@save — пометка авторов: функция сохраняется в пакет d2l для следующих разделов.

Обучение

Теперь можно реализовать обучение на нескольких GPU для одного мини-батча. Для синхронизации данных между GPU используем только что разобранные вспомогательные функции allreduce и split_and_load. Писать какой-то особый код для параллелизма не нужно: внутри мини-батча вычислительные графы на разных устройствах не зависят друг от друга, поэтому выполняются параллельно «автоматически».

def train_batch(X, y, device_params, devices, lr):
    X_shards, y_shards = split_batch(X, y, devices)
    # Потерю считаем на каждом GPU отдельно
    ls = [loss(lenet(X_shard, device_W), y_shard).sum()
          for X_shard, y_shard, device_W in zip(
              X_shards, y_shards, device_params)]
    for l in ls:  # обратное распространение — на каждом GPU отдельно
        l.backward()
    # Складываем все градиенты со всех GPU и рассылаем результат всем GPU
    with torch.no_grad():
        for i in range(len(device_params[0])):
            allreduce(
                [device_params[c][i].grad for c in range(len(devices))])
    # Параметры модели обновляем на каждом GPU отдельно
    for param in device_params:
        d2l.sgd(param, lr, X.shape[0])  # здесь используем размер полного мини-батча

Примечание переводчика

split_and_load в тексте — остаток от версии на MXNet; в PyTorch-коде роль этой функции играет split_batch. Деление на X.shape[0] (размер всего батча, а не куска) в d2l.sgd и даёт ровно тот же шаг, что на одном GPU: сумма градиентов по всем кускам делится на общий размер батча. И обратите внимание: «автоматический» параллелизм здесь обеспечивает асинхронность CUDA — Python быстро ставит в очередь работу для обеих карт, а карты выполняют её одновременно (раздел 12.2 книги).

Теперь можно определить функцию обучения. Она немного отличается от функций из предыдущих глав: нужно выделить GPU и скопировать все параметры модели на все устройства. Понятно, что каждый мини-батч обрабатывается на нескольких GPU функцией train_batch. Точность модели мы считаем только на одном GPU, а остальные в это время простаивают; это относительно неэффективно, зато удобно и код короче.

def train(num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    # Копируем параметры модели на num_gpus GPU
    device_params = [get_params(params, d) for d in devices]
    num_epochs = 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    timer = d2l.Timer()
    for epoch in range(num_epochs):
        timer.start()
        for X, y in train_iter:
            # Обучение на нескольких GPU для одного мини-батча
            train_batch(X, y, device_params, devices, lr)
            torch.cuda.synchronize()
        timer.stop()
        # Оцениваем модель на GPU 0
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(
            lambda x: lenet(x, device_params[0]), test_iter, devices[0]),))
    print(f'точность на тесте: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} с/эпоха, '
          f'на {str(devices)}')

Посмотрим, насколько хорошо всё работает на одном GPU. Сначала возьмём размер батча 256 и скорость обучения 0,2.

train(num_gpus=1, batch_size=256, lr=0.2)

Оставим размер батча и скорость обучения прежними и увеличим число GPU до 2: точность на тесте будет примерно такой же, как в предыдущем эксперименте. С точки зрения поиска оптимума алгоритм при разном числе GPU одинаков. К сожалению, никакого заметного ускорения здесь нет: модель слишком маленькая, да и датасет тоже. На таком датасете наш простой способ обучения на нескольких GPU страдает от огромных накладных расходов Python. В будущем нам встретятся более сложные модели и более сложные способы параллелизма. А пока посмотрим, что произойдёт на датасете Fashion-MNIST.

train(num_gpus=2, batch_size=256, lr=0.2)

Итоги

• Обучение глубокой сети можно разделить между несколькими GPU разными способами: между слоями, внутри слоёв или по данным. Первые два требуют тщательной оркестровки передачи данных, а последний — самая простая стратегия.
• Обучение с параллелизмом по данным само по себе несложно; оно повышает эффективность обучения за счёт увеличения эффективного размера мини-батча.
• При параллелизме по данным данные делятся между несколькими GPU, каждый GPU выполняет свои прямой и обратный проходы, затем все градиенты агрегируются в один, и результат рассылается всем GPU.
• При большем мини-батче скорость обучения нужно немного увеличить.

Упражнения

1. При обучении на $k$ GPU увеличьте размер батча с $b$ до $k \cdot b$, то есть пропорционально числу GPU.
2. Сравните точность модели при разных скоростях обучения. Как скорость обучения должна масштабироваться с ростом числа GPU?
3. Реализуйте более эффективную функцию allreduce для агрегации разных параметров на разных GPU. Почему она эффективнее?
4. Реализуйте подсчёт точности на тесте на нескольких GPU.

Обсуждение (форум d2l)

---

12.6 Краткая реализация на нескольких GPU

Реализовывать параллельные вычисления с нуля для каждой новой модели — скучно. К тому же полезно оптимизировать средства синхронизации ради высокой производительности. Ниже покажем, как сделать это высокоуровневыми API фреймворка. Математика и алгоритм — те же, что в разделе 12.5. Для запуска кода этого раздела нужны как минимум два GPU.

from d2l import torch as d2l
import torch
from torch import nn

Простая сеть

Возьмём сеть содержательнее, чем LeNet из раздела 12.5, но всё ещё лёгкую и быстро обучаемую. Мы выбрали ResNet-18 (He и др., 2016). Поскольку входные изображения маленькие, мы её немного изменили. Отличие от раздела 7.6: в начале используются меньшие ядро свёртки, шаг и дополнение, а слой max-пулинга убран.

#@save
def resnet18(num_classes, in_channels=1):
    """Немного изменённая модель ResNet-18"""
    def resnet_block(in_channels, out_channels, num_residuals,
                     first_block=False):
        blk = []
        for i in range(num_residuals):
            if i == 0 and not first_block:
                blk.append(d2l.Residual(in_channels, out_channels,
                                        use_1x1conv=True, strides=2))
            else:
                blk.append(d2l.Residual(out_channels, out_channels))
        return nn.Sequential(*blk)

    # В модели меньшие ядро свёртки, шаг и дополнение, а слой max-пулинга убран
    net = nn.Sequential(
        nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm2d(64),
        nn.ReLU())
    net.add_module("resnet_block1", resnet_block(
        64, 64, 2, first_block=True))
    net.add_module("resnet_block2", resnet_block(64, 128, 2))
    net.add_module("resnet_block3", resnet_block(128, 256, 2))
    net.add_module("resnet_block4", resnet_block(256, 512, 2))
    net.add_module("global_avg_pool", nn.AdaptiveAvgPool2d((1,1)))
    net.add_module("fc", nn.Sequential(nn.Flatten(),
                                       nn.Linear(512, num_classes)))
    return net

Инициализация сети

Мы будем инициализировать сеть внутри цикла обучения. Методы инициализации можно повторить в разделе 4.8.

net = resnet18(10)
# Получаем список GPU
devices = d2l.try_all_gpus()
# Сеть мы инициализируем в коде обучения

Обучение

Как уже говорилось, чтобы обучение шло эффективно и параллельно, код обучения должен выполнять несколько базовых функций:

• инициализировать параметры сети на всех устройствах;
• при проходе по датасету распределять мини-батчи по всем устройствам;
• параллельно на всех устройствах считать потерю и её градиент;
• агрегировать градиенты и соответственно обновлять параметры.

Наконец, параллельно считать точность и сообщать итоговое качество сети. Кроме того, что данные нужно делить и агрегировать, код обучения очень похож на реализации из предыдущих глав.

def train(net, num_gpus, batch_size, lr):
    train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)
    devices = [d2l.try_gpu(i) for i in range(num_gpus)]
    def init_weights(m):
        if type(m) in [nn.Linear, nn.Conv2d]:
            nn.init.normal_(m.weight, std=0.01)
    net.apply(init_weights)
    # Размещаем модель на нескольких GPU
    net = nn.DataParallel(net, device_ids=devices)
    trainer = torch.optim.SGD(net.parameters(), lr)
    loss = nn.CrossEntropyLoss()
    timer, num_epochs = d2l.Timer(), 10
    animator = d2l.Animator('epoch', 'test acc', xlim=[1, num_epochs])
    for epoch in range(num_epochs):
        net.train()
        timer.start()
        for X, y in train_iter:
            trainer.zero_grad()
            X, y = X.to(devices[0]), y.to(devices[0])
            l = loss(net(X), y)
            l.backward()
            trainer.step()
        timer.stop()
        animator.add(epoch + 1, (d2l.evaluate_accuracy_gpu(net, test_iter),))
    print(f'точность на тесте: {animator.Y[0][-1]:.2f}, {timer.avg():.1f} с/эпоха, '
          f'на {str(devices)}')

Посмотрим, как это работает на практике. Для разминки обучим сеть на одном GPU.

train(net, num_gpus=1, batch_size=256, lr=0.1)

Теперь обучим на 2 GPU. По сравнению с LeNet из раздела 12.5 модель ResNet-18 гораздо сложнее. Вот где видно преимущество параллелизма: вычисления занимают заметно больше времени, чем синхронизация параметров. Поскольку накладные расходы на параллелизм теперь менее существенны, масштабируемость модели улучшается.

train(net, num_gpus=2, batch_size=512, lr=0.2)

Примечание переводчика: что делает nn.DataParallel

На каждом шаге nn.DataParallel сам делает всё то, что мы писали руками в разделе 12.5: копирует модель на все указанные GPU (replicate), режет батч (scatter), запускает прямой проход параллельно в потоках, собирает выходы на GPU 0 (gather), а при backward() складывает градиенты на исходной модели. Обратите внимание: при переходе на 2 GPU авторы удвоили и батч (256 → 512), и скорость обучения (0,1 → 0,2) — то самое «линейное масштабирование» из упражнения 2 раздела 12.5.

Итоги

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

Упражнения

1. В этом разделе используется ResNet-18. Попробуйте разное число эпох, размеры батча и скорости обучения, а также больше GPU. Что будет, если попробовать на 16 GPU (например, на экземпляре AWS p2.16xlarge)?
2. Иногда устройства имеют разную вычислительную мощность: можно одновременно использовать GPU и CPU. Как тогда распределять работу? Почему?

Обсуждение (форум d2l)

---

От переводчика: как это делают сегодня и что запустить у себя

nn.DataParallel устарел — используйте DDP

Документация PyTorch давно рекомендует вместо nn.DataParallel использовать DistributedDataParallel (DDP), даже на одной машине. DataParallel работает в одном процессе Python с потоками (мешает GIL), каждый шаг заново копирует модель на все карты и собирает выходы на GPU 0 — поэтому GPU 0 загружен сильнее остальных. DDP запускает по процессу на карту, модель копируется один раз, а градиенты складываются кольцевым all-reduce (NCCL) прямо во время обратного прохода, слой за слоем.

Минимальный каркас того же обучения на DDP (запуск: torchrun --nproc_per_node=4 train_ddp.py):

# train_ddp.py — набросок переводчика; модель и данные возьмите из примеров выше
import os
import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import torchvision
from torchvision import transforms
from d2l import torch as d2l                   # нужен для d2l.Residual внутри resnet18

dist.init_process_group("nccl")                 # torchrun сам задаёт RANK, WORLD_SIZE, LOCAL_RANK
local_rank = int(os.environ["LOCAL_RANK"])      # номер карты для этого процесса: 0, 1, 2 или 3
torch.cuda.set_device(local_rank)

train_set = torchvision.datasets.FashionMNIST(
    root="./data", train=True, download=(local_rank == 0), transform=transforms.ToTensor())
dist.barrier()                                  # остальные процессы ждут, пока процесс 0 скачает данные
sampler = DistributedSampler(train_set)         # каждый процесс получает свою 1/4 датасета
loader = DataLoader(train_set, batch_size=128, sampler=sampler)  # 128 на карту → 512 всего

model = resnet18(10).cuda(local_rank)           # функция resnet18 из раздела 12.6
model = DDP(model, device_ids=[local_rank])
opt = torch.optim.SGD(model.parameters(), lr=0.2)
loss_fn = nn.CrossEntropyLoss()

for epoch in range(10):
    sampler.set_epoch(epoch)                    # новое перемешивание на каждой эпохе
    for X, y in loader:
        X, y = X.cuda(local_rank), y.cuda(local_rank)
        opt.zero_grad()
        loss_fn(model(X), y).backward()         # здесь DDP делает all-reduce градиентов
        opt.step()
    if local_rank == 0:
        print(f"эпоха {epoch + 1} готова")

dist.destroy_process_group()

🖥 На вашей машине (4×GPU)

• У вас 4 карты, связанные только PCIe 4.0 (без NVLink). Для ResNet-18 на Fashion-MNIST это не помеха: градиентов всего около 11 млн чисел (≈45 МБ в float32) на шаг. А вот для больших моделей обмен градиентами по PCIe станет заметной частью времени шага — поэтому в LLM-обучении на такой машине так важны накопление градиента (реже синхронизироваться) и bfloat16 (вдвое меньше данных). Подробно — урок 5.7.4.
• Запустите train(num_gpus=1, ...), train(num_gpus=2, ...) и train(num_gpus=4, ...) из раздела 12.6 и сравните время эпохи. Затем то же самое на DDP. Ожидайте, что DDP масштабируется лучше, а для LeNet из 12.5 ускорения почти не будет — как и пишут авторы.
• Проверить, как карты связаны между собой: nvidia-smi topo -m (у вас в матрице будут обозначения вроде NODE, PHB или SYS — связь через PCIe и процессор, — а не NV#, как при NVLink).

Словарик

Русский	English
параллелизм по данным	data parallelism
параллелизм модели	model parallelism
конвейерный параллелизм	pipeline parallelism
видеопамять	GPU memory
синхронизация	synchronization
барьерная операция	barrier
агрегация, редукция	aggregate / reduce
рассылка (broadcast)	broadcast
ускорение	speedup
линейное масштабирование	linear scaling
пропускная способность	bandwidth