torch.distributed.tensor¶

注意

torch.distributed.tensor 目前處於 alpha 狀態並正在開發中，我們致力於為文件中列出的大多數 API 提供向後相容性，但如有必要，可能會進行 API 變更。

PyTorch DTensor (分散式張量)¶

PyTorch DTensor 提供簡單且彈性的張量分片基本功能，能夠透明地處理分散式邏輯，包括分片儲存、運算子計算以及跨裝置/主機的集合通訊。DTensor 可用於建構不同的平行化解決方案，並在處理多維分片時支援分片 state_dict 表示。

請參考基於 DTensor 構建的 PyTorch 原生平行化解決方案的範例。

DTensor 遵循 SPMD (單一程式、多重資料) 程式設計模型，讓使用者能夠編寫分散式程式，如同它是具有相同收斂性的單一裝置程式。它透過指定 DeviceMesh 和 Placement，提供統一的張量分片佈局 (DTensor Layout)。

DeviceMesh 使用 n 維陣列表示裝置拓撲和叢集的通訊器。
Placement 描述邏輯張量在 DeviceMesh 上的分片佈局。DTensor 支援三種類型的 Placement：Shard、Replicate 和 Partial。

DTensor 類別 API¶

DTensor 是 torch.Tensor 的子類別。這意味著一旦創建了 DTensor，就可以使用與 torch.Tensor 非常相似的方式來使用它，包括執行不同類型的 PyTorch 運算子，如同在單個裝置中運行它們一樣，從而允許 PyTorch 運算子的正確分散式計算。

除了現有的 torch.Tensor 方法之外，它還提供了一組額外的方法來與 torch.Tensor 互動、將 DTensor Layout redistribute 重新分配到新的 DTensor、獲取所有裝置上的完整張量內容等等。

class torch.distributed.tensor.DTensor(local_tensor, spec, *, requires_grad)¶

DTensor (分散式張量) 是 torch.Tensor 的子類別，它提供類似於單一裝置的抽象，以便使用多裝置 torch.Tensor 進行程式設計。它透過 DeviceMesh 和以下類型的 Placement 來描述分散式張量分片佈局 (DTensor Layout)。

Shard：張量在 DeviceMesh 維度裝置上的張量維度 dim 上進行分片。
Replicate：張量在 DeviceMesh 維度裝置上進行複製。
Partial：張量正在等待在 DeviceMesh 維度裝置上進行歸約。

在呼叫 PyTorch 運算子時，DTensor 會覆寫 PyTorch 運算子以執行分片計算，並在必要時發出通訊。除了運算子計算之外，DTensor 還會正確地轉換或傳播放置 (DTensor Layout) (基於運算子語義本身)，並產生新的 DTensor 輸出。

為了確保在呼叫 PyTorch 運算子時 DTensor 分片計算的數值正確性，DTensor 要求運算子的每個張量參數都是 DTensor。

注意

在此直接使用 Tensor 子類別建構函式不是創建 DTensor 的建議方法 (即，它無法正確處理 autograd，因此不是公開 API)。請參閱 create_dtensor 章節，瞭解如何創建 DTensor。

回傳類型: DTensor

static from_local(local_tensor, device_mesh=None, placements=None, *, run_check=False, shape=None, stride=None)[原始碼][原始碼]¶

根據指定的 device_mesh 和 placements，從每個 rank 上的本機 torch.Tensor 創建一個 DTensor。

參數

local_tensor (torch.Tensor) – 每個 rank 上的本機 torch.Tensor。
device_mesh (DeviceMesh, optional) – 用於放置 tensor 的 DeviceMesh，如果未指定，則必須在 DeviceMesh 上下文管理器下呼叫，預設值：None
placements (List[Placement], optional) – 描述如何在 DeviceMesh 上放置本地 torch.Tensor 的 placements，必須與 device_mesh.ndim 具有相同數量的元素。

關鍵字參數

run_check (bool, optional) – 以額外的通訊為代價，跨 ranks 執行健全性檢查，以檢查每個本地 tensor 的元資訊，以確保正確性。如果 placements 中有 Replicate，則 device mesh 維度上第一個 rank 的資料將會廣播到其他 ranks。預設值：False
shape (torch.Size, optional) – 一個 int 的 List，用於指定建立在 local_tensor 之上的 DTensor 的大小。請注意，如果 local_tensor 的 shape 在各個 ranks 之間有所不同，則需要提供此參數。如果未提供，則會假定給定的分散式 tensor 在各個 ranks 之間均勻地分片，從而計算出 shape。預設值：None
stride (tuple, optional) – 一個 int 的 List，用於指定 DTensor 的 stride。如果未提供，則會假定給定的分散式 tensor 在各個 ranks 之間均勻地分片，從而計算出 stride。預設值：None

