PyTorch/XLA API

torch_xla

torch_xla.device(index: Optional[int] = None) → device

傳回 XLA 裝置的指定執行個體。

若啟用 SPMD，則傳回一個虛擬裝置，該裝置封裝了此程序可用的所有裝置。

參數

index – 要傳回的 XLA 裝置的索引。對應於 torch_xla.devices() 中的索引。

傳回

XLA torch.device。

torch_xla.devices() → List[device]

傳回目前程序中可用的所有裝置。

傳回

XLA torch.devices 的列表。

torch_xla.device_count() → int

傳回目前程序中可定址裝置的數量。

torch_xla.sync(wait: bool = False)

啟動所有擱置中的圖形運算。

參數

wait (bool) – 是否封鎖目前程序直到執行完成。

torch_xla.compile(f: Optional[Callable] = None, full_graph: Optional">[bool] = False, name: Optional">[str] = None, num_different_graphs_allowed: Optional">[int] = None)

使用 torch_xla 的 LazyTensor 追蹤模式，優化給定的模型/函式。PyTorch/XLA 將使用給定的輸入追蹤給定的函式，然後產生圖形來表示此函式內發生的 pytorch 運算。此圖形將由 XLA 編譯，並在加速器上執行（由張量的裝置決定）。對於函式的編譯區域，將停用 Eager 模式。

參數

• model (Callable) – 要優化的模組/函式，如果未傳遞，此函式將充當上下文管理器。

• full_graph (Optional[bool]) – 此編譯是否應產生單一圖形。如果設定為 True 且將產生多個圖形，torch_xla 將拋出包含偵錯資訊的錯誤並退出。

• name (Optional[name]) – 已編譯程式的名稱。如果未指定，將使用函式 f 的名稱。此名稱將用於 PT_XLA_DEBUG 訊息以及 HLO/IR 傾印檔案中。

• num_different_graphs_allowed (Optional[python:int]) – 我們允許給定的模型/函式擁有的不同追蹤圖形的數量。如果超過此限制，將引發錯誤。

範例

# usage 1
@torch_xla.compile()
def foo(x):
  return torch.sin(x) + torch.cos(x)

def foo2(x):
  return torch.sin(x) + torch.cos(x)
# usage 2
compiled_foo2 = torch_xla.compile(foo2)

# usage 3
with torch_xla.compile():
  res = foo2(x)

torch_xla.manual_seed(seed, device=None)

為目前的 XLA 裝置設定產生隨機數的種子。

參數

• seed (python:integer) – 要設定的狀態。

• device (torch.device, optional) – 需要設定 RNG 狀態的裝置。如果遺失，將設定預設裝置種子。

runtime

torch_xla.runtime.device_type() → Optional[str]

傳回目前的 PjRt 裝置類型。

如果尚未設定預設裝置，則選取一個預設裝置

傳回

裝置的字串表示形式。

torch_xla.runtime.local_process_count() → int

傳回在此主機上執行的程序數量。

torch_xla.runtime.local_device_count() → int

傳回此主機上的裝置總數。

假設每個程序都具有相同數量的可定址裝置。

torch_xla.runtime.addressable_device_count() → int

傳回此程序可見的裝置數量。

torch_xla.runtime.global_device_count() → int

傳回所有程序/主機上的裝置總數。

torch_xla.runtime.global_runtime_device_count() → int

傳回所有程序/主機上的執行階段裝置總數，對於 SPMD 特別有用。

torch_xla.runtime.world_size() → int

傳回參與作業的程序總數。

torch_xla.runtime.global_ordinal() → int

傳回此執行緒在所有程序中的全域序數。

全域序數的範圍為 [0, global_device_count)。不保證全域序數與 TPU 工作站 ID 有任何可預測的關係，也不保證它們在每個主機上是連續的。

torch_xla.runtime.local_ordinal() → int

傳回此執行緒在此主機中的本機序數。

本機序數的範圍為 [0, local_device_count)。

torch_xla.runtime.get_master_ip() → str

擷取執行階段的主工作站 IP。這會呼叫後端特定的探索 API。

傳回

主工作站的 IP 位址，以字串形式表示。

torch_xla.runtime.use_spmd(auto: Optional[bool] = False)

啟用 SPMD 模式的 API。這是啟用 SPMD 的建議方式。

如果某些張量已在非 SPMD 裝置上初始化，這會強制 SPMD 模式。這表示這些張量將在裝置之間複製。

參數

auto (bool) – 是否啟用自動分片。如需更多詳細資訊，請參閱 https://github.com/pytorch/xla/blob/master/docs/spmd_advanced.md#auto-sharding

torch_xla.runtime.is_spmd()

傳回是否為執行設定 SPMD。

torch_xla.runtime.initialize_cache(path: str, readonly: bool = False)

初始化持久編譯快取。此 API 必須在執行任何運算之前呼叫。

參數

• path (str) – 儲存持久快取的路徑。

• readonly (bool) – 此工作站是否應具有快取的寫入權限。

xla_model

torch_xla.core.xla_model.xla_device(n: Optional[int] = None, devkind: Optional[str] = None) → device

傳回 XLA 裝置的指定執行個體。

參數

• n (python:int, optional) – 要傳回的特定執行個體（序數）。如果指定，將傳回特定的 XLA 裝置執行個體。否則，將傳回 devkind 的第一個裝置。

• devkind (string..., optional) – 如果指定，裝置類型，例如 TPU、CUDA、CPU 或自訂 PJRT 裝置。已棄用。

傳回

具有請求執行個體的 torch.device。

torch_xla.core.xla_model.xla_device_hw(device: Union[str, device]) → str

傳回給定裝置的硬體類型。

參數

device (string or torch.device) – 將對應到真實裝置的 xla 裝置。

傳回

給定裝置的硬體類型的字串表示形式。

torch_xla.core.xla_model.is_master_ordinal(local: bool = True) → bool

檢查目前程序是否為主序數 (0)。

參數

local (bool) – 應檢查本機還是全域主序數。在多主機複製的情況下，只有一個全域主序數（主機 0，裝置 0），而本機主序數的數量為 NUM_HOSTS。預設值：True

傳回

布林值，指示目前程序是否為主序數。

torch_xla.core.xla_model.all_reduce(reduce_type: str, inputs: Union[Tensor, List[Tensor]], scale: float = 1.0, groups: Optional[List[List[int]]] = None, pin_layout: bool = True) → Union[Tensor, List[Tensor]]

對輸入張量執行就地縮減運算。

參數

• reduce_type (string) – xm.REDUCE_SUM、xm.REDUCE_MUL、xm.REDUCE_AND、xm.REDUCE_OR、xm.REDUCE_MIN 和 xm.REDUCE_MAX 之一。

• inputs – 要對其執行 all reduce 運算的單一 torch.Tensor 或 torch.Tensor 列表。

• scale (python:float) – 要在縮減後套用的預設縮放值。預設值：1.0

• groups (list, optional) –

  列表的列表，表示 all_reduce() 運算的複本群組。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定義兩個群組，一個具有 [0, 1, 2, 3] 複本，另一個具有 [4, 5, 6, 7] 複本。如果為 None，則只會有一個群組，其中包含所有複本。

• pin_layout (bool, optional) – 是否為此通訊運算釘選版面配置。版面配置釘選可以防止參與通訊的每個程序具有稍微不同的程式時可能發生的資料損毀，但可能會導致某些 xla 編譯失敗。當您看到類似「HloModule has a mix of layout constrained」的錯誤訊息時，請取消釘選版面配置。

傳回

如果傳遞單一 torch.Tensor，則傳回值是保存縮減值（跨複本）的 torch.Tensor。如果傳遞列表/tuple，此函式會對輸入張量執行就地 all-reduce 運算，並傳回列表/tuple 本身。

torch_xla.core.xla_model.all_gather(value: Tensor, dim: int = 0, groups: Optional[List[List[int]]] = None, output: Optional[Tensor] = None, pin_layout: bool = True) → Tensor

沿著給定的維度執行 all-gather 操作。

參數

• value (torch.Tensor) – 輸入張量。

• dim (python:int) – 收集維度。預設值：0

• groups (list, optional) –

  列表的列表，代表 all_gather() 操作的副本群組。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定義兩個群組，一個具有 [0, 1, 2, 3] 複本，另一個具有 [4, 5, 6, 7] 複本。如果為 None，則只會有一個群組，其中包含所有複本。

• output (torch.Tensor) – 選擇性輸出張量。

• pin_layout (bool, optional) – 是否為此通訊運算釘選版面配置。版面配置釘選可以防止參與通訊的每個程序具有稍微不同的程式時可能發生的資料損毀，但可能會導致某些 xla 編譯失敗。當您看到類似「HloModule has a mix of layout constrained」的錯誤訊息時，請取消釘選版面配置。

傳回

一個張量，在 dim 維度中，包含來自參與副本的所有值。

torch_xla.core.xla_model.all_to_all(value: Tensor, split_dimension: int, concat_dimension: int, split_count: int, groups: Optional[List[List[int]]] = None, pin_layout: bool = True) → Tensor

在輸入張量上執行 XLA AllToAll() 操作。

參見：https://tensorflow.dev.org.tw/xla/operation_semantics#alltoall

參數

• value (torch.Tensor) – 輸入張量。

• split_dimension (python:int) – 應該進行分割的維度。

• concat_dimension (python:int) – 應該進行串聯的維度。

• split_count (python:int) – 分割計數。

• groups (list, optional) –

  列表的列表，表示 all_reduce() 運算的複本群組。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定義兩個群組，一個具有 [0, 1, 2, 3] 複本，另一個具有 [4, 5, 6, 7] 複本。如果為 None，則只會有一個群組，其中包含所有複本。

• pin_layout (bool, optional) – 是否為此通訊運算釘選版面配置。版面配置釘選可以防止參與通訊的每個程序具有稍微不同的程式時可能發生的資料損毀，但可能會導致某些 xla 編譯失敗。當您看到類似「HloModule has a mix of layout constrained」的錯誤訊息時，請取消釘選版面配置。

傳回

all_to_all() 操作的結果 torch.Tensor。

torch_xla.core.xla_model.add_step_closure(closure: Callable[[...], Any], args: Tuple[Any, ...] = (), run_async: bool = False)

將一個閉包函數添加到步驟結束時要執行的函數列表中。

在模型訓練期間，很多時候需要列印/報告資訊（列印到控制台、發佈到 TensorBoard 等），這需要檢查中間張量的內容。在模型程式碼的不同點檢查不同張量的內容需要多次執行，並且通常會導致效能問題。添加步驟閉包函數將確保它在屏障之後運行，屆時所有即時張量都已經實體化為設備數據。即時張量將包括閉包函數參數捕獲的張量。因此，即使佇列了多個需要檢查多個張量的閉包函數，使用 add_step_closure() 也將確保只執行一次。步驟閉包函數將按照它們被加入佇列的順序依序執行。請注意，即使使用此 API 會最佳化執行，也建議每 N 個步驟限制列印/報告事件的次數。

參數

• closure (callable) – 要呼叫的函數。

• args (tuple) – 要傳遞給閉包函數的引數。

• run_async – 如果為 True，則非同步執行閉包函數。

torch_xla.core.xla_model.wait_device_ops(devices: List[str] = [])

等待給定裝置上的所有非同步操作完成。

參數

devices (string..., optional) – 需要等待其非同步操作完成的裝置。如果為空，則將等待所有本機裝置。

torch_xla.core.xla_model.optimizer_step(optimizer: Optimizer, barrier: bool = False, optimizer_args: Dict = {}, groups: Optional[List[List[int]]] = None, pin_layout: bool = True)

執行提供的最佳化器步驟，並跨所有裝置同步梯度。

參數

• optimizer (torch.Optimizer) – 需要呼叫其 step() 函數的 torch.Optimizer 執行個體。step() 函數將使用 optimizer_args 具名引數呼叫。

• barrier (bool, optional) – 是否應在此 API 中發出 XLA 張量屏障。如果使用 PyTorch XLA ParallelLoader 或 DataParallel 支援，則不需要，因為屏障將由 XLA 資料載入器迭代器 next() 呼叫發出。預設值：False

• optimizer_args (dict, optional) – optimizer.step() 呼叫的具名引數字典。

• groups (list, optional) –

  列表的列表，表示 all_reduce() 運算的複本群組。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

  定義兩個群組，一個具有 [0, 1, 2, 3] 複本，另一個具有 [4, 5, 6, 7] 複本。如果為 None，則只會有一個群組，其中包含所有複本。

• pin_layout (bool, optional) – 是否在縮減梯度時固定版面配置。有關詳細資訊，請參見 xm.all_reduce。

傳回

optimizer.step() 呼叫傳回的相同值。

範例

>>> import torch_xla.core.xla_model as xm
>>> xm.optimizer_step(self.optimizer)

torch_xla.core.xla_model.save(data: Any, file_or_path: Union[str, TextIO], master_only: bool = True, global_master: bool = False)

將輸入資料儲存到檔案中。

儲存的資料在儲存之前會傳輸到 PyTorch CPU 裝置，因此後續的 torch.load() 將載入 CPU 資料。使用檢視時必須小心。建議您在載入張量並將其移動到目標裝置後再重新建立檢視，而不是儲存檢視。

參數

• data – 要儲存的輸入資料。任何 Python 物件（列表、元組、集合、字典等）的巢狀組合。

• file_or_path – 資料儲存操作的目的地。可以是檔案路徑或 Python 檔案物件。如果 master_only 為 False，則路徑或檔案物件必須指向不同的目的地，否則來自同一主機的所有寫入都會互相覆寫。

• master_only (bool, optional) – 是否只有主裝置應儲存資料。如果為 False，則 file_or_path 引數對於參與複寫的每個序數都應是不同的檔案或路徑，否則同一主機上的所有副本都將寫入到相同的位置。預設值：True

• global_master (bool, optional) – 當 master_only 為 True 時，此旗標控制是否每個主機的主裝置（如果 global_master 為 False）儲存內容，或僅全域主裝置（序數 0）。預設值：False

範例

>>> import torch_xla.core.xla_model as xm
>>> xm.wait_device_ops() # wait for all pending operations to finish.
>>> xm.save(obj_to_save, path_to_save)
>>> xm.rendezvous('torch_xla.core.xla_model.save') # multi process context only

torch_xla.core.xla_model.rendezvous(tag: str, payload: bytes = b'', replicas: List[int] = []) → List[bytes]

等待所有網格用戶端到達具名的 rendezvous。

注意：PJRT 不支援 XRT 網格伺服器，因此這實際上是 xla_rendezvous 的別名。

參數

• tag (string) – 要加入的 rendezvous 的名稱。

• payload (bytes, optional) – 要傳送到 rendezvous 的酬載。

• replicas (list, python:int) – 參與 rendezvous 的副本序數。空列表表示網格中的所有副本。預設值：[]

傳回

所有其他核心交換的酬載，核心序數 i 的酬載位於傳回元組中的位置 i。

範例

>>> import torch_xla.core.xla_model as xm
>>> xm.rendezvous('example')

torch_xla.core.xla_model.mesh_reduce(tag: str, data, reduce_fn: Callable[[...], Any]) → Union[Any, ToXlaTensorArena]

執行圖形外的用戶端網格縮減。

參數

• tag (string) – 要加入的 rendezvous 的名稱。

• data – 要縮減的資料。reduce_fn 可呼叫物件將接收一個列表，其中包含來自所有網格用戶端程序（每個核心一個）的相同資料副本。

• reduce_fn (callable) – 一個函數，它接收 data 類物件的列表，並傳回縮減後的結果。

傳回

縮減後的值。

範例

>>> import torch_xla.core.xla_model as xm
>>> import numpy as np
>>> accuracy = xm.mesh_reduce('test_accuracy', accuracy, np.mean)

torch_xla.core.xla_model.set_rng_state(seed: int, device: Optional[str] = None)

設定隨機數字產生器狀態。

參數

• seed (python:integer) – 要設定的狀態。

• device (string, optional) – 需要設定 RNG 狀態的裝置。如果遺失，將設定預設裝置種子。

torch_xla.core.xla_model.get_rng_state(device: Optional[str] = None) → int

取得目前正在執行的隨機數字產生器狀態。

參數

device (string, optional) – 需要擷取 RNG 狀態的裝置。如果遺失，將設定預設裝置種子。

傳回

RNG 狀態，為整數。

torch_xla.core.xla_model.get_memory_info(device: Optional[device] = None) → MemoryInfo

擷取裝置記憶體用量。

參數

• device – Optional[torch.device] 要求其記憶體資訊的裝置。

• device. (如果未傳遞，將使用預設值) –

傳回

MemoryInfo 字典，其中包含給定裝置的記憶體用量。

範例

>>> xm.get_memory_info()
{'bytes_used': 290816, 'bytes_limit': 34088157184, 'peak_bytes_used': 500816}

torch_xla.core.xla_model.get_stablehlo(tensors: Optional[List[Tensor]] = None) → str

以字串格式取得計算圖形的 StableHLO。

如果 tensors 不是空的，將傾印以 tensors 作為輸出的圖形。如果 tensors 是空的，將傾印整個計算圖形。

對於推論圖形，建議將模型輸出傳遞到 tensors。對於訓練圖形，直接識別「輸出」並不容易。建議使用空的 tensors。

若要在 StableHLO 中啟用原始程式碼行資訊，請設定環境變數 XLA_HLO_DEBUG=1。

參數

tensors (list[torch.Tensor], optional) – 代表 StableHLO 圖形輸出/根目錄的張量。

傳回

字串格式的 StableHLO 模組。

torch_xla.core.xla_model.get_stablehlo_bytecode(tensors: Optional[Tensor] = None) → bytes

以位元組碼格式取得計算圖形的 StableHLO。

如果 tensors 不是空的，將傾印以 tensors 作為輸出的圖形。如果 tensors 是空的，將傾印整個計算圖形。

對於推論圖形，建議將模型輸出傳遞到 tensors。對於訓練圖形，直接識別「輸出」並不容易。建議使用空的 tensors。

參數

tensors (list[torch.Tensor], optional) – 代表 StableHLO 圖形輸出/根目錄的張量。

傳回

位元組碼格式的 StableHLO 模組。

distributed

class torch_xla.distributed.parallel_loader.MpDeviceLoader(loader, device, **kwargs)

使用背景資料上傳包裝現有的 PyTorch DataLoader。

此類別應僅用於多處理資料平行處理。它將使用 ParallelLoader 包裝傳入的資料載入器，並傳回目前裝置的 per_device_loader。

參數

• loader (torch.utils.data.DataLoader) – 要包裝的 PyTorch DataLoader。

• device (torch.device…) – 資料必須傳送到的裝置。

• kwargs – ParallelLoader 建構函式的具名引數。

範例

>>> device = torch_xla.device()
>>> train_device_loader = MpDeviceLoader(train_loader, device)

torch_xla.distributed.xla_multiprocessing.spawn(fn, args=(), nprocs=None, join=True, daemon=False, start_method='spawn')

啟用基於多處理的複寫。

參數

• fn (callable) – 要為參與複寫的每個裝置呼叫的函數。將使用複寫中程序的全域索引作為第一個引數呼叫該函數，後跟在 args 中傳遞的引數。

• args (tuple) – fn 的引數。預設值：空元組

• nprocs (python:int) – 複寫的程序/裝置數量。目前，如果指定，則可以是 1 或 None（將自動轉換為裝置的最大數量）。其他數字將導致 ValueError。

• join (bool) – 呼叫是否應封鎖，等待已產生的程序完成。預設值：True

• daemon (bool) – 產生的程序是否應設定 daemon 旗標（請參見 Python 多處理 API）。預設值：False

• start_method (string) – Python multiprocessing 程序建立方法。預設值：spawn

傳回

與 torch.multiprocessing.spawn API 傳回的物件相同的物件。如果 nprocs 為 1，將直接呼叫 fn 函數，並且 API 將傳回 None。

spmd

torch_xla.distributed.spmd.mark_sharding(t: Union[Tensor, XLAShardedTensor], mesh: Mesh, partition_spec: Tuple[Optional[Union[Tuple, int, str]], ...]) → XLAShardedTensor

使用 XLA 分割規格註解提供的張量。在內部，它將對應的 XLATensor 註解為針對 XLA SpmdPartitioner 傳遞進行分割。

參數

• t (Union[torch.Tensor, XLAShardedTensor]) – 要使用 partition_spec 註解的輸入張量。

• mesh (Mesh) – 描述邏輯 XLA 裝置拓撲和底層裝置 ID。

• partition_spec (Tuple[Tuple, python:int, str, None]) – 裝置網格維度索引或 None 的元組。每個索引都是一個整數、字串（如果網格軸已命名），或整數或字串的元組。這指定每個輸入排名如何分割（索引到 mesh_shape）或複寫 (None)。當指定元組時，對應的輸入張量軸將沿元組中的所有邏輯軸分割。請注意，元組中指定的網格軸順序會影響產生的分割。

• dynamo_custom_op (bool) – 如果設定為 True，它會呼叫 mark_sharding 的 dynamo 自訂操作變體，使其可被 dynamo 識別和追蹤。

範例

>>> import torch_xla.runtime as xr
>>> import torch_xla.distributed.spmd as xs
>>> mesh_shape = (4, 2)
>>> num_devices = xr.global_runtime_device_count()
>>> device_ids = np.array(range(num_devices))
>>> mesh = Mesh(device_ids, mesh_shape, ('x', 'y'))
>>> input = torch.randn(8, 32).to(xm.xla_device())
>>> xs.mark_sharding(input, mesh, (0, None)) # 4-way data parallel
>>> linear = nn.Linear(32, 10).to(xm.xla_device())
>>> xs.mark_sharding(linear.weight, mesh, (None, 1)) # 2-way model parallel

torch_xla.distributed.spmd.clear_sharding(t: Union[Tensor, XLAShardedTensor]) → Tensor

從輸入張量清除分割註解，並傳回 cpu 轉換後的張量。這是一個就地操作，但也會傳回相同的 torch.Tensor。

參數

t (Union[torch.Tensor, XLAShardedTensor]) – 我們要清除分割的張量

傳回

沒有分割的張量。

傳回類型

t (torch.Tensor)

範例

>>> import torch_xla.distributed.spmd as xs
>>> torch_xla.runtime.use_spmd()
>>> t1 = torch.randn(8,8).to(torch_xla.device())
>>> mesh = xs.get_1d_mesh()
>>> xs.mark_sharding(t1, mesh, (0, None))
>>> xs.clear_sharding(t1)

torch_xla.distributed.spmd.set_global_mesh(mesh: Mesh)

設定可供目前程序使用的全域網格。

參數

mesh – (Mesh) 將作為全域網格的 Mesh 物件。

範例

>>> import torch_xla.distributed.spmd as xs
>>> mesh = xs.get_1d_mesh("data")
>>> xs.set_global_mesh(mesh)

torch_xla.distributed.spmd.get_global_mesh() → Optional[Mesh]

取得目前程序的全域網格。

傳回

(Optional[Mesh]) 如果已設定全域網格，則為 Mesh 物件，否則傳回 None。

傳回類型

mesh

範例

>>> import torch_xla.distributed.spmd as xs
>>> xs.get_global_mesh()

torch_xla.distributed.spmd.get_1d_mesh(axis_name: Optional[str] = None) → Mesh

輔助函數，傳回在一個維度中包含所有裝置的網格。

參數

axis_name – (Optional[str]) 代表網格軸名稱的選擇性字串

傳回

Mesh 物件

傳回類型

Mesh

範例

>>> # This example is assuming 1 TPU v4-8
>>> import torch_xla.distributed.spmd as xs
>>> mesh = xs.get_1d_mesh("data")
>>> print(mesh.mesh_shape)
(4,)
>>> print(mesh.axis_names)
('data',)

class torch_xla.distributed.spmd.Mesh(device_ids: Union[ndarray, List], mesh_shape: Tuple[int, ...], axis_names: Optional[Tuple[str, ...]] = None)

描述邏輯 XLA 裝置拓撲網格和底層資源。

參數

• device_ids (Union[np.ndarray, List]) – 自訂順序的裝置 (ID) 平面列表。此列表會被重塑為 mesh_shape 陣列，並使用類似 C 語言的索引順序填入元素。

• mesh_shape (Tuple[python:int, ...]) – 一個整數 tuple，描述裝置網格的邏輯拓撲形狀，且每個元素描述對應軸中的裝置數量。

• axis_names (Tuple[str, ...]) – 要指派給 devices 引數維度的資源軸名稱序列。其長度應與 devices 的階數相符。

範例

>>> mesh_shape = (4, 2)
>>> num_devices = len(xm.get_xla_supported_devices())
>>> device_ids = np.array(range(num_devices))
>>> mesh = Mesh(device_ids, mesh_shape, ('x', 'y'))
>>> mesh.get_logical_mesh()
>>> array([[0, 1],
          [2, 3],
          [4, 5],
          [6, 7]])
>>> mesh.shape()
OrderedDict([('x', 4), ('y', 2)])

class torch_xla.distributed.spmd.HybridMesh(*, ici_mesh_shape: Tuple[int, ...], dcn_mesh_shape: Optional[Tuple[int, ...]] = None, axis_names: Optional[Tuple[str, ...]] = None)

建立透過 ICI 和 DCN 網路連線的裝置混合網格。

邏輯網格的形狀應依網路密集度遞增排序，例如 [replica, data, model]，其中 mdl 具有最多的網路通訊需求。

參數

• ici_mesh_shape – 內部連線裝置的邏輯網格形狀。

• dcn_mesh_shape – 外部連線裝置的邏輯網格形狀。

範例

>>> # This example is assuming 2 slices of v4-8.
>>> ici_mesh_shape = (1, 4, 1) # (data, fsdp, tensor)
>>> dcn_mesh_shape = (2, 1, 1)
>>> mesh = HybridMesh(ici_mesh_shape, dcn_mesh_shape, ('data','fsdp','tensor'))
>>> print(mesh.shape())
>>> >> OrderedDict([('data', 2), ('fsdp', 4), ('tensor', 1)])

experimental

torch_xla.experimental.eager_mode(enable: bool)

設定 torch_xla 的預設執行模式。

在 eager 模式下，只有以 `torch_xla.compile` 編譯的函式會被追蹤和編譯。其他 torch 運算會以 eager 方式執行。

debug

torch_xla.debug.metrics.metrics_report()

擷取包含完整指標和計數器報告的字串。

torch_xla.debug.metrics.short_metrics_report(counter_names: Optional[list] = None, metric_names: Optional[list] = None)

擷取包含完整指標和計數器報告的字串。

參數

• counter_names (list) – 需要列印資料的計數器名稱清單。

• metric_names (list) – 需要列印資料的指標名稱清單。

torch_xla.debug.metrics.counter_names()

擷取所有目前作用中的計數器名稱。

torch_xla.debug.metrics.counter_value(name)

傳回作用中計數器的值。

參數

name (string) – 需要擷取值的計數器名稱。

傳回

計數器值為整數。

torch_xla.debug.metrics.metric_names()

擷取所有目前作用中的指標名稱。

torch_xla.debug.metrics.metric_data(name)

傳回作用中指標的資料。

參數

name (string) – 需要擷取資料的指標名稱。

傳回

指標資料，為 (TOTAL_SAMPLES、ACCUMULATOR、SAMPLES) 的 tuple。TOTAL_SAMPLES 是已發佈到指標的樣本總數。指標僅保留給定數量的樣本 (在循環緩衝區中)。ACCUMULATOR 是 TOTAL_SAMPLES 中樣本的總和。SAMPLES 是 (TIME、VALUE) tuple 的清單。