Распараллеливание тензоров — torch.distributed.tensor.parallel

Распараллеливание тензоров (TP) основано на распределённом тензоре PyTorch (DTensor) и предоставляет несколько стилей распараллеливания: по строкам, по столбцам и по парам.

Точка входа для распараллеливания вашей nn.Module с помощью распараллеливания тензоров:

torch.distributed.tensor.parallel.parallelize_module(module, device_mesh, parallelize_plan, tp_mesh_dim=0) [source]

API для применения распараллеливания тензоров (TP) в PyTorch. Мы распараллеливаем модуль или подмодули на основе плана распараллеливания. План распараллеливания содержит ParallelStyle, который указывает, как пользователь хочет распараллелить модуль или подмодуль.

Пользователь также может указать различные стили распараллеливания для каждого имени модуля с полным квалификатором (FQN). API поддерживает 2D-распараллеливание напрямую, принимая n-мерную сеть устройств, и пользователям нужно только указать размерность, по которой мы выполняем распараллеливание тензоров.

Параметры

• module (nn.Module) – Модуль, который нужно распараллелить.
• device_mesh (DeviceMesh) – Объект, описывающий топологию сети устройств для DTensor.
• parallelize_plan (Union[ParallelStyle, Dict[str, ParallelStyle]]) – План, используемый для распараллеливания модуля. Он может быть либо объектом ParallelStyle, который содержит способ подготовки входных/выходных данных для распараллеливания тензоров, либо словарем из FQN модуля и соответствующего объекта ParallelStyle.
• tp_mesh_dim (int) – Размерность device_mesh, по которой мы выполняем распараллеливание тензоров.

Возвращаемое значение

Объект nn.Module распараллелен.

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

Модуль

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, PairwiseParallel
>>>
>>> # Define the module.
>>> m = Model(...)
>>> m = parallelize_module(m, PairwiseParallel())
>>>

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

PairwiseParallel на данный момент имеет ограничения. Если вам нужна более мелкая детализация, вам нужно передать словарь FQN модуля и стиль параллелизма.

Распараллеливание тензоров поддерживает следующие стили распараллеливания:

class torch.distributed.tensor.parallel.style.RowwiseParallel(_prepare_input=<function make_input_shard_1d_last_dim>, _prepare_output=<function make_output_tensor>) [source]

Разбиение строки модуля. Мы предполагаем, что вход представляет собой фрагментированный DTensor, а выход — torch.Tensor.

class torch.distributed.tensor.parallel.style.ColwiseParallel(_prepare_input=<function make_input_replicate_1d>, _prepare_output=<function make_sharded_output_tensor>) [source]

Разбиение столбца тензора или модуля. Мы предполагаем, что вход представляет собой дублированный DTensor, а выход — фрагментированный torch.Tensor.

class torch.distributed.tensor.parallel.style.PairwiseParallel(_prepare_input=None, _prepare_output=None) [source]