回傳值

一個 DTensor 物件

回傳類型

DTensor

注意

當 run_check=False 時，使用者有責任確保傳入的本地 tensor 在各個 ranks 之間是正確的（也就是說，tensor 對於 Shard(dim) placement 進行分片，或者對於 Replicate() placement 進行複製）。否則，所建立的 DTensor 的行為是未定義的。

注意

from_local 是可微分的，所建立的 DTensor 物件的 requires_grad 將取決於 local_tensor 是否需要 requires_grad。

full_tensor(*, grad_placements=None)[source][source]¶

傳回此 DTensor 的完整 tensor。它將執行必要的 collectives 以從其 DeviceMesh 中的其他 ranks 收集本地 tensors，並將它們串聯在一起。它是以下程式碼的語法糖

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

關鍵字參數: grad_placements (List[Placement], optional) – placements 描述了此函數傳回的完整 Tensor 的任何梯度 layout 的未來 layout。 full_tensor 將 DTensor 轉換為完整的 torch.Tensor，並且傳回的 torch.tensor 可能不會在稍後的程式碼中用作原始複製的 DTensor layout。萬一傳回的 tensor 的梯度 layout 與原始複製的 DTensor layout 不符，此引數是用戶可以提供給 autograd 的提示。如果未指定，我們將假設完整 tensor 的梯度 layout 為複製。
回傳值: 一個 torch.Tensor 物件，代表此 DTensor 的完整 tensor。
回傳類型: Tensor

注意

full_tensor 是可微分的。

redistribute(device_mesh=None, placements=None, *, async_op=False)[source][source]¶

redistribute 執行必要的 collective 運算，以將目前的 DTensor 從其目前的 placements 重新分配到新的 placements，或從其目前的 DeviceMesh 重新分配到新的 DeviceMesh。也就是說，我們可以透過為 DeviceMesh 的每個維度指定一個 Replicate placement，將 Sharded DTensor 轉換為 Replicated DTensor。

當在一個 device mesh 維度上從目前重新分配到新的 placements 時，我們將執行以下運算，包括通訊 collective 或本地運算

Shard(dim) -> Replicate(): all_gather
Shard(src_dim) -> Shard(dst_dim): all_to_all
Replicate() -> Shard(dim): 本地 chunking (也就是 torch.chunk)
Partial() -> Replicate(): all_reduce
Partial() -> Shard(dim): reduce_scatter

redistribute 可以正確地計算出在 1-D 或 N-D DeviceMesh 上建立的 DTensor 的必要重新分配步驟。

參數

device_mesh (DeviceMesh, 選擇性) – 用於放置 DTensor 的 DeviceMesh。如果未指定，則會使用目前 DTensor 的 DeviceMesh。預設值：None
placements (List[Placement], 選擇性) – 描述如何將 DTensor 放置到 DeviceMesh 中的新 placements，必須具有與 device_mesh.ndim 相同數量的元素。預設值：在所有 mesh 維度上複製

關鍵字參數

async_op (bool, 選擇性) – 是否非同步地執行 DTensor redistribute 操作。預設值：False

回傳值

一個 DTensor 物件

回傳類型

DTensor

注意

redistribute 是可微分的，這表示使用者無需擔心 redistribute 操作的反向公式。

注意

redistribute 目前僅支援在相同的 DeviceMesh 上重新分佈 DTensor，如果您需要將 DTensor 重新分佈到不同的 DeviceMesh，請提交 issue。

to_local(*, grad_placements=None)[source][source]¶

取得此 DTensor 在其目前 rank 上的本機 tensor。對於分片，它會傳回邏輯 tensor 視圖的本機 shard；對於複製，它會傳回其目前 rank 上的副本。

