Тензорный параллелизм — torch.distributed.tensor.parallel

Создано: 13 июня 2025 г. | Последнее обновление: 12 мая 2026 г.

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

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

torch.distributed.tensor.parallel.parallelize_module(module, device_mesh=None, parallelize_plan=None, *, src_data_rank=0) [источник]

Применяет тензорный параллелизм в PyTorch, распараллеливая модули или подмодули в соответствии с планом, заданным пользователем.

Мы распараллеливаем модуль или подмодули на основе parallelize_plan. parallelize_plan содержит ParallelStyle, указывающий, как пользователь хочет распараллелить модуль или подмодуль.

Пользователь также может указать разные стили параллелизма для каждого полного квалифицированного имени модуля (FQN).

Обратите внимание, что parallelize_module принимает только одномерный DeviceMesh. Если у вас есть двумерный или N-мерный DeviceMesh, сначала выполните срез DeviceMesh, чтобы получить одномерный подчинённый DeviceMesh, а затем передайте его этому API (т. е. device_mesh["tp"]).

Параметры:

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

Именованные аргументы:

src_data_rank (int, необязательный) – ранг исходных данных логического/глобального тензора. Он используется distribute_tensor() для распределения/широковещательной передачи фрагментов/реплик на другие ранги. По умолчанию в качестве источника данных для каждого измерения DeviceMesh используется group_rank=0, чтобы сохранить семантику одного устройства. Если явно передать None, parallelize_module() просто использует локальные данные, вместо того чтобы пытаться сохранить семантику одного устройства с помощью распределения/широковещательной передачи. По умолчанию: 0

Возвращает:

Распараллеленный объект nn.Module.

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

Module

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>>
>>> # Define the module.
>>> m = Model(...)
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>> m = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel(), "w2": RowwiseParallel()})
>>>

Примечание

Для сложных архитектур модулей, таких как Attention и слои MLP, рекомендуем объединять разные ParallelStyles (например, ColwiseParallel и RowwiseParallel) и передавать их в качестве parallelize_plan, чтобы добиться требуемого разбиения на шарды.

torch.distributed.tensor.parallel.api.parallelize_module(module, device_mesh=None, parallelize_plan=None, *, src_data_rank=0) [источник]

Применяет тензорный параллелизм в PyTorch, распараллеливая модули или подмодули в соответствии с планом, заданным пользователем.

Мы распараллеливаем модуль или подмодули на основе parallelize_plan. parallelize_plan содержит ParallelStyle, указывающий, как пользователь хочет распараллелить модуль или подмодуль.

Пользователь также может указать разные стили параллелизма для каждого полного квалифицированного имени модуля (FQN).

Обратите внимание, что parallelize_module принимает только одномерный DeviceMesh. Если у вас есть двумерный или N-мерный DeviceMesh, сначала выполните срез DeviceMesh, чтобы получить одномерный подчинённый DeviceMesh, а затем передайте его этому API (т. е. device_mesh["tp"]).

Параметры:

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

Именованные аргументы:

src_data_rank (int, необязательный) – ранг исходных данных логического/глобального тензора. Он используется distribute_tensor() для распределения/широковещательной передачи фрагментов/реплик на другие ранги. По умолчанию в качестве источника данных для каждого измерения DeviceMesh используется group_rank=0, чтобы сохранить семантику одного устройства. Если явно передать None, parallelize_module() просто использует локальные данные, вместо того чтобы пытаться сохранить семантику одного устройства с помощью распределения/широковещательной передачи. По умолчанию: 0

Возвращает:

Распараллеленный объект nn.Module.

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

Module

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>>
>>> # Define the module.
>>> m = Model(...)
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>> m = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel(), "w2": RowwiseParallel()})
>>>

Примечание

Для сложных архитектур модулей, таких как Attention и слои MLP, рекомендуем объединять разные ParallelStyles (например, ColwiseParallel и RowwiseParallel) и передавать их в качестве parallelize_plan, чтобы добиться требуемого разбиения на шарды.

