Параллелизм конвейера

Параллелизм конвейера был первоначально представлен в статье Gpipe и является эффективной техникой для обучения больших моделей на нескольких GPU.

Предупреждение

Параллелизм конвейера — экспериментальная функция и может быть изменён.

Параллелизм модели с использованием нескольких GPU

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

_images/no_pipe.png — Рисунок представляет модель с 4 слоями, размещёнными на 4 разных GPU (вертикальная ось). Горизонтальная ось представляет обучение этой модели со временем, демонстрируя, что только 1 GPU используется в каждый момент времени (источник изображения).

Выполнение в конвейере

Для решения этой проблемы параллелизм конвейера разбивает входной мини-батч на несколько микро-батчей и выполняет конвейерную обработку этих микро-батчей на нескольких GPU. Это показано на рисунке ниже:

_images/pipe.png — Рисунок представляет модель с 4 слоями, размещёнными на 4 разных GPU (вертикальная ось). Горизонтальная ось представляет обучение этой модели со временем, демонстрируя, что GPU используются гораздо эффективнее. Однако всё ещё существует «пузырь» (как показано на рисунке), где некоторые GPU не используются. (источник изображения).

API конвейера в PyTorch

class torch.distributed.pipeline.sync.Pipe(module, chunks=1, checkpoint='except_last', deferred_batch_norm=False) [источник]

Оборачивает произвольный модуль nn.Sequential для обучения с использованием синхронного параллелизма конвейера. Если модуль требует много памяти и не помещается на одном GPU, параллелизм конвейера является полезной техникой для обучения.

Реализация основана на статье torchgpipe.

Pipe объединяет параллелизм конвейера с сохранением контрольных точек (checkpointing) для уменьшения пиковой памяти, необходимой для обучения, минимизируя недоиспользование устройств.

Вы должны разместить все модули на соответствующих устройствах и обернуть их в модуль nn.Sequential, определяющий желаемый порядок выполнения. Если модуль не содержит параметров/буферов, предполагается, что этот модуль должен выполняться на CPU, и соответствующие входные тензоры модуля перемещаются на CPU перед выполнением. Это поведение может быть переопределено оболочкой WithDevice, которая может быть использована для явного указания устройства, на котором должен выполняться модуль.

Параметры:

module (nn.Sequential) – последовательный модуль, который нужно распараллелить с помощью конвейеризации. Каждый модуль в последовательности должен иметь все свои параметры на одном устройстве. Каждый модуль в последовательности должен быть nn.Module или nn.Sequential (для объединения нескольких последовательных модулей на одном устройстве)
chunks (int) – количество микро-батчей (по умолчанию: 1)
checkpoint (str) – когда включить сохранение контрольных точек, один из 'always', 'except_last', или 'never' (по умолчанию: 'except_last'). 'never' полностью отключает сохранение контрольных точек, 'except_last' включает сохранение контрольных точек для всех микро-батчей, кроме последнего, и 'always' включает сохранение контрольных точек для всех микро-батчей.
deferred_batch_norm (bool) – использовать отложенное BatchNorm перемещение статистик (по умолчанию: False). Если установлено True, отслеживаются статистики по нескольким микро-батчам для обновления текущих статистик за каждый мини-батч.

Исключения:

TypeError – модуль не является nn.Sequential.
ValueError – неверные аргументы

Пример::

Конвейер из двух слоёв полносвязного типа (FC) на GPU 0 и 1.

>>> # Need to initialize RPC framework first.
>>> os.environ['MASTER_ADDR'] = 'localhost'
>>> os.environ['MASTER_PORT'] = '29500'
>>> torch.distributed.rpc.init_rpc('worker', rank=0, world_size=1)
>>>
>>> # Build pipe.
>>> fc1 = nn.Linear(16, 8).cuda(0)
>>> fc2 = nn.Linear(8, 4).cuda(1)
>>> model = nn.Sequential(fc1, fc2)
>>> model = Pipe(model, chunks=8)
>>> input = torch.rand(16, 16).cuda(0)
>>> output_rref = model(input)

Примечание

Вы можете обернуть модель Pipe с torch.nn.parallel.DistributedDataParallel только когда параметр checkpoint у Pipe равен 'never'.

Примечание

Pipe в настоящее время поддерживает внутриузловую конвейеризацию, но в будущем будет расширена для поддержки межузловой конвейеризации. Функция forward возвращает RRef для поддержки межузловой конвейеризации в будущем, где результат может находиться на удалённом узле. Для внутриузловой конвейеризации вы можете использовать local_value() для получения результата локально.

Предупреждение

Pipe — экспериментальная функция и может быть изменена.

forward(*inputs) [источник]

Обрабатывает один входной мини-батч через конвейер и возвращает RRef, указывающий на результат. Pipe — это довольно прозрачный модуль-обёртка. Он не изменяет сигнатуру ввода и вывода базового модуля. Но существуют ограничения по типу. Вход и вывод должны содержать по крайней мере один тензор. Это ограничение применяется и для разбиений.

Последовательность входов подаётся в первую стадию конвейера как *inputs. В результате позиционные аргументы этой функции должны соответствовать позиционным аргументам первой стадии конвейера. То же самое условие применимо и к выводу одной стадии конвейера, который является входом для следующей стадии.

Входящий тензор разбивается на несколько микро-батчей на основе параметра chunks, используемого при инициализации Pipe. Размер пакета предполагается как первое измерение тензора, и если размер пакета меньше chunks, количество микро-батчей равно размеру пакета.