關鍵字參數: grad_placements (List[Placement], 選擇性) – 描述從此函式傳回的 Tensor 的任何梯度 layout 的未來版面的 placements。to_local 會將 DTensor 轉換為本機 tensor，且傳回的本機 tensor 可能不會在程式碼中稍後用作原始 DTensor 版面。此引數是使用者可以提供給 autograd 的提示，以防傳回的 tensor 的梯度版面與原始 DTensor 版面不符。如果未指定，我們將假設梯度版面與原始 DTensor 相同，並將其用於梯度計算。
回傳值: 一個 torch.Tensor 或 AsyncCollectiveTensor 物件。它表示其目前 rank 上的本機 tensor。當傳回 AsyncCollectiveTensor 物件時，表示本機 tensor 尚未準備好 (即，通訊尚未完成)。在這種情況下，使用者需要呼叫 wait 以等待本機 tensor 準備好。
回傳類型: Tensor

注意

to_local 是可微分的，傳回的本機 tensor 的 requires_grad 將取決於 DTensor 是否需要 requires_grad。

property device_mesh: DeviceMesh¶: 與此 DTensor 物件關聯的 DeviceMesh 屬性。

注意

device_mesh 是一個唯讀屬性，無法設定。

property placements: Tuple[Placement, ...]¶: 此 DTensor 的 placements 屬性，描述此 DTensor 在其 DeviceMesh 上的版面。

注意

placements 是一個唯讀屬性，無法設定。

DeviceMesh 作為分散式通訊器¶

DeviceMesh 是從 DTensor 建立的抽象概念，用於描述叢集的裝置拓撲並表示多維通訊器（基於 ProcessGroup）。要查看如何建立/使用 DeviceMesh 的詳細資訊，請參閱DeviceMesh 食譜。

DTensor Placement 類型¶

DTensor 支援以下類型的 Placement 在每個 DeviceMesh 維度上

class torch.distributed.tensor.placement_types.Shard(dim)[source][source]¶

Shard(dim) placement 描述在 tensor 維度 dim 上，DTensor 在對應的 DeviceMesh 維度上的分片，其中 DeviceMesh 維度上的每個 rank 僅持有全域 Tensor 的 shard/片段。Shard(dim) placement 遵循 torch.chunk(dim) 語意，其中當 tensor 維度無法在 DeviceMesh 維度上平均分割時，DeviceMesh 維度上的最後幾個 shard 可能為空。Shard placement 可用於所有 DTensor API（即 distribute_tensor、from_local 等）

參數: dim (int) – 描述 DTensor 在其對應的 DeviceMesh 維度上分片的 tensor 維度。

警告

目前，對 tensor 維度進行分片，其中 tensor 維度大小無法在 DeviceMesh 維度上平均分割，這是一個實驗性功能，可能會發生變更。

dim: int¶

class torch.distributed.tensor.placement_types.Replicate[原始碼][原始碼]¶

Replicate() placement 描述了 DTensor 在對應的 DeviceMesh 維度上進行複製，其中 DeviceMesh 維度上的每個 rank 都持有全域 Tensor 的副本。Replicate placement 可以被所有 DTensor APIs 使用 (例如 distribute_tensor、DTensor.from_local 等)。

class torch.distributed.tensor.placement_types.Partial(reduce_op='sum')[原始碼][原始碼]¶

Partial(reduce_op) placement 描述了在指定 DeviceMesh 維度上等待縮減 (reduction) 的 DTensor，其中 DeviceMesh 維度上的每個 rank 都持有全域 Tensor 的部分值。使用者可以使用 redistribute 將 Partial DTensor 重新分配到指定 DeviceMesh 維度上的 Replicate 或 Shard(dim) placement，這將觸發底層必要的通訊操作 (例如 allreduce、reduce_scatter)。

參數: reduce_op (str, optional) – 用於生成 Replicated/Sharded DTensor 的 Partial DTensor 的縮減運算 (reduction op)。僅支援逐元素的縮減運算，包括：「sum」、「avg」、「product」、「max」、「min」，預設值為「sum」。

注意

Partial placement 可以作為 DTensor 運算子的結果生成，並且只能被 DTensor.from_local API 使用。

reduce_op: str = 'sum'¶

class torch.distributed.tensor.placement_types.Placement[原始碼][原始碼]¶

Placement 類型的基底類別，它描述了 DTensor 如何放置到 DeviceMesh 上。Placement 和 DeviceMesh 一起可以描述 DTensor 的 Layout。它是三個主要的 DTensor Placement 類型的基底類別：Shard、Replicate 和 Partial。

這個類別不打算直接使用，主要作為類型提示 (typing stub) 使用。

is_partial()[原始碼][原始碼]¶

回傳類型: bool

is_replicate()[原始碼][原始碼]¶

回傳類型: bool