Тензорный параллелизм поддерживает следующие стили параллелизма:

class torch.distributed.tensor.parallel.ColwiseParallel(*, input_layouts=None, output_layouts=None, use_local_output=True) [источник]

Разделяет совместимый nn.Module по столбцам. В настоящее время поддерживаются nn.Linear и nn.Embedding. Пользователи могут объединить этот стиль с RowwiseParallel, чтобы разделить на шарды более сложные модули (например, MLP и Attention).

Именованные аргументы:

• input_layouts (Placement, необязательный) – Размещение DTensor входного тензора для nn.Module; используется для аннотирования входного тензора и преобразования его в DTensor. Если параметр не указан, предполагается, что входной тензор реплицирован.
• output_layouts (Placement, необязательный) – Размещение DTensor выходных данных nn.Module; используется для приведения выходных данных nn.Module к требуемому пользователем размещению. Если параметр не указан, выходной тензор разбивается на шарды по последнему измерению.
• use_local_output (bool, необязательный) – Использовать ли локальный torch.Tensor вместо DTensor для выходных данных модуля. Значение по умолчанию: True.

Возвращает:

Объект ParallelStyle, представляющий разбиение nn.Module на шарды по столбцам.

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m is a nn.Module that contains a "w1" nn.Linear submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # By default, the input of the "w1" Linear will be converted to Replicated DTensor
>>> # and the output of "w1" will return :class:`torch.Tensor` that shards on the last dim.
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"w1": ColwiseParallel()})
>>> ...

Примечание

По умолчанию выходные данные ColwiseParallel разбиваются на шарды по последнему измерению, если output_layouts не указано. Если существуют операции, требующие определённой формы тензора (например, перед парным RowwiseParallel), учитывайте, что при разбиении выходных данных на шарды может потребоваться адаптировать операцию под размер шарда.

class torch.distributed.tensor.parallel.RowwiseParallel(*, input_layouts=None, output_layouts=None, use_local_output=True) [источник]

Разделяет совместимый nn.Module по строкам. В настоящее время поддерживаются nn.Linear и nn.Embedding. Пользователи могут объединить этот стиль с ColwiseParallel, чтобы разделить на шарды более сложные модули (например, MLP и Attention).

Именованные аргументы:

• input_layouts (Placement, необязательный) – Размещение DTensor входного тензора для nn.Module; используется для аннотирования входного тензора и преобразования его в DTensor. Если параметр не указан, предполагается, что входной тензор разбит на шарды по последнему измерению.
• output_layouts (Placement, необязательный) – Размещение DTensor выходных данных nn.Module; используется для приведения выходных данных nn.Module к требуемому пользователем размещению. Если параметр не указан, выходной тензор реплицируется.
• use_local_output (bool, необязательный) – Использовать ли локальный torch.Tensor вместо DTensor для выходных данных модуля. Значение по умолчанию: True.

Возвращает:

Объект ParallelStyle, представляющий разбиение nn.Module на шарды по строкам.

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, RowwiseParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m is a nn.Module that contains a "w2" nn.Linear submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # By default, the input of the "w2" Linear will be converted to DTensor that shards on the last dim
>>> # and the output of "w2" will return a replicated :class:`torch.Tensor`.
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"w2": RowwiseParallel()}),
>>> ...

class torch.distributed.tensor.parallel.SequenceParallel(*, sequence_dim=1, use_local_output=False) [источник]

SequenceParallel реплицирует совместимые параметры nn.Module и выполняет вычисления с разделением на шарды, при этом входные данные разбиты на шарды по измерению последовательности. В настоящее время поддерживаются nn.LayerNorm, nn.Dropout и реализация RMSNorm на Python.

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