Только тензоры разбиваются на несколько микро-батчей, не-тензорные входные данные просто дублируются в каждом микро-баче. Для не-тензорных выходных данных на последней стадии конвейера они агрегируются как List и возвращаются пользователю. Например, если у вас есть 2 микро-бача, возвращающих целое число 5, пользователь получит объединённый результат [5, 5]

Все входные тензоры должны находиться на том же устройстве, что и первое разбиение конвейера.

Если тензор обернут с помощью NoChunk оболочки, тензор не делится между микро-батчами и дублируется как есть, так же как и не-тензоры.

Параметры:: inputs – входной мини-батч
Возвращаемое значение:: RRef к выводу мини-бача
Исключения:: TypeError – вход не содержит по крайней мере один тензор
Тип возвращаемого значения:: RRef

Соединения "skip"

Некоторые модели, такие как ResNeXt, не являются полностью последовательными и имеют соединения "skip" между слоями. При простом внедрении в параллелизм по каналам предполагается, что нам нужно копировать выходные данные определённых слоёв на несколько графических процессоров (GPU), пока мы не достигнем GPU, на котором находится слой для соединения "skip". Чтобы избежать этой накладной стоимости копирования, мы предоставляем ниже API для хранения и извлечения тензоров в различных слоях модели.

torch.distributed.pipeline.sync.skip.skippable.skippable(stash=(), pop=()) [source]

Декоратор для определения nn.Module с соединениями "skip". Модули, к которым применяется этот декоратор, называются «skippable». Данная функциональность работает без проблем даже тогда, когда модуль не обернут Pipe.

Каждый тензор "skip" управляется своим именем. Перед обработкой тензоров "skip", модуль «skippable» должен статически объявить имена тензоров "skip", используя параметры stash и/или pop. Тензоры "skip" с предварительно объявленным именем могут быть сохранены с помощью yield stash(name, tensor) или извлечены с помощью tensor = yield pop(name).

Вот пример с тремя слоями. Тензор "skip" с именем «1to3» хранится и извлекается соответственно в первом и последнем слоях:

@skippable(stash=['1to3'])
class Layer1(nn.Module):
    def forward(self, input):
        yield stash('1to3', input)
        return f1(input)

class Layer2(nn.Module):
    def forward(self, input):
        return f2(input)

@skippable(pop=['1to3'])
class Layer3(nn.Module):
    def forward(self, input):
        skip_1to3 = yield pop('1to3')
        return f3(input) + skip_1to3

model = nn.Sequential(Layer1(), Layer2(), Layer3())

Один модуль «skippable» может хранить или извлекать несколько тензоров "skip":

@skippable(stash=['alice', 'bob'], pop=['carol'])
class StashStashPop(nn.Module):
    def forward(self, input):
        yield stash('alice', f_alice(input))
        yield stash('bob', f_bob(input))
        carol = yield pop('carol')
        return input + carol

Каждый тензор "skip" должен быть связан ровно с одной парой stash и pop. Pipe автоматически проверяет это ограничение при обёртке модуля. Вы также можете проверить это ограничение, используя verify_skippables() без Pipe.

Тип возвращаемого значения:: Callable[[Тип[Модуль]], Тип[Skippable]]

class torch.distributed.pipeline.sync.skip.skippable.stash(name, tensor) [source]

Команда для хранения тензора "skip".

def forward(self, input):
    yield stash('name', input)
    return f(input)

Параметры:

name (строка) – имя тензора "skip"
input (torch.Tensor или None) – тензор для передачи в соединение "skip"

class torch.distributed.pipeline.sync.skip.skippable.pop(name) [source]

Команда для извлечения тензора "skip".

def forward(self, input):
    skip = yield pop('name')
    return f(input) + skip

Параметры:: name (строка) – имя тензора "skip"
Возвращаемое значение:: тензор "skip", предварительно сохранённый другим слоем под тем же именем
Тип возвращаемого значения:: None

torch.distributed.pipeline.sync.skip.skippable.verify_skippables(module) [source]

Проверяет, удовлетворяют ли базовые модули «skippable» требованиям целостности.

Каждый тензор "skip" должен иметь только одну пару stash и pop. Если есть одна или несколько несоответствующих пар, будет возбуждено исключение TypeError с подробными сообщениями.

Вот несколько случаев ошибок. verify_skippables() сообщит об ошибке в этих случаях:

# Layer1 stashes "1to3".
# Layer3 pops "1to3".

nn.Sequential(Layer1(), Layer2())
#               └──── ?

nn.Sequential(Layer2(), Layer3())
#                   ? ────┘

nn.Sequential(Layer1(), Layer2(), Layer3(), Layer3())
#               └───────────────────┘       ^^^^^^

nn.Sequential(Layer1(), Layer1(), Layer2(), Layer3())
#             ^^^^^^      └───────────────────┘

Чтобы использовать одно и то же имя для нескольких тензоров "skip", они должны быть изолированы в разных пространствах имён. См. isolate().

Возбуждаемые исключения:: TypeError – одна или несколько пар stash и pop не совпадают.

Учебники

Следующие учебники дают хорошее представление о том, как использовать Pipe API для обучения моделей с использованием остальных компонентов, предоставляемых PyTorch:

Благодарности

Реализация параллелизма по каналам основана на реализации fairscale pipe и torchgpipe. Мы хотели бы поблагодарить обе команды за их вклад и руководство по внедрению параллелизма по каналам в PyTorch.

© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/1.13/pipeline.html