is_shard(dim=None)[原始碼][原始碼]¶

回傳類型: bool

建立 DTensor 的不同方法¶

有三種方式可以建構 DTensor

distribute_tensor() 從每個 rank 上的邏輯或「全域」torch.Tensor 建立一個 DTensor。這可以用於對 leaf torch.Tensor s (例如模型參數/buffers 和輸入) 進行分片 (shard)。
DTensor.from_local() 從每個 rank 上的本地 torch.Tensor 建立一個 DTensor，可用於從非葉節點 (non-leaf) torch.Tensor（例如，正向/反向傳播期間的中間激活張量）建立 DTensor。
DTensor 提供了專用的張量工廠函數 (例如 empty()、ones()、randn() 等)，可以直接指定 DeviceMesh 和 Placement 來建立不同的 DTensor。與 distribute_tensor() 相比，這可以直接在設備上實例化分片記憶體，而不是在初始化邏輯張量記憶體之後執行分片。

從邏輯 torch.Tensor 建立 DTensor¶

torch.distributed 中的 SPMD（單一程式，多重資料）程式設計模型啟動多個進程（例如，透過 torchrun）來執行相同的程式，這意味著程式內部的模型首先會在不同的進程上初始化（例如，模型可能在 CPU、meta 設備上初始化，或者如果有足夠的記憶體，則直接在 GPU 上初始化）。

DTensor 提供了一個 distribute_tensor() API，可以將模型權重或張量分片到 DTensor 中，它會從每個進程上的“邏輯”張量建立一個 DTensor。這將使創建的 DTensor 能夠符合單一設備語義，這對於**數值正確性**至關重要。

torch.distributed.tensor.distribute_tensor(tensor, device_mesh=None, placements=None)[source]¶

根據指定的 placements 將葉節點 (leaf) torch.Tensor (例如 nn.Parameter/buffers) 分發到 device_mesh。 device_mesh 和 placements 的 rank 必須相同。要分發的 tensor 是邏輯或“全域”張量，並且 API 將使用 DeviceMesh 維度的第一個 rank 中的 tensor 作為真實來源，以保留單一設備語義。如果您想在 Autograd 計算過程中構建一個 DTensor，請改用 DTensor.from_local()。

參數

tensor (torch.Tensor) – 要分發的 torch.Tensor。請注意，如果您想在不能被網格維度中的設備數量均勻分割的維度上分片張量，我們使用 torch.chunk 語義來分片張量並分散分片。這種不均勻的分片行為是實驗性的，可能會發生變化。
device_mesh (DeviceMesh, optional) – 用於分發張量的 DeviceMesh，如果未指定，則必須在 DeviceMesh 上下文管理器下調用，預設值：None
placements (List[Placement], optional) – 描述如何在 DeviceMesh 上放置張量的 placements，必須具有與 device_mesh.ndim 相同數量的元素。如果未指定，我們預設會從 device_mesh 的每個維度的第一個 rank 跨 device_mesh 複製張量。

回傳值

一個 DTensor 或 XLAShardedTensor 物件。

回傳類型

DTensor

注意

當使用 xla device_type 初始化 DeviceMesh 時，distribute_tensor 會返回 XLAShardedTensor。有關更多詳細資訊，請參閱此問題。 XLA 整合是實驗性的，可能會發生變化。

除了 distribute_tensor() 之外，DTensor 還提供了一個 distribute_module() API，可以更輕鬆地在 nn.Module 層級進行分片

torch.distributed.tensor.distribute_module(module, device_mesh=None, partition_fn=None, input_fn=None, output_fn=None)[source]¶

這個函式公開三個函式來控制模組的參數/輸入/輸出

1. 在執行時期執行模組分片之前，透過指定 partition_fn (也就是允許使用者根據指定的 partition_fn，將 Module 參數轉換為 DTensor 參數)。2. 透過指定 input_fn 和 output_fn，在執行時期控制模組的輸入或輸出。(也就是將輸入轉換為 DTensor，將輸出轉換回 torch.Tensor)

參數

module (nn.Module) – 要分區的使用者模組。
device_mesh (DeviceMesh) – 用於放置模組的 device mesh。
partition_fn (Callable) – 用於分割參數的函式 (也就是跨 device_mesh 分享某些參數)。如果未指定 partition_fn，預設情況下，我們會在 mesh 中複製 module 的所有模組參數。
input_fn (Callable) – 指定輸入分佈，也就是可以控制模組的輸入如何分片。input_fn 將安裝為模組 forward_pre_hook (pre forward hook)。
output_fn (Callable) – 指定輸出分佈，也就是可以控制輸出如何分片，或將其轉換回 torch.Tensor。output_fn 將安裝為模組 forward_hook (post forward hook)。