Если переданные в этот nn.Module входные данные имеют тип torch.Tensor, предполагается, что они уже разбиты на шарды по измерению последовательности; вход преобразуется в DTensor, разбитый на шарды по этому измерению. Если переданные в этот nn.Module входные данные уже являются DTensor, но не разбиты на шарды по измерению последовательности, выполняется перераспределение входных данных для разбиения на шарды по этому измерению.

Выходные данные nn.Module будут разбиты на шарды по измерению последовательности.

Именованные аргументы:

• sequence_dim (int, необязательный) – Измерение последовательности входного тензора для nn.Module; используется для аннотирования входного тензора и преобразования его в DTensor, разбитый на шарды по измерению последовательности. Значение по умолчанию: 1.
• use_local_output (bool, необязательный) – Использовать ли локальный torch.Tensor вместо DTensor для выходных данных модуля. Значение по умолчанию: False.

Возвращает:

Объект ParallelStyle, представляющий параллелизм последовательностей для nn.Module.

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, SequenceParallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> m = Model(...)  # m is a nn.Module that contains a "norm" nn.LayerNorm submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # By default, the input of the "norm" will be converted to DTensor that shards on the sequence dim
>>> # and the output of "norm" will return a sharded on sequence dimension :class:`DTensor`.
>>>
>>> sharded_mod = parallelize_module(m, tp_mesh, {"norm": SequenceParallel()}),
>>> ...

Примечание

Стиль SequenceParallel предполагает инициализацию единицами, если в nn.Module имеются веса (например, nn.LayerNorm или RMSNorm; по умолчанию для них используется инициализация единицами). Если для весов этих модулей используются пользовательские методы инициализации, необходимо выполнить широковещательную передачу весов до или после распараллеливания, чтобы обеспечить их репликацию.

Чтобы просто настроить входные и выходные данные nn.Module с помощью размещений DTensor и выполнить необходимое перераспределение размещения, не преобразуя параметры модуля в DTensor, при вызове parallelize_module в parallelize_plan можно использовать следующие ParallelStyle:

class torch.distributed.tensor.parallel.PrepareModuleInput(*, input_layouts=None, desired_input_layouts=None, input_kwarg_layouts=None, desired_input_kwarg_layouts=None, use_local_output=False) [источник]

Настраивает входные данные nn.Module так, чтобы во время выполнения входные тензоры nn.Module преобразовывались в DTensor в соответствии с input_layouts, а перераспределение размещения выполнялось в соответствии с desired_input_layouts.

Именованные аргументы:

• input_layouts (Union[Placement, Tuple[Optional[Placement]]]) – Размещения DTensor входных тензоров nn.Module; используются для преобразования входных тензоров в DTensor. Если некоторые входные данные не являются torch.Tensor или их не нужно преобразовывать в DTensor, вместо них необходимо указать None. Значение по умолчанию: None.
• desired_input_layouts (Union[Placement, Tuple[Optional[Placement]]]) – Требуемое размещение DTensor входных тензоров nn.Module; используется для обеспечения нужных размещений DTensor для входных данных nn.Module. Длина этого аргумента должна совпадать с длиной input_layouts. Значение по умолчанию: None.
• input_kwarg_layouts (Dict[str, Placement]) – Размещения DTensor именованных входных аргументов nn.Module; используются для преобразования тензоров именованных входных аргументов в DTensor. Значение по умолчанию: None
• desired_input_kwarg_layouts – (Dict[str, Placement]): Требуемое размещение DTensor именованных входных аргументов nn.Module; используется для обеспечения нужных размещений DTensor для входных данных nn.Module. Значение по умолчанию: None.
• use_local_output (bool, необязательный) – Использовать ли локальный torch.Tensor вместо DTensor для входных данных модуля. Значение по умолчанию: False.

Возвращает:

Объект ParallelStyle, который подготавливает размещения шардинга для входных данных nn.Module.

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleInput
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> block = TransformerBlock(...)  # block is a nn.Module that contains an "attn" Attention submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # According to the style specified below, the first input of attn will be annotated to Sharded DTensor
>>> # and then redistributed to Replicated DTensor.
>>> parallelize_module(
>>>     block, # this can be a submodule or module
>>>     tp_mesh,
>>>     parallelize_plan={
>>>         "attn": PrepareModuleInput(
>>>             input_layouts=(Shard(0), None, None, ...),
>>>             desired_input_layouts=(Replicate(), None, None, ...)
>>>         ),
>>>     }
>>> )

class torch.distributed.tensor.parallel.PrepareModuleOutput(*, output_layouts, desired_output_layouts, use_local_output=True) [источник]

Настраивает выходные данные nn.Module так, чтобы во время выполнения выходные тензоры nn.Module преобразовывались в DTensor в соответствии с output_layouts, а перераспределение размещения выполнялось в соответствии с desired_output_layouts.

Именованные аргументы:

• output_layouts (Union[Placement, Tuple[Placement]]) – Размещения DTensor выходных тензоров nn.Module; используются для преобразования выходных тензоров в DTensor, если они являются torch.Tensor. Если некоторые выходные данные не являются torch.Tensor или их не нужно преобразовывать в DTensor, вместо них необходимо указать None.
• desired_output_layouts (Union[Placement, Tuple[Placement]]) – Требуемые размещения DTensor выходных тензоров nn.Module; используются для обеспечения нужных размещений DTensor для выходных данных nn.Module.
• use_local_output (bool, необязательный) – Использовать ли локальный torch.Tensor вместо DTensor для выходных данных модуля. Значение по умолчанию: True.

Возвращает:

Объект ParallelStyle, который подготавливает размещения шардинга для выходных данных nn.Module.

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleOutput
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> block = TransformerBlock(...)  # block is a nn.Module that contains an "attn" Attention submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # According to the style specified below, the output of the TransformerBlock will be converted to Replicated DTensor
>>> # and then redistributed to Sharded DTensor.
>>> parallelize_module(
>>>     block, # this can be a submodule or module
>>>     tp_mesh,
>>>     parallelize_plan = PrepareModuleOutput(
>>>         output_layouts=Replicate(),
>>>         desired_output_layouts=Shard(0)
>>>     )
>>> )

class torch.distributed.tensor.parallel.PrepareModuleInputOutput(*, input_layouts=None, desired_input_layouts=None, input_kwarg_layouts=None, desired_input_kwarg_layouts=None, use_local_input=False, output_layouts, desired_output_layouts, use_local_output=True) [источник]

Настраивает входные данные (и выходные данные) nn.Module так, чтобы во время выполнения входные тензоры (и, соответственно, выходные тензоры) nn.Module преобразовывались в DTensor в соответствии с input_layouts (и, соответственно, output_layouts), а перераспределение размещения выполнялось в соответствии с desired_input_layouts (и, соответственно, desired_output_layouts). Это сочетание PrepareModuleInput и PrepareModuleOutput.

Именованные аргументы:

• input_layouts (Union[Placement, Tuple[Optional[Placement]]]) – Размещения DTensor входных тензоров nn.Module; используются для преобразования входных тензоров в DTensor. Если некоторые входные данные не являются torch.Tensor или их не нужно преобразовывать в DTensor, вместо них необходимо указать None. Значение по умолчанию: None.
• desired_input_layouts (Union[Placement, Tuple[Optional[Placement]]]) – Требуемое размещение DTensor входных тензоров nn.Module; используется для обеспечения нужных размещений DTensor для входных данных nn.Module. Длина этого аргумента должна совпадать с длиной input_layouts. Значение по умолчанию: None.
• input_kwarg_layouts (Dict[str, Placement]) – Размещения DTensor именованных входных аргументов nn.Module; используются для преобразования тензоров именованных входных аргументов в DTensor. Значение по умолчанию: None
• desired_input_kwarg_layouts – (Dict[str, Placement]): Требуемое размещение DTensor именованных входных аргументов nn.Module; используется для обеспечения нужных размещений DTensor для входных данных nn.Module. Значение по умолчанию: None.
• use_local_input (bool, необязательный) – Использовать ли локальный torch.Tensor вместо DTensor для входных данных модуля. Значение по умолчанию: False.
• output_layouts (Union[Placement, Tuple[Placement]]) – Размещения DTensor выходных тензоров nn.Module; используются для преобразования выходных тензоров в DTensor, если они являются torch.Tensor. Если некоторые выходные данные не являются torch.Tensor или их не нужно преобразовывать в DTensor, вместо них необходимо указать None.
• desired_output_layouts (Union[Placement, Tuple[Placement]]) – Требуемые размещения DTensor выходных тензоров nn.Module; используются для обеспечения нужных размещений DTensor для выходных данных nn.Module.
• use_local_output (bool, необязательный) – Использовать ли локальный torch.Tensor вместо DTensor для выходных данных модуля. Значение по умолчанию: True.

Возвращает:

Объект ParallelStyle, который подготавливает размещения шардинга для входных и выходных данных nn.Module.

Пример::

>>> from torch.distributed.tensor.parallel import parallelize_module, PrepareModuleInputOutput
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> block = TransformerBlock(...)  # block is a nn.Module that contains an "attn" Attention submodule
>>> tp_mesh = init_device_mesh("cuda", (8,))
>>>
>>> # According to the style specified below, the first input of attn will be annotated as Sharded DTensor
>>> # and then redistributed to Replicated DTensor, and the output of the TransformerBlock will be annotated
>>> # as Replicated DTensor and then redistributed to Sharded DTensor.
>>> parallelize_module(
>>>     block, # this can be a submodule or module
>>>     tp_mesh,
>>>     parallelize_plan={
>>>         "attn": PrepareModuleInputOutput(
>>>             input_layouts=(Shard(0), None, None, ...),
>>>             desired_input_layouts=(Replicate(), None, None, ...),
>>>             output_layouts=Replicate(),
>>>             desired_output_layouts=Shard(0),
>>>         ),
>>>     }
>>> )

Для таких моделей, как Transformer, рекомендуем использовать ColwiseParallel и RowwiseParallel вместе в parallelize_plan, чтобы добиться требуемого разбиения на шарды всей модели (например, Attention и MLP).

Параллельные вычисления функции потерь кросс-энтропии (параллелизм функции потерь) поддерживаются с помощью следующего менеджера контекста:

torch.distributed.tensor.parallel.loss_parallel() [исходный код]

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

Внутри этого менеджера контекста можно как обычно использовать cross_entropy() или CrossEntropyLoss при следующих предположениях относительно входных параметров. Соответствующий вызов backward(), если он есть, также должен выполняться внутри этого менеджера контекста.

Поддерживаются как одномерные, так и многомерные DeviceMesh. В многомерной сетке ровно одно измерение сетки должно разбивать измерение классов — это измерение считается измерением «TP» и может находиться в любой позиции (не обязательно последней). Остальные измерения сетки могут быть Shard (например, разбиты по пакету данных для DP/CP) или Replicate.

Параметры:

• input (DTensor) – Входные логиты. Предполагается, что данные разбиты по измерению классов ровно в одном измерении сетки.

