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

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

Параллелизм модели с использованием нескольких графических процессоров

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

На рисунке представлена модель с 4 слоями, размещёнными на 4 разных графических процессорах (вертикальная ось). Горизонтальная ось представляет обучение этой модели во времени, демонстрируя, что используется только один графический процессор в один момент (источник изображения).

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

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

На рисунке представлена модель с 4 слоями, размещёнными на 4 разных графических процессорах (вертикальная ось). Горизонтальная ось представляет обучение этой модели во времени, демонстрируя, что графические процессоры используются гораздо эффективнее. Однако по-прежнему существует «пузырь» (как показано на рисунке), где определённые графические процессоры не используются. (источник изображения).

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

class torch.distributed.pipeline.sync.Pipe(module, chunks=1, checkpoint='except_last', deferred_batch_norm=False) [source]

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

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

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

Вы должны разместить все модули на соответствующих устройствах и обернуть их в модуль 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 на графических процессорах 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) [source]

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

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

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

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

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

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

Параметры

inputs – входной мини-пакет

Возвращает

RRef до выхода мини-пакета

Возможные исключения

TypeError – вход не содержит по крайней мере один тензор

Тип возвращаемого значения

RRef

Соединения пропуска

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

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

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

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

Вот пример с тремя слоями. Тензор пропуска с именем «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 модуль может закреплять или извлекать несколько тензоров пропуска:

@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

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

Return type

Callable[[Type[Module]], Type[Skippable]]

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

Команда для закрепления тензора пропуска.

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

Parameters

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

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

Команда для извлечения тензора пропуска.

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

Parameters

name (str) – имя тензора пропуска

Returns

тензор пропуска, ранее закреплённый другим слоем под тем же именем

Return type

None

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

Проверяет целостность skippable модулей.

Каждый тензор пропуска должен иметь только одну пару 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())
#             ^^^^^^      └───────────────────┘

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

Raises

TypeError – одна или более пар stash и pop не совпадают.

Учебники

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

• Обучение моделей Transformer с использованием параллелизма по конвейеру
• Обучение моделей Transformer с использованием Distributed Data Parallel и параллелизма по конвейеру

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

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