回傳值

一個包含參數/緩衝區的模組，這些參數/緩衝區都是 DTensor s。

回傳類型

模組

注意

當使用 xla device_type 初始化 DeviceMesh 時，distribute_module 會傳回帶有 PyTorch/XLA SPMD 標註參數的 nn.Module。有關更多詳細訊息，請參閱此問題。XLA 整合是實驗性的，可能會發生變更。

DTensor 工廠函式¶

DTensor 也提供專用的 tensor 工廠函式，允許透過額外指定 DeviceMesh 和 Placement，直接使用類似 torch.Tensor 的工廠函式 API (也就是 torch.ones、torch.empty 等) 來建立 DTensor

torch.distributed.tensor.zeros(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None)[source]¶

傳回一個填滿純量值 0 的 DTensor。

參數

size (int...) – 一個整數序列，定義輸出 DTensor 的形狀。可以是可變數量的引數，也可以是像列表或元組這樣的集合。例如：zeros(1,2,3..) 或 zeros([1,2,3..]) 或 zeros((1,2,3..))

關鍵字參數

requires_grad (bool, optional) – 如果 autograd 應該記錄傳回的 DTensor 上的操作。預設值：False。
dtype (torch.dtype, optional) – 傳回的 DTensor 的所需資料類型。預設值：如果 None，則使用全域預設值 (請參閱 torch.set_default_dtype())。
layout (torch.layout, optional) – 傳回的 DTensor 的所需版面配置。預設值：torch.strided。
device_mesh – DeviceMesh 類型，包含 ranks 的 mesh 資訊
placements – 一個 Placement 類型序列：Shard、Replicate

回傳值

每個 rank 上的一個 DTensor 物件

回傳類型

DTensor

torch.distributed.tensor.ones(*size, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None)[source]¶

傳回一個 DTensor，其元素值皆為純量值 1，形狀由變數引數 size 定義。

參數

size (int...) – 一個整數序列，定義輸出 DTensor 的形狀。可以是可變數量的引數或一個集合，如 list 或 tuple。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))

關鍵字參數

dtype (torch.dtype, optional) – 傳回的 DTensor 的所需資料類型。預設值：如果 None，則使用全域預設值 (請參閱 torch.set_default_dtype())。
layout (torch.layout, optional) – 傳回的 DTensor 的所需 layout。預設值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 應該記錄傳回的 DTensor 上的操作。預設值：False。
device_mesh – DeviceMesh 類型，包含 ranks 的 mesh 資訊
placements – 一個 Placement 類型序列：Shard、Replicate

回傳值

每個 rank 上的一個 DTensor 物件

回傳類型

DTensor

torch.distributed.tensor.empty(*size, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None)[source]¶

傳回一個 DTensor，其填充了未初始化的資料。DTensor 的形狀由變數引數 size 定義。

參數

size (int...) – 一個整數序列，定義輸出 DTensor 的形狀。可以是可變數量的引數或一個集合，如 list 或 tuple。例如：empty(1,2,3..) 或 empty([1,2,3..]) 或 empty((1,2,3..))

關鍵字參數

dtype (torch.dtype, optional) – 傳回的 DTensor 的所需資料類型。預設值：如果 None，則使用全域預設值（請參閱 torch.set_default_dtype()）。layout (torch.layout, optional): 傳回的 DTensor 的所需 layout。預設值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 應該記錄傳回的 DTensor 上的操作。預設值：False。
device_mesh – DeviceMesh 類型，包含 ranks 的 mesh 資訊
placements – 一個 Placement 類型序列：Shard、Replicate

回傳值

每個 rank 上的一個 DTensor 物件

回傳類型

DTensor

torch.distributed.tensor.full(size, fill_value, *, dtype=None, layout=torch.strided, requires_grad=False, device_mesh=None, placements=None)[source]¶

根據 device_mesh 和 placements，並以引數 size 定義形狀，傳回一個填滿 fill_value 的 DTensor。

參數

size (int...) – 一個整數序列，定義輸出 DTensor 的形狀。可以是可變數量的引數或一個集合，如 list 或 tuple。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))
fill_value (Scalar) – 用於填充輸出張量的值。

