torch.distributed.tensor

API класса DTensor

DTensor является подклассом torch.Tensor. Это означает, что созданный DTensor можно использовать почти так же, как torch.Tensor, в том числе запускать различные операторы PyTorch так, будто они выполняются на одном устройстве, обеспечивая корректные распределённые вычисления для операторов PyTorch.

Помимо существующих методов torch.Tensor, класс также предоставляет набор дополнительных методов для взаимодействия с torch.Tensor, redistribute макета DTensor в новый DTensor, получения полного содержимого тензора на всех устройствах и т. д.

class torch.distributed.tensor.DTensor(**kwargs)

DTensor (распределённый тензор) — это подкласс torch.Tensor, обеспечивающий абстракцию, аналогичную работе с одним устройством, для программирования с использованием torch.Tensor на нескольких устройствах. Он описывает макет распределённого разбиения тензора на части (макет DTensor) с помощью DeviceMesh и следующих типов Placement:

• Shard: тензор разбивается на части по измерению тензора dim на устройствах измерения DeviceMesh
• Replicate: тензор реплицируется на устройствах измерения DeviceMesh
• Partial: на устройствах измерения DeviceMesh для тензора ожидается редукция

При вызове операторов PyTorch DTensor переопределяет их для выполнения распределённых вычислений и при необходимости инициирует коммуникации. Вместе с вычислением оператора DTensor будет корректно преобразовывать или распространять размещения (макет DTensor) в соответствии с семантикой самого оператора и создавать новые выходные данные DTensor.

Чтобы обеспечить численную корректность распределённых вычислений DTensor с разбиением на части при вызове операторов PyTorch, DTensor требует, чтобы каждый аргумент оператора типа Tensor был DTensor.

Примечание

Прямое использование конструктора подкласса Tensor здесь не является рекомендуемым способом создания DTensor (то есть он некорректно обрабатывает autograd и поэтому не является публичным API). Инструкции по созданию DTensor см. в разделе create_dtensor.

__create_chunk_list__() [исходный код]

Возвращает список ChunkStorageMetadata — объекта dataclass, описывающего размер и смещение локальной части или реплики на текущем ранге. У DTensor на каждом ранге есть одна локальная часть или реплика, поэтому возвращаемый список обычно содержит только один элемент.

Этот dunder-метод используется в основном для распределённых контрольных точек.

Возвращает:

Объект List[ChunkStorageMetadata], представляющий размер и смещение части на текущем ранге.

static from_local(local_tensor, device_mesh=None, placements=None, *, run_check=False, shape=None, stride=None, grad_placements=None) [исходный код]

Создаёт DTensor из локального torch.Tensor на каждом ранге в соответствии с заданными device_mesh и placements.

Параметры:

• local_tensor (torch.Tensor) – локальный torch.Tensor на каждом ранге.
• device_mesh (DeviceMesh, необязательный) – DeviceMesh, в который помещается тензор. Если не задан, вызов должен выполняться в контекстном менеджере DeviceMesh. Значение по умолчанию: None
• placements (List[Placement], необязательный) – размещения, описывающие, как поместить локальный torch.Tensor в DeviceMesh; количество элементов должно совпадать с device_mesh.ndim.

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

• run_check (bool, необязательный) – выполнить проверку согласованности метаданных каждого локального тензора между рангами, чтобы подтвердить корректность; это потребует дополнительных коммуникаций. Если в placements имеется Replicate, данные с первого ранга измерения сетки устройств будут переданы остальным рангам. Значение по умолчанию: False
• shape (torch.Size, необязательный) – список целых чисел, задающий размер DTensor, построенного на основе local_tensor. Укажите его, если размеры local_tensor различаются на разных рангах. Если значение не задано, shape вычисляется исходя из предположения, что заданный распределённый тензор равномерно разбит на части между рангами. Значение по умолчанию: None
• stride (tuple, необязательный) – список целых чисел, задающий шаг DTensor. Если значение не задано, stride вычисляется исходя из предположения, что заданный распределённый тензор равномерно разбит на части между рангами. Значение по умолчанию: None
• grad_placements (List[Placement], необязательный) – задаёт ожидаемые размещения входного градиента. Перед выходом из DTensor входной градиент (обычный тензор) будет перераспределён в соответствии с этим размещением. Если значение не задано, используются приведённые ниже гарантии размещения по умолчанию. Значение по умолчанию: None

Возвращает:

Объект DTensor

Вызывает исключения:

ValueError – Если placements содержит разные типы редукции Partial (например, одновременно Partial("sum") и Partial("max")). Для всех размещений Partial должна использоваться одна и та же операция редукции.

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

DTensor

Примечание

При вызове run_check=False пользователь должен самостоятельно убедиться, что переданный локальный тензор корректен для всех рангов (то есть тензор разбит на части для размещения Shard(dim) или реплицирован для размещения Replicate()). В противном случае поведение созданного DTensor не определено.

Примечание