• target (Union[torch.Tensor, DTensor]) –

  Должен содержать индексы истинных классов (вероятности классов в настоящее время не поддерживаются). Размещения выводятся из размещений input следующим образом: (1) замена Shard(class_dim) измерения TP на Replicate и (2) сдвиг любого другого Shard(d) с d > class_dim на единицу вниз (так как измерение классов тензора удаляется из результата):

  input tensor shape       input placements             target shape      target placements
  ----------------------   --------------------------   ---------------   ------------------------
  (batch, class)           (Shard(0), Shard(1))         (batch,)          (Shard(0), Replicate())
  (batch, class)           (Replicate(), Shard(1))      (batch,)          (Replicate(), Replicate())
  (batch, class, seq)      (Shard(2), Shard(1))         (batch, seq)      (Shard(1), Replicate())
  (batch, class, seq)      (Shard(0), Shard(1))         (batch, seq)      (Shard(0), Replicate())
  

  Обычный torch.Tensor принимается только тогда, когда все выведенные размещения target являются Replicate (например, в одномерной сетке); в противном случае необходимо явно передать DTensor, чтобы предполагаемое разбиение было однозначным.

• weight (Union[torch.Tensor, DTensor], optional) – Если задан, предполагается, что он реплицирован по DeviceMesh.
• label_smoothing – В настоящее время не поддерживается.

Возвращает:

• reduction="none" — функция потерь для каждого примера; наследует размещения target (Shard в измерениях параллельной обработки пакетов, Replicate в TP).
• reduction="sum" — скаляр; Replicate в измерении TP (пользовательская операция all-reduce обеспечивает глобальную корректность локального значения каждого ранга по TP), Partial("sum") в каждом другом измерении сетки Shard, чтобы сведение между рангами выполнялось отложенно при материализации/перераспределении, и Replicate в измерениях сетки Replicate.
• reduction="mean" — поддерживается только в одномерной сетке; возвращает полностью реплицированный DTensor.

В одномерной сетке все три варианта сводятся к полностью реплицированному DTensor (исходное поведение).

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

Объект DTensor, размещения которого зависят от reduction

Примечание

reduction="mean" поддерживается только в одномерной DeviceMesh. В многомерной сетке деление на каждом ранге на total_weight использует локальное количество, поэтому агрегирование по измерениям, отличным от TP, не даёт правильного глобального среднего; вместо этого используйте reduction="sum" или "none" и при необходимости самостоятельно делите на глобальное количество.

Пример

Здесь вручную создаётся разбитый DTensor, чтобы показать, как его использовать. На практике это обычно результат работы модуля TP.

>>> from torch.distributed.tensor.parallel import loss_parallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> device_mesh = init_device_mesh("cuda", (8,))
>>> input = torch.randn(4, 16, device="cuda", requires_grad=True)
>>> dist_input = distribute_tensor(input, device_mesh, placements=[Shard(1)])
>>> target = torch.randint(16, (4,), device="cuda")
>>> with loss_parallel():
>>>     loss = F.cross_entropy(dist_input, target, reduction="mean")
>>>     loss.backward()
>>> ...

The loss_parallel API is experimental and subject to change.

torch.distributed.tensor.parallel.loss.loss_parallel() [исходный код]

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

Внутри этого менеджера контекста можно как обычно использовать cross_entropy() или CrossEntropyLoss при следующих предположениях относительно входных параметров. Соответствующий вызов backward(), если он есть, также должен выполняться внутри этого менеджера контекста.

Поддерживаются как одномерные, так и многомерные DeviceMesh. В многомерной сетке ровно одно измерение сетки должно разбивать измерение классов — это измерение считается измерением «TP» и может находиться в любой позиции (не обязательно последней). Остальные измерения сетки могут быть Shard (например, разбиты по пакету данных для DP/CP) или Replicate.

Параметры:

• input (DTensor) – Входные логиты. Предполагается, что данные разбиты по измерению классов ровно в одном измерении сетки.