關鍵字參數

dtype (torch.dtype, optional) – 傳回的 DTensor 的所需資料類型。預設值：如果 None，則使用全域預設值 (請參閱 torch.set_default_dtype())。
layout (torch.layout, optional) – 傳回的 DTensor 的所需 layout。預設值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 應該記錄傳回的 DTensor 上的操作。預設值：False。
device_mesh – DeviceMesh 類型，包含 rank 的網格資訊。
placements – 一個 Placement 類型序列：Shard、Replicate

回傳值

每個 rank 上的一個 DTensor 物件

回傳類型

DTensor

torch.distributed.tensor.rand(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None)[source]¶

傳回一個 DTensor，其中填滿了來自區間 [0, 1) 上均勻分佈的隨機數字。張量的形狀由變數引數 size 定義。

參數

size (int...) – 一個整數序列，定義輸出 DTensor 的形狀。可以是可變數量的引數或一個集合，如 list 或 tuple。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))

關鍵字參數

dtype (torch.dtype, optional) – 傳回的 DTensor 的所需資料類型。預設值：如果 None，則使用全域預設值 (請參閱 torch.set_default_dtype())。
layout (torch.layout, optional) – 傳回的 DTensor 的所需 layout。預設值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 應該記錄傳回的 DTensor 上的操作。預設值：False。
device_mesh – DeviceMesh 類型，包含 rank 的網格資訊。
placements – 一個 Placement 類型序列：Shard、Replicate

回傳值

每個 rank 上的一個 DTensor 物件

回傳類型

DTensor

torch.distributed.tensor.randn(*size, requires_grad=False, dtype=None, layout=torch.strided, device_mesh=None, placements=None)[source]¶

傳回一個 DTensor，其中填滿了來自平均值為 0 且變異數為 1 的常態分佈的隨機數字。張量的形狀由變數引數 size 定義。

參數

size (int...) – 一個整數序列，定義輸出 DTensor 的形狀。可以是可變數量的引數或一個集合，如 list 或 tuple。例如：ones(1,2,3..) 或 ones([1,2,3..]) 或 ones((1,2,3..))

關鍵字參數

dtype (torch.dtype, optional) – 傳回的 DTensor 的所需資料類型。預設值：如果 None，則使用全域預設值 (請參閱 torch.set_default_dtype())。
layout (torch.layout, optional) – 傳回的 DTensor 的所需 layout。預設值：torch.strided。
requires_grad (bool, optional) – 如果 autograd 應該記錄傳回的 DTensor 上的操作。預設值：False。
device_mesh – DeviceMesh 類型，包含 rank 的網格資訊。
placements – 一個 Placement 類型序列：Shard、Replicate

回傳值

每個 rank 上的一個 DTensor 物件

回傳類型

DTensor

除錯¶

記錄¶

啟動程式時，您可以使用來自 torch._logging 的 TORCH_LOGS 環境變數開啟額外的記錄功能。

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)[source][source]¶

產生詳細的表格，顯示模組層級上的運算和集合追蹤資訊。資訊量取決於 noise_level

列印模組層級的集合計數
列印未包含在基本運算中的 dTensor 運算、模組資訊
列印未包含在基本運算中的運算
列印所有運算

generate_json_dump(file_name='comm_mode_log.json', noise_level=3)[source][source]¶: 建立用於建構瀏覽器視覺效果的 json 檔案 0. 列印模組層級的集合計數 1. 列印未包含在基本運算中的 dTensor 運算 2. 列印未包含在基本運算中的運算 3. 列印所有運算

get_comm_counts()[source][source]¶

以字典形式傳回通訊計數。

回傳值: 以字典形式呈現的通訊計數。
回傳類型: Dict[Any, int]

get_parameter_info()[source][source]¶

回傳類型: Dict[str, Dict[str, Any]]

get_sharding_info()[source][source]¶

回傳類型: Dict[str, Dict[str, Any]]

get_total_counts()[source][source]¶

回傳類型: int

log_comm_debug_tracing_table_to_file(file_name='comm_mode_log.txt', noise_level=3)[source][source]¶: CommDebugMode 主控台輸出的替代方案，將輸出寫入使用者指定檔案。

為了視覺化維度小於 3 的 DTensor 的分片，DTensor 提供了 visualize_sharding()

torch.distributed.tensor.debug.visualize_sharding(dtensor, header='')[source]¶: 在終端機中視覺化 1D 或 2D 的 DTensor 的分片。