PairwiseParallel объединяет стили по столбцам и по строкам как фиксированную пару, как это делает Megatron-LM (https://arxiv.org/abs/1909.08053). Мы предполагаем, что как входные, так и выходные данные должны быть дублированными DTensor.

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

PairwiseParallel на данный момент не поддерживает nn.MultiheadAttention, nn.Transformer хорошо. Одним из решений является применение ColwiseParallel и RowwiseParallel к компонентам трансформера. Мы рекомендуем использовать PairwiseParallel только для MLP с чётным числом слоёв.

class torch.distributed.tensor.parallel.style.SequenceParallel [source]

SequenceParallel объединяет стили по столбцам и по строкам как фиксированную пару вместе с последовательным распараллеливанием, как это делает Megatron-LM Sequence parallel (https://arxiv.org/pdf/2205.05198.pdf). Мы предполагаем, что как входные, так и выходные данные должны быть фрагментированными DTensor.

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

SequenceParallel на данный момент не поддерживает nn.MultiheadAttention, nn.Transformer хорошо. Одним из решений является применение ColwiseParallel и RowwiseParallel к компонентам трансформера. Мы рекомендуем использовать SequenceParallel только для MLP с чётным числом слоёв.

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

torch.distributed.tensor.parallel.style.make_input_replicate_1d(input, device_mesh=None) [source]

Дублирование входного тензора по одномерной сети устройств. Эта функция будет использоваться в ParallelStyle.

Параметры

• input (Union[torch.Tensor, DTensor]) – Этот входной тензор будет дублирован по одномерной DeviceMesh.
• device_mesh (DeviceMesh, optional) – Одномерная сеть устройств, где input будет дублирован. Если сеть устройств не передана и input является DTensor, будет использоваться input.device_mesh. Если DeviceMesh не одномерная, будет выброшено исключение. По умолчанию: None

Возвращаемое значение

DTensor, дублированный по device_mesh.

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

DTensor

torch.distributed.tensor.parallel.style.make_input_reshard_replicate(input, device_mesh) [source]

Для построения фрагментированного DTensor из тензора на разных рангах и последующего преобразования в дублированный DTensor.

Параметры

• input (torch.Tensor) – Входной тензор на каждом ранге, состоящий из глобального фрагментированного DTensor по размерности 0 по одномерной DeviceMesh, а затем фрагментированный DTensor преобразуется в дублированный DTensor.
• device_mesh (DeviceMesh, optional) – Одномерная сеть устройств, где input будет фрагментирован. Если DeviceMesh не одномерная, будет выброшено исключение. По умолчанию: None

Возвращаемое значение

A DTensor sharded on dimension 0 over device_mesh

и затем преобразуется в дублированный.

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

DTensor

torch.distributed.tensor.parallel.style.make_input_shard_1d(input, device_mesh=None, dim=0) [source]

Фрагментация входного тензора по dim по одномерной сети устройств. Эта функция будет использоваться в ParallelStyle.

Параметры

• input (Union[torch.Tensor, DTensor]) – Одиночный тензор будет фрагментирован по размерности dim по одномерной DeviceMesh.
• device_mesh (DeviceMesh, optional) – Одномерная сеть устройств, где input будет фрагментирован. Если сеть устройств не передана и input является DTensor, будет использоваться input.device_mesh. Если DeviceMesh не одномерная, будет выброшено исключение. По умолчанию: None
• dim (int, optional) – Размерность фрагментации тензора input. По умолчанию: 0

Возвращаемое значение

Фрагментированный DTensor по размерности dim по device_mesh.

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

DTensor

torch.distributed.tensor.parallel.style.make_input_shard_1d_last_dim(input, device_mesh=None) [source]

Функция-обёртка для make_input_shard_1d с dim = -1.

Параметры

• input (Union[torch.Tensor, DTensor]) – Этот единственный тензор будет фрагментирован по последнему измерению по одномерной DeviceMesh.
• device_mesh (DeviceMesh, optional) – Одномерная сетка устройств, на которой input будет фрагментирован. Если сетка устройств не передаётся, и input — это DTensor, будет использоваться input.device_mesh. Если DeviceMesh не является одномерной, будет выброшено исключение. По умолчанию: None

Возвращаемое значение

Фрагментированный DTensor по последнему измерению по device_mesh.

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

DTensor

torch.distributed.tensor.parallel.style.make_output_replicate_1d(output, device_mesh=None) [source]

Преобразует выходной DTensor в дублирующий DTensor. Это будет использовано в ParallelStyle.

Параметры

• output (DTensor) – Выход модуля, который нужно преобразовать.
• device_mesh (DeviceMesh, optional) – Объект, необходимый для дублирования вывода. Он должен быть одномерной device_mesh, и при передаче неодномерной device_mesh будет выброшено исключение. Если сетка устройств не передаётся, будет использована сетка из выходного значения. По умолчанию: None

Возвращаемое значение

Объект DTensor , сделанный дублирующим.

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

DTensor

torch.distributed.tensor.parallel.style.make_output_reshard_tensor(output, device_mesh=None) [source]

Преобразует выходной DTensor в фрагментированный DTensor и возвращает локальный тензор.

Параметры

• output (DTensor) – Выход модуля, который нужно преобразовать.
• device_mesh (DeviceMesh, optional) – Объект, необходимый для фрагментации вывода. Он должен быть одномерной device_mesh, и при передаче неодномерной device_mesh будет выброшено исключение. Если сетка устройств не передаётся, будет использована сетка из выходного значения. По умолчанию: None

Возвращаемое значение

Объект torch.Tensor, преобразованный из выходного DTensor.

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

Tensor

torch.distributed.tensor.parallel.style.make_output_shard_1d(output, device_mesh=None, dim=0) [source]

Преобразует выходной DTensor в фрагментированный DTensor. Это будет использовано в ParallelStyle.

Параметры

• output (DTensor) – Выход модуля, который нужно преобразовать.
• device_mesh (DeviceMesh, optional) – Объект, необходимый для фрагментации вывода. Он должен быть одномерной device_mesh, и при передаче неодномерной device_mesh будет выброшено исключение. Если сетка устройств не передаётся, будет использована сетка из выходного значения. По умолчанию: None
• dim (int) – Измерение фрагментации вывода. По умолчанию: 0

Возвращаемое значение

Объект DTensor , фрагментированный по заданному измерению.

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

DTensor

torch.distributed.tensor.parallel.style.make_output_tensor(output, device_mesh=None) [source]

Преобразует выходной DTensor в дублирующий DTensor, а затем в Tensor.

Параметры

• output (DTensor) – Выход модуля, который нужно преобразовать.
• device_mesh (DeviceMesh, optional) – Объект, необходимый для дублирования вывода. Он должен быть одномерной device_mesh, и при передаче неодномерной device_mesh будет выброшено исключение. Если сетка устройств не передаётся, будет использована сетка из выходного значения. По умолчанию: None

Возвращаемое значение

Объект torch.Tensor, преобразованный из выходного DTensor.

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

Tensor

В настоящее время существуют некоторые ограничения, которые затрудняют работу модуля MultiheadAttention с параллелизмом тензоров "из коробки", поэтому мы рекомендуем пользователям попробовать ColwiseParallel и RowwiseParallel для каждого параметра. Возможно, потребуется изменить код, так как теперь параллелизация выполняется по первому измерению модуля MultiheadAttention.

Мы также поддерживаем 2D-параллелизм, где мы комбинируем параллелизм тензоров с параллелизмом данных. Чтобы интегрировать с FullyShardedDataParallel, пользователям нужно явно вызвать следующую API:

torch.distributed.tensor.parallel.fsdp.enable_2d_with_fsdp() [source]

API регистрирует расширение, необходимое для работы параллелизма тензоров (TP) с FullyShardedDataParallel (FSDP). Сначала мы параллелизуем параметры внутри одного модуля или подмодулей на основе плана параллелизации и позволим FSDP перефрагментировать локальный тензор распределённого параметра, который по сути является DTensor.

Возвращаемое значение

Логическое значение bool, указывающее на успех или неудачу регистрации расширения.

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

bool

Чтобы интегрировать с DistributedDataParallel, пользователям нужно явно вызвать следующую API:

torch.distributed.tensor.parallel.ddp.pre_dp_module_transform(module) [source]

Включает возможность совместного использования параллелизма тензоров (TP) и параллелизма данных (DP) в PyTorch при использовании DDP. Мы должны преобразовать параметры, являющиеся DTensor, в локальные тензоры перед применением API параллелизма данных. Затем мы регистрируем два обработчика: один для преобразования локальных тензоров обратно в DTensor до вычисления и один для преобразования DTensor обратно в тензоры после вычисления. Таким образом, мы избегаем специального обработки параметров DTensor в DDP и получаем распространение градиентов DTensor обратно к DP, например, градиентные буферы DDP.

Пока эта API работает только с DistributedDataParallel. В дальнейшем она будет поддерживать и другие методы DP, такие как FSDP.

Параметры

module (nn.Module) – Модуль, к которому применён TP.

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, PairwiseParallel
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> from torch.distributed.tensor.parallel.ddp import pre_dp_module_transform
>>>
>>> # Define the module.
>>> m = module(...)
>>> parallelize_module(m, PairwiseParallel())
>>> m = pre_dp_module_transform(m)
>>> m = DDP(m)
>>>