from_local является дифференцируемым; requires_grad созданного объекта DTensor зависит от того, установлен ли для local_tensor флаг requires_grad.

Примечание

При обратном проходе from_local обеспечивает следующие гарантии размещения градиента. Для каждого измерения сетки размещение градиента определяется следующим образом:

Размещение при прямом проходе	Размещение градиента
Shard	Shard
Replicate	Replicate
Partial	Replicate

Если размещение при прямом проходе — Partial, мы всегда перераспределяем градиент в Replicate, а не оставляем его в Partial. Это может быть не самым эффективным вариантом, но позволяет избежать неоднозначности и обеспечивает более ясную семантику градиентов для пользователей.

full_tensor(*, grad_placements=None) [исходный код]

Возвращает полный тензор этого DTensor. Для этого выполняются необходимые коллективные операции: сбор локальных тензоров с других рангов в его DeviceMesh и их объединение. Это синтаксический сахар для следующего кода:

dtensor.redistribute(placements=[Replicate()] * mesh.ndim).to_local()

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

grad_placements (List[Placement], необязательный) – размещения задают будущий макет градиента для полного тензора, возвращаемого этой функцией. full_tensor преобразует DTensor в полный torch.Tensor, который впоследствии может использоваться в коде не в исходном реплицированном макете DTensor. Этот аргумент служит подсказкой для autograd на случай, если макет градиента возвращённого тензора не совпадает с исходным реплицированным макетом DTensor. Если значение не задано, предполагается, что макет градиента полного тензора будет реплицированным.

Возвращает:

Объект torch.Tensor, представляющий полный тензор этого DTensor.

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

Tensor

Примечание

full_tensor является дифференцируемым.

redistribute(device_mesh=None, placements=None, *, async_op=False, forward_dtype=None, backward_dtype=None) [исходный код]

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

При перераспределении в новые размещения по одному измерению сетки устройств выполняются следующие операции, включающие коллективные коммуникации или локальные действия:

1. Shard(dim) -> Replicate(): all_gather
2. Shard(src_dim) -> Shard(dst_dim): all_to_all
3. Replicate() -> Shard(dim): локальное разбиение на части (то есть torch.chunk)
4. Partial() -> Replicate(): all_reduce
5. Partial() -> Shard(dim): reduce_scatter
6. Shard(dim) -> Partial("sum"): локальное заполнение нулями. Это преобразование, выполняемое только по явному запросу, создаёт на каждом ранге локальный тензор глобального размера и не требует коммуникаций.

Из преобразований в Partial этот публичный API поддерживает только Shard(dim) -> Partial("sum"). Каждый ранг размещает свою локальную часть в тензоре нулей логической формы по соответствующему смещению; таким образом, суммирование по рангам восстанавливает исходный логический тензор.

redistribute корректно определяет необходимые этапы перераспределения для DTensor, созданных как на одномерной, так и на N-мерной сетке DeviceMesh.

Параметры:

• device_mesh (DeviceMesh, необязательный) – DeviceMesh, в который помещается DTensor. Если значение не задано, используется DeviceMesh текущего DTensor. Значение по умолчанию: None
• placements (List[Placement], необязательный) – новые размещения, описывающие, как поместить DTensor в DeviceMesh; количество элементов должно совпадать с device_mesh.ndim. По умолчанию: репликация по всем измерениям сетки

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

• async_op (bool, необязательный) – выполнять ли операцию перераспределения DTensor асинхронно. Значение по умолчанию: False
• forward_dtype (torch.dtype, необязательный) – преобразовать локальный тензор к типу forward_dtype перед выполнением коллективной операции прямого прохода. Результирующий DTensor будет иметь тип forward_dtype. Если задано None, преобразование не выполняется. Значение по умолчанию: None
• backward_dtype (torch.dtype, необязательный) – преобразовать градиент к типу backward_dtype перед выполнением коллективной операции обратного прохода. После операции градиент всегда преобразуется обратно к типу исходного DTensor. Если задано None, коллективная операция обратного прохода выполняется с типом исходного DTensor (а не forward_dtype). Значение по умолчанию: None

Возвращает:

Объект DTensor

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

DTensor

Примечание

redistribute допускает вычисление производных второго порядка: пользователю не нужно беспокоиться о формуле обратного прохода для операции перераспределения или её совместимости с autograd при вычислении производных второго порядка. Дифференцирование более высокого порядка не тестировалось (но может работать).

Примечание

redistribute в настоящее время поддерживает перераспределение DTensor только в пределах одной и той же DeviceMesh. Если вам нужно перераспределять DTensor в другую DeviceMesh, создайте issue.

to_local(*, grad_placements=None) [исходный код]

Возвращает локальный тензор этого DTensor на текущем ранге. При разбиении на части возвращается локальная часть представления логического тензора, а при репликации — реплика на текущем ранге.

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