注意

這需要 tabulate 套件。空張量不會印出任何分片資訊

實驗性功能¶

DTensor 也提供了一組實驗性功能。這些功能可能處於原型階段，或者基本功能已經完成，但正在尋找使用者回饋。如果您對這些功能有任何回饋，請提交 Issue 到 PyTorch。

torch.distributed.tensor.experimental.context_parallel(mesh, *, buffers=None, buffer_seq_dims=None, no_restore_buffers=None)[source]¶

context_parallel 是一個實驗性的 API，用於啟用上下文平行處理 (CP)。這個 API 執行兩個動作：1) 使用支援 CP 的版本修補 SDPA (torch.nn.functional.scaled_dot_product_attention)，2) 沿著序列維度對 buffers 進行分片，並且每個 rank 將根據 mesh 保留對應的分片。

參數

mesh (DeviceMesh) – 用於上下文平行處理的裝置網格。
buffers (Optional[List[torch.Tensor]]) – 用法取決於序列維度的 buffers。範例包括輸入批次、標籤和位置嵌入 buffers。這些 buffers 必須沿著序列維度進行分片，以確保準確性。分片將會就地發生，buffer 的形狀將會在上下文中變更。在上下文結束後，buffers 將會被還原。可以使用 no_restore_buffers 來指定哪些 buffers 不需要還原。請注意，buffers 不應包含任何 nn.Parameter。
buffer_seq_dims (Optional[List[int]]) – buffers 的序列維度。
no_restore_buffers (Optional[Set[torch.Tensor]]) – 在此集合中的 buffers 不會在上下文結束後還原。此集合必須是 buffers 的子集。如果在上下文結束後不會使用這些 buffers，則可以將這些 buffers 放入此清單中，以避免額外的還原時間。

回傳類型

Generator[None, None, None]

警告

torch.distributed._tensor.experimental.attention.context_parallel 是 PyTorch 中的原型功能。 API 可能會變更。

torch.distributed.tensor.experimental.local_map(func, out_placements, in_placements=None, device_mesh=None, *, redistribute_inputs=False)[source]¶

local_map() 是一個實驗性的 API，允許使用者將 DTensor 傳遞到一個設計用於在 torch.Tensor 上應用的函式。其運作方式是提取 DTensor 的本地元件，呼叫該函式，並根據 out_placements 將輸出包裝成 DTensor。

參數

func (Callable) – 要應用於 DTensor 的每個本地分片的函式。
out_placements (Union[PlacementType, Tuple[PlacementType, …]]) – func 的扁平化輸出中 DTensor 的所需放置位置。如果扁平化的 output 是一個單一值，則 out_placements 應為 PlacementType 類型。否則，如果扁平化的 output 具有多個值，則 out_placements 應為 PlacementType 值的 tuple，與扁平化的 output 一一對應。此外，對於 Tensor 輸出，我們使用 PlacementType 作為其放置位置（Tuple[Placement] 值）。對於非 Tensor 輸出，PlacementType 應為 None。請注意，唯一的例外是沒有傳入任何 DTensor 參數。在這種情況下，即使 out_placements 不是 None，結果函式也應忽略所需的放置位置，因為該函式不是使用 DTensor 執行的。
in_placements (Tuple[PlacementType, …], optional) – func 的扁平化輸入中 DTensor 的所需放置位置。如果指定了 in_placements，local_map() 將檢查每個 DTensor 參數的放置位置是否與所需的放置位置相同。如果放置位置不同且 redistribute_inputs 為 False，則會引發例外。否則，如果 redistribute_inputs 為 True，則會先將參數重新分佈到所需的分片放置位置，然後再將其本地 tensor 傳遞給 func。唯一的例外是當所需的放置位置不是 None 且參數為 torch.Tensor 時。在這種情況下，將跳過放置位置檢查，並將參數直接傳遞給 func。如果 in_placements 為 None，則不會執行任何放置位置檢查。預設值：None
device_mesh (DeviceMesh, optional) – 所有 DTensor 所在的裝置網格。如果未指定，將從輸入 DTensor 的裝置網格推斷。 local_map 要求每個 DTensor 都位於同一個裝置網格上。預設值：None。
redistribute_inputs (bool, optional) – 指示是否在輸入 DTensor 的放置位置與所需的輸入放置位置不同時，重新分片輸入 DTensor 的布林值。如果此值為 False 且某些 DTensor 輸入具有不同的放置位置，則會引發例外。預設值：False。