• target (Union[torch.Tensor, DTensor]) –

  Должен содержать индексы истинных классов (вероятности классов в настоящее время не поддерживаются). Размещения выводятся из размещений input следующим образом: (1) замена Shard(class_dim) измерения TP на Replicate и (2) сдвиг любого другого Shard(d) с d > class_dim на единицу вниз (так как измерение классов тензора удаляется из результата):

  input tensor shape       input placements             target shape      target placements
  ----------------------   --------------------------   ---------------   ------------------------
  (batch, class)           (Shard(0), Shard(1))         (batch,)          (Shard(0), Replicate())
  (batch, class)           (Replicate(), Shard(1))      (batch,)          (Replicate(), Replicate())
  (batch, class, seq)      (Shard(2), Shard(1))         (batch, seq)      (Shard(1), Replicate())
  (batch, class, seq)      (Shard(0), Shard(1))         (batch, seq)      (Shard(0), Replicate())
  

  Обычный torch.Tensor принимается только тогда, когда все выведенные размещения target являются Replicate (например, в одномерной сетке); в противном случае необходимо явно передать DTensor, чтобы предполагаемое разбиение было однозначным.

• weight (Union[torch.Tensor, DTensor], optional) – Если задан, предполагается, что он реплицирован по DeviceMesh.
• label_smoothing – В настоящее время не поддерживается.

Возвращает:

• reduction="none" — функция потерь для каждого примера; наследует размещения target (Shard в измерениях параллельной обработки пакетов, Replicate в TP).
• reduction="sum" — скаляр; Replicate в измерении TP (пользовательская операция all-reduce обеспечивает глобальную корректность локального значения каждого ранга по TP), Partial("sum") в каждом другом измерении сетки Shard, чтобы сведение между рангами выполнялось отложенно при материализации/перераспределении, и Replicate в измерениях сетки Replicate.
• reduction="mean" — поддерживается только в одномерной сетке; возвращает полностью реплицированный DTensor.

В одномерной сетке все три варианта сводятся к полностью реплицированному DTensor (исходное поведение).

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

Объект DTensor, размещения которого зависят от reduction

Примечание

reduction="mean" поддерживается только в одномерной DeviceMesh. В многомерной сетке деление на каждом ранге на total_weight использует локальное количество, поэтому агрегирование по измерениям, отличным от TP, не даёт правильного глобального среднего; вместо этого используйте reduction="sum" или "none" и при необходимости самостоятельно делите на глобальное количество.

Пример

Здесь вручную создаётся разбитый DTensor, чтобы показать, как его использовать. На практике это обычно результат работы модуля TP.

>>> from torch.distributed.tensor.parallel import loss_parallel
>>> from torch.distributed.device_mesh import init_device_mesh
>>> ...
>>> device_mesh = init_device_mesh("cuda", (8,))
>>> input = torch.randn(4, 16, device="cuda", requires_grad=True)
>>> dist_input = distribute_tensor(input, device_mesh, placements=[Shard(1)])
>>> target = torch.randint(16, (4,), device="cuda")
>>> with loss_parallel():
>>>     loss = F.cross_entropy(dist_input, target, reduction="mean")
>>>     loss.backward()
>>> ...

torch.distributed.tensor.parallel.input_reshard.input_reshard(module, tp_device_mesh, input_reshard_dim=None) [исходный код]

Регистрирует хуки для nn.Module, обеспечивающие переразбиение входных данных и позволяющие разбивать и восстанавливать их во время обратного прохода.

Регистрирует хуки для nn.Module с переразбиением входных данных, чтобы можно было разбивать данные согласно заданным tp_device_mesh и input_reshard_dim, а затем восстанавливать входные данные при повторном вычислении активаций во время обратного прохода. Это возможно, потому что при параллелизме тензоров (TP) входные данные одинаковы на всех рангах TP.

Параметры:

• module (nn.Module) – Модуль, для которого регистрируется переразбиение входных данных.
• tp_device_mesh (DeviceMesh) – Объект, описывающий топологию сетки устройств для параллелизма тензоров.
• input_reshard_dim (Optional[int]) – Измерение, по которому выполняется разбиение входных данных. Если задано None, входные данные не разбиваются. По умолчанию: None

Возвращает:

Объект nn.Module, для которого зарегистрировано переразбиение входных данных TP.

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

Модуль