grad_placements (List[Placement], необязательный) – размещения задают будущий макет градиента для тензора, возвращаемого этой функцией. to_local преобразует DTensor в локальный тензор, который впоследствии может использоваться в коде не в исходном макете DTensor. Этот аргумент служит подсказкой для autograd на случай, если макет градиента возвращённого тензора не совпадает с исходным макетом DTensor. Если значение не задано, предполагается, что макет градиента останется таким же, как у исходного DTensor, и будет использоваться для вычисления градиента.

Возвращает:

Объект torch.Tensor или AsyncCollectiveTensor, представляющий локальный тензор на текущем ранге. Если возвращён объект AsyncCollectiveTensor, локальный тензор ещё не готов (то есть коммуникация не завершена). В этом случае необходимо вызвать wait и дождаться готовности локального тензора.

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

Tensor

Примечание

to_local является дифференцируемым; requires_grad возвращённого локального тензора зависит от того, установлен ли для DTensor флаг requires_grad.

property device_mesh: DeviceMesh

Атрибут DeviceMesh, связанный с этим объектом DTensor.

Примечание

device_mesh — свойство только для чтения, его нельзя задать.

property placements: tuple[Placement, ...]

Атрибут placements этого DTensor, описывающий его макет в соответствующей DeviceMesh.

Примечание

placements — свойство только для чтения, его нельзя задать.

DeviceMesh как средство распределённой коммуникации

DeviceMesh была создана на основе DTensor как абстракция для описания топологии устройств кластера и представления многомерных средств коммуникации (поверх ProcessGroup). Подробные инструкции по созданию и использованию DeviceMesh см. в руководстве по DeviceMesh.

Типы размещения DTensor

DTensor поддерживает следующие типы Placement для каждого измерения DeviceMesh:

class torch.distributed.tensor.placement_types.Shard [источник]

Размещение Shard(dim) описывает разбиение DTensor по измерению тензора dim относительно соответствующего измерения DeviceMesh, при котором каждый ранг в измерении DeviceMesh содержит только часть/фрагмент глобального тензора. Размещение Shard(dim) следует семантике torch.chunk(dim): последние фрагменты в измерении DeviceMesh могут быть пустыми, если размер измерения тензора не делится нацело на размер измерения DeviceMesh. Размещение Shard можно использовать во всех API DTensor (например, distribute_tensor, from_local и т. д.).

Параметры:

dim (int) – Измерение тензора, по которому DTensor разбивается относительно соответствующего измерения DeviceMesh.

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

Разбиение по измерению тензора, размер которого не делится нацело на размер измерения DeviceMesh, в настоящее время является экспериментальным и может измениться.

Примечание

При проверке, похоже ли размещение на shard, используйте _is_shard_like() вместо isinstance(p, Shard), чтобы также учитывать _StridedShard.

static local_shard_size_and_offset(curr_local_size, num_chunks, rank) [источник]

По размеру текущего локального тензора (который уже может быть разбит по некоторым измерениям) вычисляет размер и смещение нового локального фрагмента с учетом заданного числа частей (num_chunks обычно равно размеру текущего измерения разбиения).

Примечание: смещение нового локального фрагмента отсчитывается относительно текущего разбитого тензора, а не глобального. См. _utils.compute_local_shape_and_global_offset, чтобы вычислить глобальное смещение.

Возвращает (размер нового локального фрагмента, смещение)

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

tuple[_RankTypeT, _RankTypeT]

class torch.distributed.tensor.placement_types._StridedShard(dim, *, split_factor=1, sf=None) [источник]

_StridedShard изначально был добавлен для поддержки разбиения 2D FSDP2 + TP, при котором тензор сначала разбивается по измерению TP-сетки, а затем по измерению FSDP-сетки. Мы называем это разбиением справа налево, противоположным разбиению слева направо, используемому по умолчанию. См. пример ниже:

tensor shape: [8, 8]
mesh: [[0, 1], [2, 3]], names=("dp", "tp")
placements: [Shard(0), Shard(0)]

При разбиении по умолчанию тензор сначала разбивается по измерению сетки «dp», а затем по измерению «tp». Результат разбиения будет следующим:

Rank    |   Mesh Coordinate |   Shard Index
------------------------------------------------
0       |   (0, 0)          |   0 (row 0-1)
1       |   (0, 1)          |   1 (row 2-3)
2       |   (1, 0)          |   2 (row 4-5)
3       |   (1, 1)          |   3 (row 6-7)

При разбиении FSDP2 + TP порядок обратный: тензор сначала разбивается по измерению сетки «tp», а затем по измерению «dp». Такое разбиение справа налево даст следующий результат:

Rank    |   Mesh Coordinate |   Shard Index
------------------------------------------------
0       |   (0, 0)          |   0 (row 0-1)
1       |   (0, 1)          |   2 (row 4-5)
2       |   (1, 0)          |   1 (row 2-3)
3       |   (1, 1)          |   3 (row 6-7)