回傳值

一個 Callable，將 func 應用於輸入 DTensor 的每個本地分片，並返回一個從 func 的傳回值建構的 DTensor。

引發

AssertionError – 如果輸入 DTensor 未放置在同一個裝置網格上，或者如果它們放置在與傳入的 device_mesh 參數不同的裝置網格上。
AssertionError – 對於任何非 DTensor 輸出，我們要求其在 out_placements 中的相應輸出放置位置為 None。如果不是這種情況，將引發 AssertionError。
ValueError – 如果 redistribute_inputs=False 但根據 in_placements 輸入 DTensor 需要重新分佈。

範例

>>> def mm_allreduce_forward(device_mesh, W, X):
>>>     partial_sum_tensor = torch.mm(W, X)
>>>     reduced_tensor = funcol.all_reduce(partial_sum_tensor, "sum", device_mesh)
>>>     return reduced_tensor
>>>
>>> W = torch.randn(12, 8, requires_grad=False)
>>> X = torch.randn(8, 16, requires_grad=False)
>>> Y = torch.mm(W, X)
>>> row_wise = [Shard(0)]  # row-wise sharding placements on 1-d mesh
>>> col_wise = [Shard(1)]  # col-wise sharding placements on 1-d mesh
>>>
>>> # local_mm_allreduce_forward is the function wrapped with DTensor/Tensor convertion
>>> local_mm_allreduce_forward = local_map(
>>>     mm_allreduce_forward,
>>>     out_placements=[Replicate()],
>>>     in_placements=[col_wise, row_wise],
>>>     device_mesh=device_mesh,
>>> )
>>>
>>> W_dt = distribute_tensor(W, device_mesh, (col_wise))  # col-wisely sharded W tensor
>>> X_dt = distribute_tensor(X, device_mesh, (row_wise))  # row-wisely sharded X tensor
>>> Y_dt = local_mm_allreduce_forward(device_mesh, W_dt, X_dt)  # apply local_mm_allreduce_forward to DTensors

注意

此 API 目前處於實驗階段，可能會有所變更

torch.distributed.tensor.experimental.register_sharding(op)[source]¶

register_sharding() 是一個實驗性的 API，允許使用者在 tensor 輸入和輸出為 DTensor 時，為運算子註冊分片策略。當以下情況時，它可能很有用：（1）op 沒有預設的分片策略，例如當 op 是 DTensor 不支援的自訂運算子時；（2）當使用者想要覆寫現有運算子的預設分片策略時。

參數: op (Union[OpOverload, List[OpOverload]]) – 要註冊自訂分片函式的運算子或運算子列表。
回傳值: 一個函數裝飾器，可用於封裝一個函數，該函數定義了 op 中指定的運算符的分片策略。定義的分片策略將註冊到 DTensor，如果 DTensor 已經實現了該運算符，則會覆蓋預設的分片策略。自定義分片函數接受與原始運算符相同的輸入（除非參數是 torch.Tensor，否則它將被 DTensor 在內部使用的類張量物件取代）。該函數應返回一個 2 元組序列，每個元組指定可接受的輸出放置及其相應的輸入放置。

範例

>>> @register_sharding(aten._softmax.default)
>>> def custom_softmax_sharding(x, dim, half_to_float):
>>>     softmax_dim = dim if dim >= 0 else dim + x.ndim
>>>     acceptable_shardings = []
>>>
>>>     all_replicate = ([Replicate()], [Replicate(), None, None])
>>>     acceptable_shardings.append(all_replicate)
>>>
>>>     for sharding_dim in range(x.ndim):
>>>         if sharding_dim != softmax_dim:
>>>             all_sharded = (
>>>                 [Shard(sharding_dim)],
>>>                 [Shard(sharding_dim), None, None],
>>>             )
>>>             acceptable_shardings.append(all_sharded)
>>>
>>>     return acceptable_shardings

注意

此 API 目前處於實驗階段，可能會有所變更

torch.distributed.tensor¶

PyTorch DTensor (分散式張量)¶

DTensor 類別 API¶

DeviceMesh 作為分散式通訊器¶

DTensor Placement 類型¶

建立 DTensor 的不同方法¶

從邏輯 torch.Tensor 建立 DTensor¶

DTensor 工廠函式¶

除錯¶

記錄¶

除錯工具¶

實驗性功能¶

文件

教學課程

資源