В результате любая попытка перераспределить этот DTensor в полную реплику приведет к неправильному результату, поскольку перераспределение из фрагментов в реплику всегда выполняется справа налево, независимо от того, использовалось ли разбиение слева направо или справа налево. Чтобы решить эту проблему, мы используем размещение _StridedShard, которое позволяет согласовать разбиение справа налево с соглашением слева направо как при распределении тензора, так и при его перераспределении.

С помощью _StridedShard разбиение справа налево, показанное выше, можно представить так:

tensor shape: [8, 8]
mesh: [[0, 1], [2, 3]], names=("dp", "tp")
placements: [_StridedShard(0, split_factor=2), Shard(0)]

Обработка placements слева направо даст тот же результат, отличный от результата при использовании размещения Shard:

Rank    |   Mesh Coordinate |   Shard Index
------------------------------------------------
0       |   (0, 0)          |   0 (row 0-1)
1       |   (0, 1)          |   2 (row 4-5)
2       |   (1, 0)          |   1 (row 2-3)
3       |   (1, 1)          |   3 (row 6-7)

Аргумент split_factor — это число существующих фрагментов по измерению разбиения тензора до обработки размещения _StridedShard, как если бы разбиение выполнялось справа налево. В приведенном выше примере тензор сначала нужно разбить на 2 фрагмента по измерению «tp», а затем разбить по измерению «dp». Поэтому split_factor размещения _StridedShard по измерению «dp» равно 2.

local_shard_size_and_offset(curr_local_size, num_chunks, rank, offset_mode=_StridedShardOffsetMode.FIRST) [источник]

Вычисляет размер и смещение(я) локального фрагмента для размещения _StridedShard.

В отличие от обычного размещения Shard, при котором смещения непрерывны, _StridedShard создает разреженные (страйдовые) смещения из-за семантики разбиения справа налево. Этот метод вычисляет фактические индексы, принадлежащие локальному фрагменту.

Параметры:

• self (_StridedShard) – Экземпляр размещения _StridedShard.
• curr_local_size (int) – Текущий размер измерения тензора, который нужно разбить.
• num_chunks (int) – Число частей, на которые нужно разделить измерение (обычно равно размеру измерения сетки).
• rank (RankType) – Индекс ранга, для которого вычисляется фрагмент.
• offset_mode (_StridedShardOffsetMode) – Управляет материализацией смещений. FIRST возвращает только первое смещение как int, ALL возвращает все смещения как список, а NONE пропускает материализацию смещений. Значение по умолчанию — FIRST.

Возвращает:

Кортеж, содержащий:

• local_shard_size (int): Число элементов в локальном фрагменте для этого ранга.
• offset (int | list[int] | None): Определяется параметром offset_mode. FIRST возвращает первое смещение как int, ALL возвращает список всех смещений, а NONE возвращает None.

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

tuple

property split_factor: int | SymInt

class torch.distributed.tensor.placement_types.Replicate [источник]

Размещение Replicate() описывает репликацию DTensor по соответствующему измерению DeviceMesh: каждый ранг в измерении DeviceMesh содержит реплику глобального тензора. Размещение Replicate можно использовать во всех API DTensor (например, distribute_tensor, DTensor.from_local и т. д.).

class torch.distributed.tensor.placement_types.Partial [источник]

ALL_REDUCE_OPS: tuple[str, ...] = ('sum', 'avg', 'min', 'max', 'product')

Размещение Partial(reduce_op) описывает DTensor, для которого ожидается редукция по указанному измерению DeviceMesh: каждый ранг в измерении DeviceMesh содержит частичное значение глобального тензора. Пользователь может перераспределить DTensor Partial в размещение Replicate или Shard(dim) по указанному измерению DeviceMesh с помощью redistribute; при этом необходимые операции обмена данными будут выполнены автоматически (например, allreduce, reduce_scatter).

Параметры:

reduce_op (str, optional) –

Операция редукции, используемая для преобразования частичного DTensor в реплицированный/разделенный DTensor. Соответствует операциям редукции, поддерживаемым torch.distributed.ReduceOp. Значение по умолчанию: «sum».

Поддерживаемые значения:

• "sum": Поэлементное суммирование по всем рангам.
• "avg": Поэлементное усреднение по всем рангам.
• "min": Поэлементный минимум по всем рангам.
• "max": Поэлементный максимум по всем рангам.
• "product": Поэлементное произведение по всем рангам.
• "band": Побитовое И по всем рангам (только для целочисленных тензоров).
• "bor": Побитовое ИЛИ по всем рангам (только для целочисленных тензоров).
• "bxor": Побитовое исключающее ИЛИ по всем рангам (только для целочисленных тензоров).

Примечание

Размещение Partial может быть результатом операций DTensor, вызвано DTensor.from_local или получено при перераспределении размещения Shard в Partial("sum").

LINEAR_REDUCE_OPS: tuple[str, ...] = ('sum', 'avg')

class torch.distributed.tensor.placement_types._MaskPartial(reduce_op=None, mask_buffer=None, offset_shape=None, offset_dim=0, *args, **kwargs) [источник]

Размещение частичной маски, предназначенное для операции встраивания с разбиением по строкам, где необходимо замаскировать индексы и скорректировать их для локального фрагмента встраивания. Маскирование встраивания — это особый тип размещения Partial.

ПРИМЕЧАНИЕ: жизненный цикл этого размещения _MaskPartial соответствует жизненному циклу связанного DTensor, то есть indices_mask существует только в течение времени жизни DTensor.

mask_buffer: MaskBuffer

offset_dim: int = 0

offset_shape: Size | None = None

class torch.distributed.tensor.placement_types.Placement

Базовый класс типа Placement, описывающий размещение DTensor в DeviceMesh. Placement и DeviceMesh вместе могут описывать компоновку DTensor. Это базовый класс для трех основных типов размещения DTensor: Shard, Replicate и Partial.

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

is_partial(self: torch._C._distributed.Placement, reduce_op: str | None = None) → bool

is_replicate(self: torch._C._distributed.Placement) → bool

is_shard(self: torch._C._distributed.Placement, dim: SupportsInt | SupportsIndex | None = None) → bool

Создание DTensor из логического torch.Tensor

Модель программирования SPMD (одна программа, множество данных) в torch.distributed запускает несколько процессов (например, с помощью torchrun) для выполнения одной и той же программы. Это означает, что модель внутри программы сначала инициализируется в разных процессах (например, модель может быть инициализирована на CPU, на устройстве meta или непосредственно на GPU, если памяти достаточно).

DTensor предоставляет API distribute_tensor(), который может разбивать веса модели или тензоры на DTensor, создавая DTensor из «логического» тензора в каждом процессе. Благодаря этому созданные DTensor соответствуют семантике одного устройства, что критически важно для численной корректности.

torch.distributed.tensor.distribute_tensor(tensor, device_mesh=None, placements=None, *, src_data_rank=0) [источник]

Распределяет конечный torch.Tensor (например, nn.Parameter/буферы) по device_mesh в соответствии с заданным placements. Ранг device_mesh и placements должен совпадать. Распределяемый tensor является логическим, или «глобальным», тензором; API использует tensor первого ранга измерения DeviceMesh как источник достоверных данных, чтобы сохранить семантику одного устройства. Если нужно создать DTensor в ходе вычислений Autograd, используйте вместо этого DTensor.from_local().

Параметры:

• tensor (torch.Tensor) – torch.Tensor для распределения. Если требуется разбить тензор по измерению, размер которого не делится нацело на число устройств в соответствующем измерении сетки, для разбиения и рассылки фрагментов используется семантика torch.chunk. Поведение при неравномерном разбиении является экспериментальным и может измениться.
• device_mesh (DeviceMesh, optional) – DeviceMesh для распределения тензора. Если параметр не задан, вызов должен выполняться в контекстном менеджере DeviceMesh. Значение по умолчанию: None.
• placements (List[Placement], optional) – размещения, описывающие размещение тензора в DeviceMesh; число элементов должно совпадать с device_mesh.ndim. Если параметр не задан, тензор по умолчанию реплицируется между device_mesh, начиная с первого ранга каждого измерения device_mesh.

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

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

Возвращает:

Объект DTensor или XLAShardedTensor.

Вызывает исключения:

ValueError – Если placements содержит разные типы редукции Partial (например, одновременно Partial("sum") и Partial("max")). Для всех размещений Partial должна использоваться одна и та же операция редукции.

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

DTensor

Примечание

При инициализации DeviceMesh с типом устройства xla функция distribute_tensor вместо этого возвращает XLAShardedTensor. Подробнее см. это обсуждение. Интеграция с XLA является экспериментальной и может измениться.

Помимо distribute_tensor(), DTensor также предоставляет API distribute_module(), упрощающий разбиение на уровне nn.Module.

torch.distributed.tensor.distribute_module(module, device_mesh=None, partition_fn=None, input_fn=None, output_fn=None) [источник]

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

Поведение трех функций обратного вызова:

1. partition_fn разбивает параметры/буферы модуля до выполнения (то есть преобразует обычные параметры torch.Tensor в параметры DTensor). Все параметры/буферы, не преобразованные функцией partition_fn, после этого реплицируются между device_mesh. Если partition_fn равен None, реплицируются все параметры/буферы.
2. input_fn / output_fn преобразуют входы/выходы модуля во время выполнения (например, преобразуют обычные входные тензоры в DTensor или преобразуют выходы DTensor обратно в обычные тензоры torch.Tensor).

Параметры:

• module (nn.Module) – пользовательский модуль для разбиения.
• device_mesh (DeviceMesh) – сетка устройств для размещения модуля. Если None, используется текущая сетка из окружающего контекста сетки.
• partition_fn (Callable) – функция с сигнатурой partition_fn(name: str, module: nn.Module, device_mesh: DeviceMesh) -> None. Она вызывается для каждого подмодуля (включая сам module), где name — полное квалифицированное имя подмодуля ("" для модуля верхнего уровня), возвращаемое функцией nn.Module.named_modules(). Функция должна изменять module на месте, обычно вызывая distribute_tensor() и register_parameter() для замены выбранных параметров на параметры DTensor. Возвращаемое значение игнорируется. Если параметр не задан, все параметры модуля реплицируются между device_mesh.
• input_fn (Callable) – функция с сигнатурой input_fn(module: nn.Module, inputs: Inputs, device_mesh: DeviceMesh) -> Inputs, где Inputs — кортеж позиционных аргументов для forward (по одному элементу на каждый позиционный аргумент, например (x,) для forward(self, x); именованные аргументы в input_fn не передаются). Устанавливается на module верхнего уровня как forward_pre_hook и возвращает новый кортеж аргументов той же структуры, что и inputs, обычно с преобразованием каждой записи torch.Tensor в DTensor. Этот кортеж заменяет inputs при фактическом вызове forward.
• output_fn (Callable) – функция с сигнатурой output_fn(module: nn.Module, outputs: Outputs, device_mesh: DeviceMesh) -> Outputs, где Outputs — тип возвращаемого значения module.forward (один тензор или любая вложенная структура тензоров, например кортеж/словарь); обычно эти тензоры являются объектами DTensor. Устанавливается на module верхнего уровня как forward_hook и возвращает заменяющий результат той же структуры — обычно с преобразованием каждого DTensor обратно в обычный torch.Tensor с помощью DTensor.to_local() или DTensor.full_tensor() либо с перераспределением в другие размещения.

Возвращает:

Модуль, все параметры/буферы которого являются объектами DTensor (тот же объект module, измененный на месте).

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

Module

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

Ранее input_fn и output_fn принимали два аргумента (inputs, device_mesh) / (outputs, device_mesh). Эта форма устарела; используйте описанную выше форму с тремя аргументами.

Пример:

import torch.nn as nn
from torch.distributed.tensor import (
    DTensor,
    Replicate,
    Shard,
    distribute_module,
    distribute_tensor,
    init_device_mesh,
)


class MyModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.fc = nn.Linear(8, 8)

    def forward(self, x):
        return self.fc(x)


mesh = init_device_mesh("cuda", (4,))


def shard_params(name, mod, mesh):
    if isinstance(mod, nn.Linear):
        for pname, param in mod.named_parameters(recurse=False):
            mod.register_parameter(
                pname, nn.Parameter(distribute_tensor(param, mesh, [Shard(0)]))
            )


# Replicate the input across the mesh before forward.
def to_replicated(mod, inputs, mesh):
    return tuple(
        distribute_tensor(x, mesh, [Replicate()])
        if isinstance(x, torch.Tensor)
        else x
        for x in inputs
    )


# Convert the DTensor output back to a plain torch.Tensor after forward.
def to_local(mod, outputs, mesh):
    return outputs.full_tensor() if isinstance(outputs, DTensor) else outputs


sharded = distribute_module(
    MyModule(),
    mesh,
    partition_fn=shard_params,
    input_fn=to_replicated,
    output_fn=to_local,
)

Примечание

При инициализации DeviceMesh с типом устройства xla функция distribute_module возвращает nn.Module с параметрами, аннотированными для PyTorch/XLA SPMD. Подробнее см. это обсуждение. Интеграция с XLA является экспериментальной и может измениться.

Функции-фабрики DTensor

DTensor также предоставляет специальные функции-фабрики тензоров, позволяющие создавать DTensor напрямую с помощью API функций-фабрик, аналогичных torch.Tensor (например, torch.ones, torch.empty и т. д.), дополнительно указывая DeviceMesh и Placement для создаваемого DTensor:

torch.distributed.tensor.zeros(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor, заполненный скалярным значением 0.

Параметры:

size (int...) – последовательность целых чисел, задающая форму выходного DTensor. Можно передать переменное число аргументов или коллекцию, например список или кортеж. Например: zeros(1,2,3..) или zeros([1,2,3..]) или zeros((1,2,3..))

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

• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого DTensor. По умолчанию: если None, используется глобальное значение по умолчанию (см. torch.set_default_dtype()).
• layout (torch.layout, необязательно) – требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

torch.distributed.tensor.ones(*size, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor, заполненный скалярным значением 1, с формой, заданной переменным аргументом size.

Параметры:

size (int...) – последовательность целых чисел, задающая форму выходного DTensor. Можно передать переменное число аргументов или коллекцию, например список или кортеж. Например: ones(1,2,3..) или ones([1,2,3..]) или ones((1,2,3..))

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

• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого DTensor. По умолчанию: если None, используется глобальное значение по умолчанию (см. torch.set_default_dtype()).
• layout (torch.layout, необязательно) – требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

torch.distributed.tensor.empty(*size, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor, заполненный неинициализированными данными. Форма DTensor задаётся переменным аргументом size.

Параметры:

size (int...) – последовательность целых чисел, задающая форму выходного DTensor. Можно передать переменное число аргументов или коллекцию, например список или кортеж. Например: empty(1,2,3..) или empty([1,2,3..]) или empty((1,2,3..))

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

• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого DTensor. По умолчанию: если None, используется глобальное значение по умолчанию (см. torch.set_default_dtype()). layout (torch.layout, необязательно): требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

torch.distributed.tensor.full(size, fill_value, *, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor, заполненный значением fill_value в соответствии с device_mesh и placements, с формой, заданной аргументом size.

Параметры:

• size (int...) – последовательность целых чисел, задающая форму выходного DTensor. Можно передать переменное число аргументов или коллекцию, например список или кортеж. Например: ones(1,2,3..) или ones([1,2,3..]) или ones((1,2,3..))
• fill_value (Scalar) – значение, которым заполняется выходной тензор.

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

• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого DTensor. По умолчанию: если None, используется глобальное значение по умолчанию (см. torch.set_default_dtype()).
• layout (torch.layout, необязательно) – требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов.
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

torch.distributed.tensor.linspace(start, end, steps, *, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor размера steps, значения которого равномерно распределены от start до end.

Параметры:

• start (float или DTensor) – начальное значение набора точек. Если это DTensor, он должен быть нульмерным
• end (float или DTensor) – конечное значение набора точек. Если это DTensor, он должен быть нульмерным
• steps (int) – размер создаваемого DTensor

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

• dtype (torch.dtype, необязательно) – тип данных, используемый для вычислений. По умолчанию: если None, для вычислений используется глобальный тип данных по умолчанию (см. torch.set_default_dtype()), если start и end — вещественные числа; если одно из них комплексное, используется соответствующий комплексный тип данных.
• layout (torch.layout, необязательно) – требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

torch.distributed.tensor.logspace(start, end, steps, base=10.0, *, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor размера steps, значения которого равномерно распределены от base ^start до base ^end включительно в логарифмическом масштабе с основанием base.

Параметры:

• start (float или DTensor) – начальное значение набора точек. Если это DTensor, он должен быть нульмерным
• end (float или DTensor) – конечное значение набора точек. Если это DTensor, он должен быть нульмерным
• steps (int) – размер создаваемого DTensor

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

• base (float, необязательно) – основание логарифма. По умолчанию: 10.0.
• dtype (torch.dtype, необязательно) – тип данных, используемый для вычислений. По умолчанию: если None, для вычислений используется глобальный тип данных по умолчанию (см. torch.set_default_dtype()), если start и end — вещественные числа; если одно из них комплексное, используется соответствующий комплексный тип данных.
• layout (torch.layout, необязательно) – требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

torch.distributed.tensor.rand(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor, заполненный случайными числами из равномерного распределения на интервале [0, 1). Форма тензора задаётся переменным аргументом size.

Параметры:

size (int...) – последовательность целых чисел, задающая форму выходного DTensor. Можно передать переменное число аргументов или коллекцию, например список или кортеж. Например: ones(1,2,3..) или ones([1,2,3..]) или ones((1,2,3..))

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

• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого DTensor. По умолчанию: если None, используется глобальное значение по умолчанию (см. torch.set_default_dtype()).
• layout (torch.layout, необязательно) – требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов.
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

torch.distributed.tensor.randn(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None) [исходный код]

Возвращает DTensor, заполненный случайными числами из нормального распределения со средним 0 и дисперсией 1. Форма тензора задаётся переменным аргументом size.

Параметры:

size (int...) – последовательность целых чисел, задающая форму выходного DTensor. Можно передать переменное число аргументов или коллекцию, например список или кортеж. Например: ones(1,2,3..) или ones([1,2,3..]) или ones((1,2,3..))

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

• dtype (torch.dtype, необязательно) – требуемый тип данных возвращаемого DTensor. По умолчанию: если None, используется глобальное значение по умолчанию (см. torch.set_default_dtype()).
• layout (torch.layout, необязательно) – требуемый формат возвращаемого DTensor. По умолчанию: torch.strided.
• requires_grad (bool, необязательно) – следует ли autograd записывать операции над возвращаемым DTensor. По умолчанию: False.
• device_mesh – тип DeviceMesh, содержит информацию о сетке рангов.
• placements – последовательность типа Placement: Shard, Replicate

Возвращает:

Объект DTensor на каждом ранге

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

DTensor

Случайные операции

DTensor предоставляет функциональность распределённого генератора случайных чисел (RNG), обеспечивающую уникальные значения для случайных операций над секционированными тензорами и одинаковые значения для случайных операций над реплицированными тензорами. Эта система требует, чтобы все участвующие ранги (например, ранги SPMD) начинали с одинакового состояния генератора перед выполнением каждой случайной операции с DTensor. При соблюдении этого условия после завершения каждой такой операции все они оказываются в одном и том же состоянии. Во время случайных операций обмен данными для синхронизации состояний RNG не выполняется.

Операторы, принимающие именованный аргумент generator, используют переданный пользователем генератор, если он задан; в противном случае используется генератор устройства по умолчанию. Используемый генератор продвигается после операции с DTensor. Допускается использовать один и тот же генератор как для операций с DTensor, так и для операций без DTensor, однако в этом случае необходимо следить, чтобы операции без DTensor одинаково продвигали состояние генератора на всех рангах.

При использовании DTensor совместно с конвейерным параллелизмом ранги каждого этапа конвейера должны использовать разные начальные значения, а ранги внутри одного этапа конвейера — одинаковые.

Инфраструктура RNG в DTensor основана на алгоритме RNG на основе Philox и поддерживает любые серверные части на основе Philox (CUDA и другие устройства, подобные CUDA), но, к сожалению, пока не поддерживает серверную часть CPU.

Ведение журналов

При запуске программы можно включить дополнительное ведение журналов с помощью переменной окружения TORCH_LOGS из torch._logging :

• TORCH_LOGS=+dtensor будет отображать сообщения logging.DEBUG и все уровни выше.
• TORCH_LOGS=dtensor будет отображать сообщения logging.INFO и выше.
• TORCH_LOGS=-dtensor будет отображать сообщения logging.WARNING и выше.

Средства отладки

Для отладки программы, применившей DTensor, и более подробного изучения коллективных операций, выполняемых под капотом, DTensor предоставляет CommDebugMode:

class torch.distributed.tensor.debug.CommDebugMode

CommDebugMode — это менеджер контекста, который подсчитывает количество функциональных коллективных операций в своем контексте. Для этого используется TorchDispatchMode.

Примечание

Пока поддерживаются не все коллективные операции.

Пример использования

mod = ...
comm_mode = CommDebugMode()
with comm_mode:
    mod.sum().backward()
print(comm_mode.get_comm_counts())

generate_comm_debug_tracing_table(noise_level=3) [исходный код]

Создает подробную таблицу с информацией об операциях и трассировке коллективных операций на уровне модуля. Объем информации зависит от noise_level

0. выводит количество коллективных операций на уровне модуля
1. выводит операции dTensor, не относящиеся к тривиальным операциям, и информацию о модуле
2. выводит операции, не относящиеся к тривиальным операциям
3. выводит все операции

generate_json_dump(file_name='comm_mode_log.json', noise_level=3) [исходный код]

Создает JSON-файл для построения визуализации в браузере: 0. выводит количество коллективных операций на уровне модуля; 1. выводит операции dTensor, не относящиеся к тривиальным операциям; 2. выводит операции, не относящиеся к тривиальным операциям; 3. выводит все операции

get_comm_counts() [исходный код]

Возвращает количество операций обмена данными в виде словаря.

Возвращает:

Количество операций обмена данными в виде словаря.

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

Dict[Any, int]

get_parameter_info() [исходный код]

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

dict[str, dict[str, Any]]

get_sharding_info() [исходный код]

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

dict[str, dict[str, Any]]

get_total_counts() [исходный код]

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

int

log_comm_debug_tracing_table_to_file(file_name='comm_mode_log.txt', noise_level=3) [исходный код]

Альтернатива выводу CommDebugMode в консоль: записывает данные в указанный пользователем файл

Для визуализации разбиения DTensor размерности менее 3 DTensor предоставляет visualize_sharding():

torch.distributed.tensor.debug.visualize_sharding(dtensor, header='', use_rich=False) [исходный код]

Визуализирует разбиение в терминале для DTensor размерности 1D или 2D.

Примечание

Для этого необходим пакет tabulate либо rich и matplotlib. Для пустых тензоров информация о разбиении не выводится.

Случай 1: ошибка в пользовательском коде

Чаще всего эта ошибка возникает, когда создается обычный Tensor (с помощью фабричной функции), а затем выполняется операция с Tensor и DTensor, например:

tensor = torch.arange(10)
return tensor + dtensor

Операции со смешанными Tensor и DTensor запрещены: если входными данными какой-либо операции (например, torch.add) является DTensor, все входные данные типа Tensor также должны быть DTensor. Это связано с неоднозначностью семантики. Неизвестно, одинаков ли tensor на всех рангах или различается, поэтому мы предлагаем пользователю самостоятельно определить, как создать DTensor с корректными размещениями из tensor.

Если tensor действительно одинаков на каждом ранге, создайте реплицированный DTensor:

tensor = torch.arange(10)
tensor = DTensor.from_local(tensor, placements=(Replicate(),))
return tensor + dtensor

Ниже показано, как создать DTensor с фрагментами. Семантически это означает, что данные Tensor разделены между фрагментами, а операции выполняются над «полными объединенными данными».

tensor = torch.full([], RANK)
tensor = DTensor.from_local(tensor, placements=(Shard(0),))
return tensor + dtensor

Возможно, вы захотите выполнить с тензором и другие действия (это не единственные два варианта!).