XLA 設備上的 PyTorch¶

PyTorch 使用 torch_xla 套件在 XLA 設備（例如 TPU）上執行。本文件說明如何在這些設備上執行模型。

建立 XLA 張量¶

PyTorch/XLA 為 PyTorch 新增了一個新的 xla 設備類型。此設備類型的工作方式與其他 PyTorch 設備類型相同。例如，以下是如何建立和列印 XLA 張量

import torch
import torch_xla
import torch_xla.core.xla_model as xm

t = torch.randn(2, 2, device=xm.xla_device())
print(t.device)
print(t)

這段程式碼看起來應該很熟悉。 PyTorch/XLA 使用與常規 PyTorch 相同的介面，並新增了一些功能。匯入 torch_xla 會初始化 PyTorch/XLA，而 xm.xla_device() 會傳回目前的 XLA 設備。根據您的環境，這可能是 CPU 或 TPU。

XLA 張量是 PyTorch 張量¶

PyTorch 操作可以像 CPU 或 CUDA 張量一樣在 XLA 張量上執行。

例如，可以將 XLA 張量加在一起

t0 = torch.randn(2, 2, device=xm.xla_device())
t1 = torch.randn(2, 2, device=xm.xla_device())
print(t0 + t1)

或矩陣相乘

print(t0.mm(t1))

或與神經網路模組一起使用

l_in = torch.randn(10, device=xm.xla_device())
linear = torch.nn.Linear(10, 20).to(xm.xla_device())
l_out = linear(l_in)
print(l_out)

與其他設備類型一樣，XLA 張量僅適用於相同設備上的其他 XLA 張量。因此，程式碼如下

l_in = torch.randn(10, device=xm.xla_device())
linear = torch.nn.Linear(10, 20)
l_out = linear(l_in)
print(l_out)
# Input tensor is not an XLA tensor: torch.FloatTensor

會擲回錯誤，因為 torch.nn.Linear 模組在 CPU 上。

在 XLA 設備上執行模型¶

建立新的 PyTorch 網路或轉換現有網路以在 XLA 設備上執行只需要幾行程 XLA 特定的程式碼。以下程式碼片段重點說明了在單一設備上以及使用 XLA 多程序在多個設備上執行時，這些行是什麼。

在單一 XLA 設備上執行¶

以下程式碼片段顯示了在單一 XLA 設備上訓練的網路

import torch_xla.core.xla_model as xm

device = xm.xla_device()
model = MNIST().train().to(device)
loss_fn = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

for data, target in train_loader:
  optimizer.zero_grad()
  data = data.to(device)
  target = target.to(device)
  output = model(data)
  loss = loss_fn(output, target)
  loss.backward()

  optimizer.step()
  xm.mark_step()

此程式碼片段重點說明了將模型切換到 XLA 上執行的容易程度。模型定義、資料載入器、優化器和訓練迴圈可以在任何設備上工作。唯一特定於 XLA 的程式碼是幾行程，用於取得 XLA 設備並標記步驟。在每次訓練迭代結束時呼叫 xm.mark_step() 會導致 XLA 執行其目前的圖形並更新模型的參數。如需 XLA 如何建立圖形和執行操作的詳細資訊，請參閱 XLA 張量深入探討。

使用多程序在多個 XLA 設備上執行¶

PyTorch/XLA 可以透過在多個 XLA 設備上執行來輕鬆加速訓練。以下程式碼片段顯示了如何執行

import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl
import torch_xla.distributed.xla_multiprocessing as xmp

def _mp_fn(index):
  device = xm.xla_device()
  mp_device_loader = pl.MpDeviceLoader(train_loader, device)

  model = MNIST().train().to(device)
  loss_fn = nn.NLLLoss()
  optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)

  for data, target in mp_device_loader:
    optimizer.zero_grad()
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()
    xm.optimizer_step(optimizer)

if __name__ == '__main__':
  xmp.spawn(_mp_fn, args=())

此多設備程式碼片段與先前的單一設備程式碼片段之間存在三個差異。讓我們逐一介紹。

xmp.spawn()
- 建立每個都執行 XLA 設備的程序。
- 每個程序只能存取分配給目前程序的設備。例如，在 TPU v4-8 上，將會產生 4 個程序，並且每個程序都將擁有一個 TPU 設備。
- 請注意，如果您在每個程序上列印 xm.xla_device()，您將會在所有設備上看到 xla:0。這是因為每個程序只能看到一個設備。這並不意味著多程序不起作用。只有在 TPU v2 和 TPU v3 上使用 PJRT 執行階段時才會執行，因為將會有 #devices/2 個程序，並且每個程序都將有 2 個執行緒（如需詳細資訊，請查看此文件）。
MpDeviceLoader
- 將訓練資料載入每個設備。
- MpDeviceLoader 可以包裝在 torch 資料載入器上。它可以將資料預載入設備，並將資料載入與設備執行重疊，以提高效能。
- MpDeviceLoader 也會在產生每 batches_per_execution（預設為 1）批次時為您呼叫 xm.mark_step。
xm.optimizer_step(optimizer)
- 整合設備之間的梯度，並發出 XLA 設備步驟計算。
- 它幾乎是一個 all_reduce_gradients + optimizer.step() + mark_step，並傳回減少的損失。

模型定義、優化器定義和訓練迴圈保持不變。

**注意：**請務必注意，在使用多程序時，使用者只能從 xmp.spawn() 的目標函式（或任何在呼叫堆疊中具有 xmp.spawn() 作為父項的函式）中開始擷取和存取 XLA 設備。

如需使用多程序在多個 XLA 設備上訓練網路的詳細資訊，請參閱完整的多程序範例。

在 TPU Pod 上執行¶

不同加速器的多主機設定可能非常不同。本文件將討論多主機訓練中與設備無關的部分，並以 TPU + PJRT 執行階段（目前在 1.13 和 2.x 版本中可用）為例。

在您開始之前，請先查看我們在這裡的使用者指南，其中將說明一些 Google Cloud 基礎知識，例如如何使用 gcloud 命令以及如何設定您的專案。您也可以查看這裡以取得所有 Cloud TPU 的操作說明。本文件將重點介紹 PyTorch/XLA 對設定的觀點。

假設您在 train_mnist_xla.py 中有上述來自上一節的 mnist 範例。如果它是單一主機多設備訓練，您將 ssh 到 TPUVM 並執行如下命令

PJRT_DEVICE=TPU python3 train_mnist_xla.py

現在，為了在 TPU v4-16（具有 2 個主機，每個主機具有 4 個 TPU 設備）上執行相同的模型，您需要

確保每個主機都可以存取訓練腳本和訓練資料。這通常是透過使用 gcloud scp 命令或 gcloud ssh 命令將訓練腳本複製到所有主機來完成的。
同時在所有主機上執行相同的訓練命令。

gcloud alpha compute tpus tpu-vm ssh $USER-pjrt --zone=$ZONE --project=$PROJECT --worker=all --command="PJRT_DEVICE=TPU python3 train_mnist_xla.py"

以上 gcloud ssh 命令將 ssh 到 TPUVM Pod 中的所有主機，並同時執行相同的命令。

**注意：**您需要在 TPUVM vm 之外執行上述 gcloud 命令。

多程序訓練和多主機訓練使用相同的模型代碼和訓練腳本。PyTorch/XLA 和底層基礎架構將確保每個設備都知道全局拓撲以及每個設備的本地和全局序號。跨設備通信將跨所有設備而不是本地設備進行。

有關 PJRT 運行時以及如何在 Pod 上運行的更多詳細信息，請參閱此文件。有關 PyTorch/XLA 和 TPU Pod 的更多信息，以及在 TPU Pod 上使用虛擬數據運行 resnet50 的完整指南，請參閱此指南。

XLA 張量深入探討¶

使用 XLA 張量和設備只需要更改幾行程式碼。但是，即使 XLA 張量的行為與 CPU 和 CUDA 張量非常相似，但它們的內部結構卻不同。本節介紹 XLA 張量的獨特之處。

XLA 張量是惰性的¶

CPU 和 CUDA 張量會立即或 急切地 啟動操作。另一方面，XLA 張量是 惰性的。它們會將操作記錄在圖表中，直到需要結果為止。像這樣延遲執行可以讓 XLA 優化執行。例如，可以將多個獨立操作的圖表融合到單個優化操作中。

惰性執行通常對呼叫者不可見。PyTorch/XLA 會自動構建圖表、將它們發送到 XLA 設備，並在 XLA 設備和 CPU 之間複製數據時進行同步。在執行優化器步驟時插入屏障會顯式同步 CPU 和 XLA 設備。有關我們的惰性張量設計的更多信息，您可以閱讀這篇論文。

XLA 張量和 bFloat16¶

PyTorch/XLA 在 TPU 上運行時可以使用 bfloat16 數據類型。事實上，PyTorch/XLA 在 TPU 上處理浮點類型（torch.float 和 torch.double）的方式不同。此行為由 XLA_USE_BF16 和 XLA_DOWNCAST_BF16 環境變數控制

默認情況下，torch.float 和 torch.double 在 TPU 上都是 torch.float。
如果設置了 XLA_USE_BF16，則 torch.float 和 torch.double 在 TPU 上都是 bfloat16。
如果設置了 XLA_DOWNCAST_BF16，則 torch.float 在 TPU 上是 bfloat16，而 torch.double 在 TPU 上是 float32。
如果 PyTorch 張量具有 torch.bfloat16 數據類型，它將直接映射到 TPU bfloat16（XLA BF16 基元類型）。

開發人員應注意，無論 TPU 上的 XLA 張量實際使用什麼數據類型，它們始終會報告其 PyTorch 數據類型。此轉換是自動且不透明的。如果將 TPU 上的 XLA 張量移回 CPU，它將從其實際數據類型轉換為其 PyTorch 數據類型。根據您的程式碼如何運作，由處理單元類型觸發的這種轉換可能很重要。

内存佈局¶

XLA 張量的內部數據表示對用戶是不透明的。與 CPU 和 CUDA 張量不同，它們不會公開其存儲，並且始終顯示為連續的。這允許 XLA 調整張量的内存佈局以獲得更好的性能。

在 CPU 和 XLA 設備之間移動 XLA 張量¶

XLA 張量可以從 CPU 移動到 XLA 設備，也可以從 XLA 設備移動到 CPU。如果移動了視圖，則它正在查看的數據也會被複製到另一個設備，並且視圖關係不會保留。換句話說，一旦數據被複製到另一個設備，它就與其以前的設備或其上的任何張量沒有關係。同樣，根據您的程式碼如何運作，理解和適應這種轉變可能很重要。

保存和加載 XLA 張量¶

在保存之前，應將 XLA 張量移動到 CPU，如下面的程式碼片段所示

import torch
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()

t0 = torch.randn(2, 2, device=device)
t1 = torch.randn(2, 2, device=device)

tensors = (t0.cpu(), t1.cpu())

torch.save(tensors, 'tensors.pt')

tensors = torch.load('tensors.pt')

t0 = tensors[0].to(device)
t1 = tensors[1].to(device)

這使您可以將加載的張量放在任何可用的設備上，而不僅僅是初始化它們的設備上。

根據上面關於將 XLA 張量移動到 CPU 的說明，在使用視圖時必須小心。建議您在加載張量並將其移動到目標設備後重新創建視圖，而不是保存視圖。

提供了一個實用程序 API，可以通過先將數據移動到 CPU 來保存數據

import torch
import torch_xla
import torch_xla.core.xla_model as xm

xm.save(model.state_dict(), path)

在多個設備的情況下，上述 API 將僅保存主設備序號 (0) 的數據。

如果内存與模型參數的大小相比有限，則提供了一個 API 來減少主機上的内存占用

import torch_xla.utils.serialization as xser

xser.save(model.state_dict(), path)

此 API 一次將一個 XLA 張量流式傳輸到 CPU，從而減少了使用的主機内存量，但它需要匹配的加載 API 才能恢復

import torch_xla.utils.serialization as xser

state_dict = xser.load(path)
model.load_state_dict(state_dict)

可以直接保存 XLA 張量，但不建議這樣做。XLA 張量始終會加載回保存它們的設備，如果該設備不可用，則加載將失敗。PyTorch/XLA 和所有 PyTorch 一樣，都處於積極開發中，這種行為在將來可能會發生變化。

進一步閱讀¶

更多文檔可在 PyTorch/XLA 存儲庫中找到。有關在 TPU 上運行網絡的更多示例，請參見此處。

PyTorch/XLA API¶

xla_model¶

torch_xla.core.xla_model.xla_device(n=None, devkind=None)[原始碼]¶

返回給定 XLA 設備實例。

參數

n (python:int, 可選) – 要返回的特定實例（序號）。如果指定，則將返回特定的 XLA 設備實例。否則，將返回 devkind 的第一個設備。
devkind (string..., 可選) – 如果指定，則為 TPU、GPU 或 CPU 之一。

返回

具有請求實例的 torch.device。

torch_xla.core.xla_model.get_xla_supported_devices(devkind=None, max_devices=None)[原始碼]¶

返回給定類型支持的設備列表。

參數

devkind (string..., 可選) – 如果指定，則為 TPU、GPU 或 CPU 之一（“GPU” XLA 設備當前未實現）。
max_devices (python:int, 可選) – 要返回的該類型設備的最大數量。

返回

設備字符串列表。

torch_xla.core.xla_model.xla_device_hw(device)[原始碼]¶

返回給定設備的硬件類型。

參數: device (string or torch.device) – 將映射到真實設備的 xla 設備。
返回: 給定設備的硬件類型（CPU、TPU、GPU）的字符串表示形式。

torch_xla.core.xla_model.get_ordinal(defval=0)[原始碼]¶

檢索當前線程的複製序號。

序號範圍從 0 到 xrt_world_size() 減 1。

參數: defval (python:int, 可選) – 如果沒有可用的複製信息，則返回的默認值。對於 PjRt，將忽略此值。默認值：0
返回: 當前線程的複製序號。

torch_xla.core.xla_model.get_local_ordinal(defval=0)[原始碼]¶

檢索當前線程的本地複製序號。

本地序號範圍從 0 到本地設備數量減 1。

參數: defval (python:int, 可選) – 如果沒有可用的複製信息，則返回的默認值。對於 PjRt，將忽略此值。默認值：0
返回: 當前線程的本地複製序號。

torch_xla.core.xla_model.is_master_ordinal(local=True)[原始碼]¶

檢查當前進程是否是主序號 (0)。

參數: local (bool) – 是否應檢查本地或全局主序號。在多主機複製的情況下，只有一個全局主序號（主機 0，設備 0），而有 NUM_HOSTS 個本地主序號。默認值：True
返回: 一個布爾值，指示當前進程是否是主序號。

torch_xla.core.xla_model.xrt_world_size(defval=1)[原始碼]¶

檢索參與複製的設備數量。

參數: defval (python:int, 可選) – 如果沒有可用的複製信息，則返回的默認值。默認值：1
返回: 參與複製的設備數量。

torch_xla.core.xla_model.all_reduce(reduce_type, inputs, scale=1.0, groups=None, cctx=None, pin_layout=True)[原始碼]¶

對輸入張量執行就地 reduce 操作。

參數

reduce_type (字串) – xm.REDUCE_SUM、xm.REDUCE_MUL、xm.REDUCE_AND、xm.REDUCE_OR、xm.REDUCE_MIN 和 xm.REDUCE_MAX 其中之一。
inputs – 單個 torch.Tensor 或要對其執行 all reduce 操作的 torch.Tensor 清單。
scale (python:float) – 在 reduce 後套用的預設縮放值。預設值：1.0
groups (清單，可選) –
表示 all_reduce() 操作的副本群組的清單。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

定義了兩個群組，一個包含 [0, 1, 2, 3] 副本，另一個包含 [4, 5, 6, 7] 副本。如果為 None，則只有一個群組，其中包含所有副本。
pin_layout (布林值，可選) – 是否要為此通訊操作固定佈局。固定佈局可以防止參與通訊的每個程序的程式略有不同時可能發生的資料損毀，但它可能會導致某些 xla 編譯失敗。當您看到錯誤訊息「HloModule has a mix of layout constrained」時，請取消固定佈局。

返回

如果傳入單個 torch.Tensor，則傳回值是一個包含 reduce 值（跨副本）的 torch.Tensor。如果傳入清單／元組，則此函數會對輸入張量執行就地 all-reduce 操作，並傳回清單／元組本身。

torch_xla.core.xla_model.all_gather(value, dim=0, groups=None, output=None, pin_layout=True)[原始碼]¶

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

參數

value (torch.Tensor) – 輸入張量。
dim (python:int) – 收集維度。預設值：0
groups (清單，可選) –
表示 all_gather() 操作的副本群組的清單。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

定義了兩個群組，一個包含 [0, 1, 2, 3] 副本，另一個包含 [4, 5, 6, 7] 副本。如果為 None，則只有一個群組，其中包含所有副本。
output (torch.Tensor) – 可選的輸出張量。
pin_layout (布林值，可選) – 是否要為此通訊操作固定佈局。固定佈局可以防止參與通訊的每個程序的程式略有不同時可能發生的資料損毀，但它可能會導致某些 xla 編譯失敗。當您看到錯誤訊息「HloModule has a mix of layout constrained」時，請取消固定佈局。

返回

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

torch_xla.core.xla_model.all_to_all(value, split_dimension, concat_dimension, split_count, groups=None, pin_layout=True)[原始碼]¶

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

請參閱：https://www.tensorflow.org/xla/operation_semantics#alltoall

參數

value (torch.Tensor) – 輸入張量。
split_dimension (python:int) – 應該在其上進行分割的維度。
concat_dimension (python:int) – 應該在其上進行串聯的維度。
split_count (python:int) – 分割計數。
groups (清單，可選) –
表示 all_reduce() 操作的副本群組的清單。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

定義了兩個群組，一個包含 [0, 1, 2, 3] 副本，另一個包含 [4, 5, 6, 7] 副本。如果為 None，則只有一個群組，其中包含所有副本。
pin_layout (布林值，可選) – 是否要為此通訊操作固定佈局。固定佈局可以防止參與通訊的每個程序的程式略有不同時可能發生的資料損毀，但它可能會導致某些 xla 編譯失敗。當您看到錯誤訊息「HloModule has a mix of layout constrained」時，請取消固定佈局。

返回

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

torch_xla.core.xla_model.add_step_closure(closure, args=(), run_async=False)[原始碼]¶

將一個 closure 新增到要在步驟結束時執行的 closure 清單中。

在模型訓練期間，很多時候需要列印／報告（列印到主控台、發佈到 TensorBoard 等）需要檢查中 intermediary 張量內容的資訊。在模型程式碼的不同點檢查不同張量的內容需要多次執行，通常會導致效能問題。新增步驟 closure 將確保它在 barrier 之後執行，屆時所有活動張量都將已具體化為裝置資料。活動張量將包含 closure 參數捕獲的那些張量。因此，使用 add_step_closure() 將確保只執行一次，即使有多個 closure 排隊，需要檢查多個張量。步驟 closure 將按照它們排隊的順序依序執行。請注意，即使使用此 API 會最佳化執行，但建議您每 N 個步驟限制一次列印／報告事件。

參數

closure (可呼叫物件) – 要呼叫的函數。
args (元組) – 要傳遞給 closure 的參數。
run_async – 如果為 True，則非同步執行 closure。

torch_xla.core.xla_model.wait_device_ops(devices=[])[原始碼]¶

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

參數: devices (字串...，可選) – 需要等待其非同步操作的裝置。如果為空，則將等待所有本地裝置。

torch_xla.core.xla_model.optimizer_step(optimizer, barrier=False, optimizer_args={}, groups=None, pin_layout=True)[原始碼]¶

執行提供的最佳化器步驟，並發出 XLA 裝置步驟計算。

參數

optimizer (torch.Optimizer) – 需要呼叫其 step() 函數的 torch.Optimizer 執行個體。將使用 optimizer_args 命名的參數呼叫 step() 函數。
barrier (布林值，可選) – 是否應該在此 API 中發出 XLA 張量 barrier。如果使用 PyTorch XLA ParallelLoader 或 DataParallel 支援，則不需要這樣做，因為 barrier 將由 XLA 資料載入器迭代器 next() 呼叫發出。預設值：False
optimizer_args (字典，可選) – optimizer.step() 呼叫的命名參數字典。
groups (清單，可選) –
表示 all_reduce() 操作的副本群組的清單。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

定義了兩個群組，一個包含 [0, 1, 2, 3] 副本，另一個包含 [4, 5, 6, 7] 副本。如果為 None，則只有一個群組，其中包含所有副本。
pin_layout (布林值，可選) – 在減少梯度時是否固定佈局。詳細資訊請參閱 xm.all_reduce。

返回

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

torch_xla.core.xla_model.save(data, file_or_path, master_only=True, global_master=False)[原始碼]¶

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

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

參數

data – 要儲存的輸入資料。Python 物件（清單、元組、集合、字典等）的任何巢狀組合。
file_or_path – 資料儲存操作的目的地。檔案路徑或 Python 檔案物件。如果 master_only 為 False，則路徑或檔案物件必須指向不同的目的地，否則來自同一個主機的所有寫入都將互相覆蓋。
master_only (布林值，可選) – 是否只有主裝置應該儲存資料。如果為 False，則 file_or_path 參數應該是參與複製的每個序號的不同檔案或路徑，否則同一個主機上的所有副本都將寫入同一個位置。預設值：True
global_master (布林值，可選) – 當 master_only 為 True 時，此旗標控制每個主機的主裝置（如果 global_master 為 False）是否儲存內容，還是只有全域主裝置（序號 0）儲存內容。預設值：False

torch_xla.core.xla_model.rendezvous(tag, payload=b'', replicas=[])[原始碼]¶

等待所有網格客戶端到達指定的集合點。

參數

tag (字串) – 要加入的集合點的名稱。
payload (位元組，可選) – 要傳送到集合點的有效負載。
replicas (清單， python:int) – 參與集合點的副本序號。空表示網格中的所有副本。預設值：[]

返回

由所有其他核心交換的有效負載，核心序號 i 的有效負載位於傳回元組中的位置 i。

torch_xla.core.xla_model.do_on_ordinals(target, data=(), ordinals=(0, ))[原始碼]¶

僅在給定的序號集合上執行函數。

參數

target (可呼叫物件) – 要在 ordinals 上執行的函數。
data – target 函數的任何輸入資料，其中包含張量。 target 函數使用的所有 XLA 張量都必須在此參數中傳遞。函數使用的所有其他資料都可以像往常一樣由 Python 直譯器捕獲。預設值：()
ordinals (清單， python:int) – 應該執行 target 函數的序號清單／集合。預設值：(0,)

返回

在執行 target 函數的序號中，函數傳回值，否則為 None。

torch_xla.core.xla_model.mesh_reduce(tag, data, reduce_fn)[原始碼]¶

執行圖外客戶端網格 reduce。

參數

tag (字串) – 要加入的集合點的名稱。
data – 要 reduce 的資料。 reduce_fn 可呼叫物件將接收一個清單，其中包含來自所有網格客戶端程序（每個核心一個）的相同資料的副本。
reduce_fn (可呼叫物件) – 接收 data 類物件清單並傳回 reduce 結果的函數。

返回

reduce 的值。

torch_xla.core.xla_model.set_rng_state(seed, device=None)[原始碼]¶

設定亂數產生器狀態。

參數

seed (python:integer) – 要設定的狀態。
device (string, optional) – 需要設定 RNG 狀態的裝置。如果遺漏，將設定預設裝置種子。

torch_xla.core.xla_model.get_rng_state(device=None)[source]¶

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

參數: device (string, optional) – 需要擷取其 RNG 狀態的裝置。如果遺漏，將設定預設裝置種子。
返回: RNG 狀態，以整數表示。

torch_xla.core.xla_model.get_memory_info(device)[source]¶

擷取裝置記憶體資訊。

參數: device (string) – 要求其記憶體資訊的裝置。
返回: 一個字典，包含 kb_free（可用記憶體，以 KB 為單位）和 kb_total（總記憶體，以 KB 為單位）鍵值。

torch_xla.core.functions.all_reduce(reduce_type, value, scale=1.0, groups=None)[source]¶

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

這與 xm.all_reduce() 相同，但支援 Autograd 微分。

參數

reduce_type (string) – REDUCE_SUM、REDUCE_MUL、REDUCE_AND、REDUCE_OR、REDUCE_MIN 和 REDUCE_MAX 之一。
value (torch.Tensor) – 要對其執行全縮減運算的張量。
scale (python:float) – 在 reduce 後套用的預設縮放值。預設值：1.0
groups (清單，可選) –
表示 all_reduce() 操作的副本群組的清單。範例：[[0, 1, 2, 3], [4, 5, 6, 7]]

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

返回

跨選定副本縮減後的值。

torch_xla.core.functions.all_gather(value, dim=0)[source]¶

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

這與 xm.all_gather() 相同，但支援 Autograd 微分。

參數

value (torch.Tensor) – 輸入張量。
dim (python:int) – 收集維度。預設值：0

返回

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

torch_xla.core.functions.nms(boxes, scores, score_threshold, iou_threshold, output_size)[source]¶

執行非極大值抑制運算。

參數

boxes (torch.Tensor) – 形狀為 [N, 4] 的 torch.Tensor，以 (y0, x0, y1, x1) 形式列出方框座標。
scores (torch.Tensor) – 形狀為 [N] 的 torch.Tensor，列出每個方框的分數。
score_threshold (torch.Tensor) – 方框被視為有效的最低分數。
iou_threshold (torch.Tensor) – 觸發重疊邏輯的最小 IOU（交集超過聯集）分數。
output_size (python:int) – 傳回索引的最大數量（必須小於或等於 N）。

返回

一個 torch.Tensor 元組，第一個元素是選定的方框索引，第二個元素是有效方框的數量。

distributed¶

class torch_xla.distributed.parallel_loader.ParallelLoader(loader, devices, batchdim=0, batches_per_execution=1, loader_prefetch_size=8, device_prefetch_size=4, host_to_device_transfer_threads=1)[source]¶

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

參數

loader (torch.utils.data.DataLoader) – 要包裝的 PyTorch DataLoader。
devices (torch.device…) – 資料必須傳送到的裝置清單。 loader 傳回的第 i 個樣本將傳送到 devices[i % len(devices)]。
batchdim (python:int, optional) – 保存批次大小的維度。預設值：0
loader_prefetch_size (python:int, optional) – 從 loader 讀取樣本的執行緒所使用的佇列的最大容量，這些樣本將由將資料上傳到裝置的工作執行緒處理。預設值：8
device_prefetch_size (python:int, optional) – 每個裝置佇列的最大大小，工作執行緒會將已傳送到裝置的張量存放到這些佇列中。預設值：4
host_to_device_transfer_threads (python:int, optional) – 平行工作以將資料從載入器佇列傳輸到裝置佇列的執行緒數。預設值：1

per_device_loader(device)[source]¶

擷取指定裝置的載入器迭代器物件。

參數: device (torch.device) – 正在請求其載入器的裝置。
返回: device 的載入器迭代器物件。這不是 torch.utils.data.DataLoader 介面，而是一個 Python 迭代器，它傳回與包裝的 torch.utils.data.DataLoader 相同的張量資料結構，但位於 XLA 裝置上。

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

啟用基於多程序的複製。

參數

fn (callable) – 要為參與複製的每個裝置呼叫的函式。將使用第一個參數作為複製中程序的全域索引來呼叫該函式，後跟在 args 中傳遞的參數。
args (tuple) – fn 的參數。預設值：空元組
nprocs (python:int) – 複製的程序/裝置數量。目前，如果指定，可以是 1 或最大裝置數量。
join (bool) – 呼叫是否應該阻塞，等待已產生的程序完成。預設值：True
daemon (bool) – 正在產生的程序是否應該設定 daemon 旗標（請參閱 Python 多程序 API）。預設值：False
start_method (string) – Python multiprocessing 程序建立方法。預設值：spawn

返回

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

class torch_xla.distributed.xla_multiprocessing.MpModelWrapper(model)[source]¶

包裝模型，以便在使用 fork 方法時將主機記憶體使用量降至最低。

此類別應與 spawn(…, start_method=’fork’) API 一起使用，以最大程度地減少主機記憶體的使用。模型不是在每個多程序程序上建立，從而複製模型的初始主機記憶體，而是在全域範圍內建立一次，然後在 spawn() 目標函式內移動到每個裝置中。範例

WRAPPED_MODEL = xmp.MpModelWrapper(MyNetwork())

def _mp_fn(index, ...):
  device = xm.xla_device()
  model = WRAPPED_MODEL.to(device)
  ...

xmp.spawn(_mp_fn, ..., start_method='fork')

此方法有兩個優點。首先，它僅使用一份記憶體頁面副本來存放原始模型權重，其次，它通過在過程中降低系統記憶體的負載，將包裝模型移動到每個裝置中。

to(device)[source]¶

擷取已移動到指定裝置的模型。

參數: device (torch.device) – 模型應移動到的裝置。
返回: 指定裝置上的模型。

class torch_xla.distributed.xla_multiprocessing.MpSerialExecutor[source]¶

用於在多核心程序之間以序列化方式執行函式的工具。

範例

# At global scope.
SERIAL_EXEC = xmp.MpSerialExecutor()

def load_dataset(path):
  return maybe_download_and_load(path)

def _mp_fn(index, ...):
  # Avoid all cores downloading the same data with the serial executor.
  dataset = SERIAL_EXEC.run(lambda: load_dataset('/tmp/mnist-data'))
  ...

xmp.spawn(_mp_fn, ...)

run(fn)[source]¶

針對每個核心程序序列化執行提供的函式。

參數: fn (callable) – 以序列化方式執行的函式。
返回: fn 的傳回值。

utils¶

class torch_xla.utils.tf_record_reader.TfRecordReader(path, compression='', buffer_size=16777216, transforms=None)[source]¶

讀取 TfRecords 或 TfExamples。

參數

path (string) – 包含 TfRecords 的檔案路徑。
compression (string, optional) – 壓縮類型。空字串表示不壓縮，否則為 ZLIB 或 GZIP。預設值：不壓縮。
buffer_size (python:int, optional) – 用於讀取 TfRecords 的緩衝區大小。預設值：16 * 1024 * 1024
transforms (dict, optional) – 一個字典，其鍵值與 TfExample 標籤名稱匹配，值則是可以呼叫來轉換匹配張量資料的 callable，或者表示字串轉換的 STR。

class torch_xla.utils.utils.SampleGenerator(data, sample_count)[source]¶

迭代器，它傳回給定輸入資料的多個樣本。

可用來取代 PyTorch DataLoader 以產生合成數據。

參數

data – 每次迭代步驟應返回的數據。
sample_count – 要返回的最大 data 樣本數。

類別 torch_xla.utils.utils.DataWrapper[來源]¶: 用於包裝要發送到設備的數據結構的實用程序類別。

torch_xla.utils.serialization.save(data, path, master_only=True, global_master=False)[來源]¶

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

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

參數

data – 要儲存的輸入資料。Python 物件（清單、元組、集合、字典等）的任何巢狀組合。
path – 數據保存操作的目標文件。如果 master_only 為 False，則路徑必須指向不同的目標，否則來自同一主機的所有寫入都將覆蓋彼此。
master_only (bool, 可選) – 是否只有主設備應保存數據。如果為 False，則 path 參數對於參與複製的每個序號應該是不同的路徑，否則同一主機上的所有副本都將寫入相同的位置。默認值：True
global_master (布林值，可選) – 當 master_only 為 True 時，此旗標控制每個主機的主裝置（如果 global_master 為 False）是否儲存內容，還是只有全域主裝置（序號 0）儲存內容。預設值：False

torch_xla.utils.serialization.load(path)[來源]¶

載入先前使用 save() API 保存的數據。

參數: path (str) – 傳遞給 save() API 的路徑。
返回: 載入的數據。

torch_xla.utils.gcsfs.open(path, mode='r', encoding=None)[來源]¶

開啟 Google Cloud Storage (GCS) 文件以進行讀取或寫入。

參數

path (string) – 文件的 GCS 路徑。必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 存储空間的名稱，而 PATH 是以 / 分隔的路徑。
mode (string, 可選) – 開啟模式，類似於 open() API。默認值：'r'
encoding (string, 可選) – 以文字模式開啟時用於將字節解碼為字串的字元編碼。默認值：None

返回

GCS 文件物件。

torch_xla.utils.gcsfs.list(path)[來源]¶

列出 GCS 存储空間的內容。

參數: path (string) – 文件的 GCS 路徑。必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 存储空間的名稱，而 PATH 是以 / 分隔的路徑。
返回: GcsBlob 物件的清單。

torch_xla.utils.gcsfs.stat(path)[來源]¶

擷取 GCS 文件的資訊。

參數: path (string) – 文件的 GCS 路徑。必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 存储空間的名稱，而 PATH 是以 / 分隔的路徑。
返回: 一個 GcsBlob 物件。

torch_xla.utils.gcsfs.remove(path)[來源]¶

移除 GCS Blob。

參數: path (string) – 文件的 GCS 路徑。必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 存储空間的名稱，而 PATH 是以 / 分隔的路徑。

torch_xla.utils.gcsfs.rmtree(path)[來源]¶

移除給定路徑內的所有 GCS Blob。

參數

path (string) –

文件模式或資料夾的 GCS 路徑。必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 的名稱

存储空間，而 PATH 是以 / 分隔的路徑。

torch_xla.utils.gcsfs.read(path)[來源]¶

讀取 GCS Blob 的全部內容。

參數: path (string) – 文件的 GCS 路徑。必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 存储空間的名稱，而 PATH 是以 / 分隔的路徑。
返回: GCS Blob 中存儲的字節。

torch_xla.utils.gcsfs.write(path, content)[來源]¶

將字串/字節或文件寫入 GCS Blob。

參數

path (string) – 文件的 GCS 路徑。必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 存储空間的名稱，而 PATH 是以 / 分隔的路徑。
content (string, bytes 或 文件物件) – 要寫入 path 的內容。

torch_xla.utils.gcsfs.generic_open(path, mode='r', encoding=None)[來源]¶

開啟文件（GCS 或非 GCS）以進行讀取或寫入。

參數

path (string) –
要開啟的文件的路徑。如果是 GCS 路徑，則必須是“gs://BUCKET_NAME/PATH”，其中 BUCKET_NAME 是 GCS 的名稱

存储空間，而 PATH 是以 / 分隔的路徑。
mode (string, 可選) – 開啟模式，類似於 open() API。默認值：'r'
encoding (string, 可選) – 以文字模式開啟時用於將字節解碼為字串的字元編碼。默認值：None

返回

開啟的文件物件。

torch_xla.utils.gcsfs.generic_read(path)[來源]¶

讀取所提供位置的全部內容。

參數: path (string) – 要讀取的 GCS 路徑或本機路徑。
返回: GCS Blob 或本機文件中存儲的字節。

torch_xla.utils.gcsfs.generic_write(output_string, path, makedirs=False)[來源]¶

將字串/字節或文件寫入 GCS Blob 或本機磁碟。

根據傳入的路徑，此 API 可以寫入本機或 GCS 文件。檢查 path 是否以“gs://”前綴開頭，如果不是，則使用 open。

參數

output_string (string) – 要寫入輸出的字串。
path (string) – 輸出的 GCS 路徑或本機路徑。
makedirs (bool) – 如果缺少 path 父資料夾，是否應建立它們。默認值：False

torch_xla.utils.gcsfs.is_gcs_path(path)[來源]¶

檢查路徑是否為 GCS 路徑。

參數: path (string) – 要檢查的路徑。
返回: path 是否為 GCS 路徑。

類別 torch_xla.utils.cached_dataset.CachedDataset(data_set, path, max_files_per_folder=1000, compress=True)[來源]¶

通過提供文件緩存來包裝現有的 torch.utils.data.Dataset。

CachedDataset 可用於交換處理原始數據集所需的 CPU/RAM 資源與存儲/網絡資源。範例

train_dataset = datasets.MNIST(
    FLAGS.datadir,
    train=True,
    download=True,
    transform=transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))]))
train_dataset = CachedDataset(train_dataset, FLAGS.dscache_dir)

CachedDataset 將透明地緩存原始 Dataset 樣本，以便在第一次之後的每次運行都不會觸發與原始樣本處理相關的任何更多 CPU/RAM 使用。完全緩存 CachedDataset 後，可以將其導出（即 tar.gz）並在不同的機器上使用。只需解壓 tar.gz 並傳遞 None 作為原始 Dataset：範例

train_dataset = CachedDataset(None, FLAGS.dscache_dir)

要完全緩存 CachedDataset，只需運行 warmup() API。保存在 GCS 上的 CachedDataset 具備無需顯式導出即可從不同機器使用的優勢。

參數

data_set (torch.utils.data.Dataset) – 要緩存的原始 torch.utils.data.Dataset。如果所有輸入樣本都存儲在 path 文件夾中，則可以將其設置為 None。
path (string) – 應存儲/加載數據集樣本的路徑。除非已存儲所有樣本，否則 path 需要可寫入。path 可以是 GCS 路徑（以 gs:// 為前綴）。
max_files_per_folder (python:int) – 要存儲在單個文件夾中的最大文件數。如果 data_set 為 None，則忽略此值並從緩存的元數據中獲取。默認值：1000
compress (bool) – 是否應壓縮保存的樣本。壓縮以 CPU 用於壓縮/解壓縮為代價來節省空間。如果 data_set 為 None，則忽略此值並從緩存的元數據中獲取。默認值：True

test¶

疑難排解¶

請注意，本節中的信息可能會在未來版本的*PyTorch/XLA*軟件中刪除，因為其中許多信息是特定於可能會更改的給定內部實現的。

為了診斷問題，我們可以使用*PyTorch/XLA*提供的執行指標和計數器。當模型速度緩慢時，要檢查的**第一件事**是生成指標報告。

指標報告對診斷問題非常有幫助。如果您有指標報告，請嘗試將其包含在發送給我們的錯誤報告中。

執行 A 自動指標分析¶

我們提供自動分析指標報告並提供摘要的方法。只需使用 PT_XLA_DEBUG=1 運行您的工作負載即可。一些示例輸出將是

pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss, aten::_ctc_loss_backward,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps
pt-xla-profiler: TransferFromServerTime too frequent: 12 counts during 12 steps

以下部分將說明如何獲取和理解更詳細的指標報告。

獲取指標報告¶

將以下行放入您的程序中以生成報告

import torch_xla.debug.metrics as met

# For short report that only contains a few key metrics.
print(met.short_metrics_report())
# For full report that includes all metrics.
print(met.short_metrics_report())

了解指標報告¶

該報告包括以下內容

我們發出 *XLA* 編譯的次數以及發出所花費的時間。
我們執行的次數以及執行所花費的時間。
我們建立/銷毀等的裝置資料控制代碼數量。

此資訊會以樣本的百分位數形式報告。例如：

Metric: CompileTime
  TotalSamples: 202
  Counter: 06m09s401ms746.001us
  ValueRate: 778ms572.062us / second
  Rate: 0.425201 / second
  Percentiles: 1%=001ms32.778us; 5%=001ms61.283us; 10%=001ms79.236us; 20%=001ms110.973us; 50%=001ms228.773us; 80%=001ms339.183us; 90%=001ms434.305us; 95%=002ms921.063us; 99%=21s102ms853.173us

我們也提供計數器，它們是用來追蹤內部軟體狀態的已命名整數變數。例如：

Counter: CachedSyncTensors
  Value: 395

在本報告中，任何以 aten:: 開頭的計數器都表示 *XLA* 裝置和 CPU 之間的上下文切換，這可能是模型程式碼中潛在的效能最佳化區域。

計數器可用於瞭解哪些操作被路由回 *PyTorch* 的 CPU 引擎。它們使用 C++ 命名空間進行完整限定。

Counter: aten::nonzero
  Value: 33

如果您看到 aten:: 作業，而不是 nonzero 和 _local_scalar_dense，這通常表示 PyTorch/XLA 中缺少降低。歡迎在 GitHub 問題上為其開啟功能請求。

清除指標報告¶

如果您想清除步驟/時期之間的指標，可以使用：

import torch_xla.debug.metrics as met

met.clear_all()

效能分析¶

若要深入分析您的工作負載以瞭解瓶頸，請查看以下資源：

已知的效能注意事項¶

PyTorch/XLA 的語義行為與常規 PyTorch 相似，並且 XLA 張量與 CPU 和 GPU 張量共用完整的張量介面。但是，XLA/硬體中的限制和惰性求值模型表明某些模式可能會導致效能不佳。

如果您的模型顯示效能不佳，請牢記以下注意事項：

如果重新編譯次數過多，XLA/TPU 的效能會下降。

XLA 編譯成本很高。每次遇到新形狀時，PyTorch/XLA 都會自動重新編譯圖表。通常，模型應在幾個步驟內穩定下來，並且您可以看到其餘訓練的巨大加速。

為了避免重新編譯，不僅形狀必須恆定，而且所有主機中 XLA 裝置上的計算也必須恆定。

可能來源:
- nonzero 的直接或間接使用會引入動態形狀；例如，遮罩索引 base[index]，其中 index 是遮罩張量。
- 步驟之間迭代次數不同的迴圈可能會導致不同的執行圖表，因此需要重新編譯。
解決方案:
- 張量形狀在迭代之間應保持一致，或者應使用少量形狀變化。
- 盡可能將張量填充到固定大小。
某些操作沒有到 XLA 的原生轉換。

對於這些操作，PyTorch/XLA 會自動轉移到 CPU 記憶體，在 CPU 上進行評估，然後將結果轉移回 XLA 裝置。在訓練步驟中執行過多此類操作可能會導致顯著的減速。

可能來源:
- item() 操作明確要求評估結果。除非必要，否則不要使用它。
解決方案:
- 對於大多數操作，我們可以將它們降低到 XLA 來修復它。查看指標報告章節以找出缺少的操作，並在 GitHub 上開啟功能請求。
- 即使已知 PyTorch 張量是標量，也要避免使用 tensor.item()。將其保留為張量並對其使用張量操作。
- 在適用的情況下，使用 torch.where 來替換控制流程。例如，clip_grad*norm* 中使用的帶有 item() 的控制流程是有問題的，並且會影響效能，因此我們已通過呼叫 torch.where 來修補 clip_grad_norm_，這顯著提高了效能。 .. 程式碼區塊:: python
  
  … else
  
  device = parameters[0].device total_norm = torch.zeros([], device=device if parameters else None) for p in parameters
  
  param_norm = p.grad.data.norm(norm_type) ** norm_type total_norm.add_(param_norm)
  
  total_norm = (total_norm ** (1. / norm_type))
  
  clip_coef = torch.tensor(max_norm, device=device) / (total_norm + 1e-6) for p in parameters
  
  p.grad.data.mul_(torch.where(clip_coef < 1, clip_coef, torch.tensor(1., device=device)))
``torch_xla.distributed.data_parallel`` 中的迭代器可能會丟棄輸入迭代器中的最後幾批。

這是為了確保我們在所有 XLA 裝置上完成相同的工作量。

解決方案:
- 當資料集很小且步驟太少時，這可能會導致無效時期。因此，在這些情況下最好使用較小的批次大小。

XLA 張量怪癖¶

**XLA 張量內部是不透明的。** XLA 張量始終顯示為連續的並且沒有儲存空間。網路不應嘗試檢查 XLA 張量的步幅。
**XLA 張量在儲存之前應移至 CPU。** 直接儲存 XLA 張量會導致它們從儲存它們的裝置重新載入。如果在載入時裝置不可用，則載入將失敗。在儲存 XLA 張量之前將其移至 CPU 可讓您決定將載入的張量放在哪些裝置上。如果要在沒有 XLA 裝置的機器上載入張量，則必須執行此操作。但是，在儲存 XLA 張量之前，應小心地將其移至 CPU，因為跨裝置類型移動張量不會保留檢視關係。相反，應在載入張量後根據需要重建檢視。
**使用 Python 的 copy.copy 複製 XLA 張量會返回深層副本，而不是淺層副本。** 使用 XLA 張量的檢視來獲取其淺層副本。
**處理共用權重。** 模組可以通過將一個模組的參數設置為另一個模組來共用權重。這種模組權重的“綁定”應該在模組移至 XLA 裝置*之後*完成。否則，將在 XLA 裝置上建立共用張量的兩個獨立副本。

使用 debug_run.py 收集除錯資訊¶

scripts/debug_run.py 中提供了一個實用程式，可用於建立包含除錯 *PyTorch/XLA* 執行所需資訊的 tar.gz 封存檔。

範例

./scripts/debug_run.py --outfile /tmp/debug_run.tar.gz -- python -u SCRIPT [ARGS...]

建議使用 *python* -u 旗標來停用緩衝，以便正確地交錯擷取的日誌（否則 STDOUT 將在所有 STDERR 之後呈現）。

上面的命令列範例會將包含封存資訊的臨時資料夾保留在檔案系統中。使用 --tidy 旗標可以在退出時將其刪除

./scripts/debug_run.py --tidy --outfile /tmp/debug_run.tar.gz -- python -u SCRIPT [ARGS...]

然後應在必要時將 debug_run.tar.gz 檔案附加到錯誤報告中。

由於腳本會收集大量資料，因此通常應該讓它運行不超過一百步左右。

如果 SCRIPT 具有控制步數的參數，則應使用這些參數，否則按下 CTRL^C 將中斷執行。

還建議在單核心模式下執行，以最大程度地減少資料量。在除錯執行問題時，也強烈建議在單核心模式下執行。

常見問題¶

缺少 XLA 設定 錯誤訊息：如果您使用的是 TPU，則需要設定 XRT_TPU_CONFIG。如果使用的是 GPU，請為 N 個 GPU 設定 GPU_NUM_DEVICES=N。如果使用的是 CPU，請設定 XRT_DEVICE_MAP="CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0" 和 XRT_WORKERS="localservice:0;grpc://:9002"

PJRT 執行階段（測試版）¶

*本文件反映了當前每晚建置中 PJRT 支援的當前狀態*。有關最新穩定版本中的狀態，請參閱 r2.0 分支上的相同文件。

PyTorch/XLA 團隊目前正在從目前支援的 XRT 執行階段遷移到 JAX 使用的 PJRT 執行階段。

PJRT 在 PyTorch/XLA 2.0 中以預覽版提供。**我們計畫將 PJRT 作為我們官方支援的執行階段**，因此我們鼓勵所有使用者進行試用。我們的目標是在 2.1 版中穩定 PJRT，因此如果您在使用 PJRT 時遇到錯誤，請在 GitHub 上提交帶有 runtime 標籤的問題。

PyTorch/XLA r2.0 中的新功能:

如果您沒有傳入任何其他執行階段設定，則預設情況下會設定 PJRT。如果您繼續設定 XRT 設定 (XRT_TPU_CONFIG)，則此變更沒有影響
libtpu 中新的 TPU 執行階段實作將效能提高了高達 30%。
新的 xm.rendezvous 實作，可擴展至數千個 TPU 核心
[實驗性] TPU v2 和 v3 的 torch.distributed 支援，包括 pjrt:// init_method
[實驗性] PJRT 中的單主機 GPU 支援。多主機支援即將推出！

TL;DR¶

要使用 PJRT 預覽版執行階段，請將 PJRT_DEVICE 環境變數設定為 CPU、TPU 或GPU`
在 XRT 中，所有分散式工作負載都是多程序的，每個裝置一個程序。在 PJRT 中的 TPU v2 和 v3 上，工作負載是多程序和多執行緒的（4 個程序，每個程序 2 個執行緒），因此您的工作負載應該是執行緒安全的。有關更多資訊，請參閱 TPU v2/v3 上的多執行緒和 API 指南的多程序章節。要牢記的關鍵差異
- 若要以執行緒安全的方式初始化模型，請在初始化後跨複本廣播參數 (torch_xla.experimental.pjrt.broadcast_master_param) 或從通用檢查點載入每個複本的參數。
- 對於其他隨機數生成，請盡可能使用 torch.Generator。全域 torch RNG *不是* 執行緒安全的，即使您在複本中設定了相同的 torch.manual_seed。
- 要使用 torch.distributed，請匯入 torch_xla.experimental.pjrt_backend 並使用 pjrt:// init_method。
- 這些步驟對於 GPU 和 TPU v4 是可選的。

從 XRT 到 PJRT 的範例差異

 import os

 import torch
 import torch.nn as nn
 from torch.nn.parallel import DistributedDataParallel as DDP
 import torch.optim as optim
 import torch.distributed as dist
 import torch_xla.core.xla_model as xm
 import torch_xla.distributed.parallel_loader as pl
 import torch_xla.distributed.xla_backend
 import torch_xla.distributed.xla_multiprocessing as xmp
+import torch_xla.experimental.pjrt_backend
+import torch_xla.experimental.pjrt as pjrt


 def _mp_fn(index):
   device = xm.xla_device()
-  dist.init_process_group('xla', rank=xm.get_ordinal(), world_size=xm.xrt_world_size())
+  dist.init_process_group('xla', init_method='pjrt://')

   torch.manual_seed(42)
   model = nn.Linear(128, 10).to(device)

+  # Optional for TPU v4 and GPU
+  pjrt.broadcast_master_param(model)
   model = DDP(model, gradient_as_bucket_view=True)

   loss_fn = nn.MSELoss()
   optimizer = optim.SGD(model.parameters(), lr=.001)

   for i in range(10):
     data, target = torch.randn((128, 128), device=device), torch.randn((128, 10), device=device)

     optimizer.zero_grad()
     output = model(data)
     loss = loss_fn(output, target)
     loss.backward()

     optimizer.step()
     xm.mark_step()

   # Print mean parameters so we can confirm they're the same across replicas
   print([p.mean() for p in model.parameters()])

 if __name__ == '__main__':
-  os.environ['XRT_TPU_CONFIG'] = 'localservice;0;localhost:51011'
-  os.environ['MASTER_ADDR'] = 'localhost'
-  os.environ['MASTER_PORT'] = '12355'

+  # Recommended: set PJRT_DEVICE to your local device type
+  os.environ['PJRT_DEVICE'] = 'TPU'

   xmp.spawn(_mp_fn)

優點¶

簡單的執行階段設定：只需將 PJRT_DEVICE 設定為 TPU、CPU 或 GPU，然後即可開始使用 XLA！或者，讓 PJRT 根據您的環境自動選擇裝置。
改進的效能：減少 gRPC 的開銷意味著更快的端到端執行。在 TorchBench 2.0 上，我們觀察到 TPU v4 上的訓練時間縮短了 >35%。
輕鬆的 Pod 執行：只需將您的程式碼複製到每個 TPU 工作節點，然後使用 gcloud compute tpus tpuvm ssh --worker=all 同時執行它們。
更好的擴展性：消除了 XRT 對參數大小的限制，並支援多達 2048 個 TPU 晶片。

快速入門¶

若要開始在 PyTorch/XLA 中使用 PJRT，您只需要設定 PJRT_DEVICE 環境變數。如果您使用的是 TPU v2 或 v3，請继续阅读以了解 TPU v2 和 v3 與 v4 之间的差異。

CPU¶

在安裝了 PyTorch/XLA 的任何機器上，您都可以在 CPU 上執行我們的 MNIST 範例，如下所示

PJRT_DEVICE=CPU python3 xla/test/test_train_mp_mnist.py --fake_data

TPU¶

若要建立安裝了 PyTorch/XLA r1.13 的新 TPU

gcloud alpha compute tpus tpu-vm create $USER-pjrt --accelerator-type=v4-8 --version=tpu-vm-v4-pt-1.13 --zone=us-central2-b --project=$PROJECT

在 v4-8 上，您可以執行我們的 ResNet50 範例，如下所示

git clone --depth=1 --branch r1.13 https://github.com/pytorch/xla.git
PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1

預設情況下，PJRT 將使用所有 TPU 晶片。若要僅使用一個 TPU 晶片，請設定 TPU_PROCESS_BOUNDS 和 TPU_VISIBLE_CHIPS

TPU_PROCESS_BOUNDS=1,1,1 TPU_VISIBLE_CHIPS=0 PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1

Pods¶

在 TPU Pod 上，使用 gcloud 在每個 TPU 上並行執行您的命令

gcloud alpha compute tpus tpu-vm ssh $USER-pjrt --zone=us-central2-b --project=$PROJECT --worker=all --command="git clone --depth=1 --branch r1.13 https://github.com/pytorch/xla.git"
gcloud alpha compute tpus tpu-vm ssh $USER-pjrt --zone=us-central2-b --project=$PROJECT --worker=all --command="PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1"

Docker¶

您也可以使用 Docker 在預先安裝了 PyTorch/XLA 的容器中執行您的工作負載

export DOCKER_IMAGE=gcr.io/...

# Optional: authenticate docker if your image is in a private GCP repository
gcloud compute tpus tpu-vm ssh $USER-pjrt --zone=us-central2-b --project=$PROJECT --worker=all --command "sudo gcloud auth configure-docker"

# Run your workload
gcloud compute tpus tpu-vm ssh $USER-pjrt --zone=us-central2-b --project=$PROJECT --worker=all --command "sudo docker run --rm --privileged --net=host -e PJRT_DEVICE=TPU $DOCKER_IMAGE python pytorch/xla/test/test_train_mp_imagenet.py --fake_data"

請注意，docker run 需要對主機具有特殊權限 (--privileged) 才能將 TPU 裝置公開給容器。目前，TPU Pod 上的 Docker 僅支援主機網路 --net=host。有關更多資訊，請參閱 Cloud TPU 文件。

GPU¶

警告：GPU 支援仍處於高度實驗階段！

若要將 GPU 與 PJRT 搭配使用，只需設定 PJRT_DEVICE=GPU 並將 GPU_NUM_DEVICES 設定為主機上的裝置數量。例如

PJRT_DEVICE=GPU GPU_NUM_DEVICES=4 python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=128 --num_epochs=1

目前，僅支援單一主機，多主機 GPU 叢集支援將在未來版本中新增。

與 XRT 的差異¶

儘管在大多數情況下，我們預計 PJRT 和 XRT 從最終使用者的角度來看可以互換使用（尤其是在 TPU v4 上），但仍有一些細微的差異需要注意。重要的是，XRT 是圍繞 TPU 節點架構設計的，因此即使在 TPU VM 上，它也會始終產生一個客戶端程序和一個伺服器程序。因此，每批輸入都有額外的延遲，因為需要序列化和反序列化資料才能透過網路傳送資料。

PJRT 直接使用本機裝置，沒有中間伺服器程序。在預設設定中，PJRT 將為每個 TPU 晶片建立一個程序，或為每個 TPU 主機建立 4 個程序。有關 TPU 架構的更多資訊，請參閱 Cloud TPU 文件。

對於受限於. 的工作負載，可以提高效能。
在 XRT 下，伺服器程序是唯一與 TPU 裝置互動的程序，用戶端程序無法直接存取 TPU 裝置。在分析單一主機 TPU（例如 v3-8 或 v4-8）時，您通常會看到 8 個裝置追蹤（每個 TPU 核心一個）。使用 PJRT 時，每個程序都有一個晶片，而且來自該程序的分析只會顯示 2 個 TPU 核心。
- 出於相同的原因，分析在使用 XRT 的 TPU Pod 上無法運作，因為伺服器程序的執行獨立於使用者的模型程式碼。PJRT 沒有此限制，因此可以在 TPU Pod 中的每個程序分析 2 個 TPU 核心。
PJRT 只支援 TPU VM 架構，我們沒有計畫使用 PJRT 支援 TPU 節點架構。
使用 PJRT 時，執行階段設定會簡單得多。執行 TPU Pod 工作負載不需要 xla_dist。您只需將程式碼複製到每個 TPU 主機（[gcloud compute tpus tpu-vm scp](https://cloud.google.com/sdk/gcloud/reference/alpha/compute/tpus/tpu-vm/scp)），並在每個主機上平行執行程式碼（例如 [gcloud compute tpus tpu-vm ssh --workers=all --command="PJRT_DEVICE=TPU python run.py"](https://cloud.google.com/sdk/gcloud/reference/alpha/compute/tpus/tpu-vm/ssh)）。
已使用 XLA 原生集體通訊重新實作 xm.rendezvous，以增強大型 TPU Pod 的穩定性。詳情請參閱下方說明。

TPU v2/v3 上的多執行緒處理¶

在 TPU v2 和 v3 上，**分散式工作負載一律會以多執行緒執行**，因為每個 TPU 核心都會將兩個 TPU 核心公開為裝置，而且一次只有一個程序可以開啟 TPU 晶片。在其預設配置中，xmp.spawn 會自動產生盡可能多的程序（每個 TPU 主機 4 個），並為每個程序建立兩個執行緒（每個 TPU 核心一個）。

注意：在 TPU v4 上，每個 TPU 晶片都表示為一個 PyTorch 裝置，因此分散式工作負載將跨 4 個程序執行，每個程序只有一個執行緒。這與 XRT 的行為相同。

在大多數情況下，這不需要對現有程式碼進行大量變更。在大多數情況下，您必須進行的主要變更是模型初始化。因為 torch 的全域 RNG 會在執行緒之間共用，即使您在每個副本中將 torch.manual_seed 設定為相同的值，結果也會因執行緒和執行而異。若要讓副本之間的參數保持一致，請使用 torch_xla.experimental.pjrt.broadcast_master_param 將一個副本的參數廣播到所有其他副本，或從通用檢查點載入每個副本的參數。

xm.rendezvous 的變更¶

PyTorch/XLA r2.0 的新功能

使用 XRT 時，工作節點 0 會執行網格主服務，而所有工作節點上的所有程序都會透過 gRPC 連線到該服務。在實務上，我們發現由於工作節點 0 的輸入連線數量眾多，因此在擁有數千個晶片的 TPU Pod 上執行單一網格主程序並不可靠。單一用戶端程序逾時可能會導致失敗，並強制整個工作負載重新啟動。

因此，我們已使用原生 XLA 集體通訊重新實作 xm.rendezvous，這在大型 TPU Pod 上更加穩定且經過充分測試。與 XRT 實作相比，這帶來了兩個新的限制

因為酬載必須成為 XLA 圖形的一部分，所以會在傳輸資料之前和之後呼叫 xm.mark_step。在模型程式碼中間呼叫 xm.rendezvous可能會強制進行不必要的編譯。
因為 XLA 不允許在工作節點子集上執行集體作業，所以所有工作節點都必須參與 rendezvous。

如果您需要 xm.rendezvous 的舊行為（即在不改變 XLA 圖形的情況下通訊資料和/或同步工作節點子集），請考慮使用 ``torch.distributed.barrier` <https://pytorch.com.tw/docs/stable/distributed.html#torch.distributed.barrier>`_ 或使用 [torch.distributed.all_gather_object](https://pytorch.com.tw/docs/stable/distributed.html#torch.distributed.all_gather_object) 的 gloo 程序群組。如果您也使用 xla torch.distributed 後端，則可以使用 torch.new_group 建立 gloo 子群組。請參閱 PyTorch 文件中的這個範例。請記住這些限制

torch.distributed 在 TPU v2/v3 上不完全支援。只實作了 xla 後端的部分作業，而 gloo 在多程序環境中可能無法如預期般運作。
在我們的實驗中，gloo 無法擴展到數千個 TPU 晶片，因此預期這種替代方案的可靠性不如使用 PJRT 的 xm.rendezvous。

PJRT 和 torch.distributed¶

PyTorch/XLA r2.0 的新功能

當將 PJRT 與 torch.distributed 和 [torch.nn.parallel.DistributedDataParallel](https://github.com/pytorch/xla/blob/master/docs/ddp.md) 搭配使用時，我們強烈建議使用新的 pjrt:// init_method，它會透過查詢執行階段自動找到副本 ID、世界大小和主 IP。例如：

import torch
import torch.distributed as dist
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp
from torch_xla.experimental import pjrt

# Required for `pjrt://` init_method
import torch_xla.experimental.pjrt_backend

def _all_gather(index: int):
  # No need to pass in `rank` or `world_size`
  dist.init_process_group('xla', init_method='pjrt://')

  t = torch.tensor([index], dtype=torch.int32, device=xm.xla_device())
  output = [torch.zeros_like(t) for _ in range(dist.get_world_size())]
  dist.all_gather(output, t)

  xm.mark_step()
  print(output)

if __name__ == '__main__':
  xmp.spawn(_all_gather)

注意：雖然 TPU v4 並不要求使用 pjrt:// init_method，但我們仍然建議使用。如果您使用 env://，則必須將 MASTER_ADDR 設定為擁有裝置 0 的 IP 主機，而該主機*並非*一律是工作節點 0。 pjrt:// init_method 會自動找到這個 IP，並支援 TPU v2/v3。

如需在 PyTorch/XLA 上使用 DistributedDataParallel 的詳細資訊，請參閱 TPU V4 上的 ``ddp.md` <./ddp.md>`_。如需同時使用 DDP 和 PJRT 的範例，請在 TPU 上執行下列範例指令碼

PJRT_DEVICE=TPU python xla/test/test_train_mp_mnist.py --ddp --pjrt_distributed --fake_data --num_epochs 1

效能¶

TorchBench 顯示，與 XRT 相比，PJRT 在各項任務中的平均訓練時間都有所改善，在 TPU v4-8 上的平均改善幅度超過 35%。改善幅度因任務和模型類型而異，從 0% 到 175% 不等。下圖顯示了按任務劃分的明細

新的 TPU 執行階段¶

PyTorch/XLA r2.0 的新功能

PyTorch/XLA r2.0 版本引入了對PJRT 外掛 API的支援，該 API 用於存取 libtpu 中基於 TFRT 的新 TPU 執行階段。現在，當設定 PJRT_DEVICE=TPU 時，這是預設的執行階段。在 2.0 版本中，1.13 版中使用的基於 StreamExecutor 的舊版 TPU 執行階段仍然可以使用 PJRT_DEVICE=TPU_LEGACY，但它將在未來的版本中被移除。如果您遇到只發生在 TPU 上而沒有發生在 TPU_LEGACY 上的問題，請在 GitHub 上提交問題。

在大多數情況下，我們預計這兩個執行階段的效能會相似，但在某些情況下，新的執行階段可能會快達 30%。下圖顯示了按任務劃分的明細

注意：此圖表中顯示的改進也包含在 PJRT 與 XRT 的比較中。

PyTorch XLA 中的 TorchDynamo(torch.compile) 整合¶

Torchdynamo 是一個 Python 層級的 JIT 編譯器，旨在讓未修改的 PyTorch 程式更快。它為編譯器後端提供了一個乾淨的 API 來掛鉤，其最大的特點是在執行之前動態修改 Python 位元組碼。在 pytorch/xla 2.0 版本中，PyTorch/XLA 為 TorchDynamo 提供了一個實驗性的後端，用於推理和訓練。

XLA 橋接的工作方式是，當 Dynamo 識別出模型模式時，它將提供一個 TorchFX 圖形，而 PyTorch/XLA 將使用現有的 Lazy Tensor 技術來編譯 FX 圖形並返回編譯後的函數。

推理¶

以下是一個使用 torch.compile 執行 resnet18 的小程式碼範例

import torch
imprt torchvision
import torch_xla.core.xla_model as xm

def eval_model(loader):
  device = xm.xla_device()
  xla_resnet18 = torchvision.models.resnet18().to(device)
  xla_resnet18.eval()
  dynamo_resnet18 = torch.compile(
      xla_resnet18, backend='torchxla_trace_once')
  for data, _ in loader:
    output = dynamo_resnet18(data)

**注意：**推理後端名稱 torchxla_trace_once 可能會更改。

使用 torch.compile 時，您會看到 PyTorch/XLA 只在初始化時追蹤一次 resent18 模型，並且每次調用 dynamo_resnet18 時都執行編譯後的二進制文件，而不是每次都追蹤模型。請注意，目前 Dynamo 不支援回退，因此如果有一個操作無法被 XLA 追蹤，它就會出錯。我們將在即將發布的 2.1 版本中修復這個問題。以下是在 Cloud TPU v4-8 上使用 torch bench 比較 Dynamo 和 Lazy 的推理速度分析

訓練¶

PyTorch/XLA 也支援 Dynamo 進行訓練，但它還處於實驗階段，我們正在與 PyTorch 編譯器團隊合作迭代實作。在 2.0 版本中，它只支援正向和反向傳遞，但不支援優化器。以下是一個使用 torch.compile 訓練 resnet18 的範例

import torch
imprt torchvision
import torch_xla.core.xla_model as xm

def train_model(model, data, target):
  loss_fn = torch.nn.CrossEntropyLoss()
  pred = model(data)
  loss = loss_fn(pred, target)
  loss.backward()
  return pred

def train_model_main(loader):
  device = xm.xla_device()
  xla_resnet18 = torchvision.models.resnet18().to(device)
  xla_resnet18.train()
  dynamo_train_model = torch.compile(
        train_model, backend='aot_torchxla_trace_once')
  for data, target in loader:
    output = dynamo_train_model(xla_resnet18, data, target)

**注意：**我們這裡使用的後端是 aot_torchxla_trace_once（可能會更改），而不是 torchxla_trace_once

我們預計每個訓練步驟會提取並執行 3 個圖形，而不是像使用 Lazy Tensor 那樣只執行一個訓練步驟。以下是在 Cloud TPU v4-8 上使用 torch bench 比較 Dynamo 和 Lazy 的訓練速度分析。

**注意：**我們針對每個模型的正向和反向傳遞運行一個步驟，然後收集端到端時間。在現實世界中，我們會在每個訓練作業中運行多個步驟，這可以很容易地隱藏執行中的追蹤成本（因為它是異步的）。在這種情況下，Lazy Tensor 的性能會好得多。

我們目前正在開發對優化器的支援，它將很快在 nightly 版本中提供，但不會在 2.0 版本中提供。

結論¶

TorchDynamo 為編譯器後端提供了一種非常有前景的方式，可以對用戶隱藏複雜性，並以圖形格式輕鬆檢索建模代碼。與 PyTorch/XLA 傳統的 Lazy Tensor 提取圖形的方式相比，TorchDynamo 可以跳過每次迭代的圖形追蹤，因此提供更好的推理響應時間。然而，TorchDynamo 還沒有追蹤通訊操作（如 all_reduce 和 all_gather），並且它為正向和反向提供了單獨的圖形，這損害了 xla 的效能。與 Lazy Tensor 相比，這些功能差距使其在實際訓練用例中的效率較低，特別是在訓練中追蹤成本可以與執行重疊。PyTorch/XLA 團隊將繼續投資 TorchDynamo，並與上游合作，以使訓練故事更加成熟。

PyTorch XLA 中的全分片數據並行 (FSDP)¶

PyTorch XLA 中的全分片數據並行 (FSDP) 是一個用於在數據並行工作器之間分片模組參數的工具。

使用範例

import torch
import torch_xla.core.xla_model as xm
from torch_xla.distributed.fsdp import XlaFullyShardedDataParallel as FSDP

model = FSDP(my_module)
optim = torch.optim.Adam(model.parameters(), lr=0.0001)
output = model(x, y)
loss = output.sum()
loss.backward()
optim.step()

也可以單獨分片各層，並使用外部包裝器處理任何剩餘參數。

注意事項

XlaFullyShardedDataParallel 類別支援 https://arxiv.org/abs/1910.02054 中的 ZeRO-2 優化器（分片梯度和優化器狀態）和 ZeRO-3 優化器（分片參數、梯度和優化器狀態）。
- ZeRO-3 優化器應通過嵌套 FSDP 實現，並使用 reshard_after_forward=True。範例請參閱 test/test_train_mp_mnist_fsdp_with_ckpt.py 和 test/test_train_mp_imagenet_fsdp.py。
- 對於無法放入單個 TPU 內存或主機 CPU 內存的大型模型，應該將子模組構造與內部 FSDP 包裝交織在一起。範例請參閱 ``FSDPViTModel` <https://github.com/ronghanghu/vit_10b_fsdp_example/blob/master/run_vit_training.py>`_。
提供了一個簡單的包裝器 checkpoint_module（基於 https://github.com/pytorch/xla/pull/3524 中的 torch_xla.utils.checkpoint.checkpoint），用於對給定的 nn.Module 實例執行梯度檢查點。範例請參閱 test/test_train_mp_mnist_fsdp_with_ckpt.py 和 test/test_train_mp_imagenet_fsdp.py。
自動包裝子模組：除了手動嵌套 FSDP 包裝之外，還可以指定 auto_wrap_policy 參數以使用內部 FSDP 自動包裝子模組。torch_xla.distributed.fsdp.wrap 中的 size_based_auto_wrap_policy 是 auto_wrap_policy 可調用對象的一個範例，此策略包裝參數數量大於 1 億的層。torch_xla.distributed.fsdp.wrap 中的 transformer_auto_wrap_policy 是適用於類似變壓器的模型架構的 auto_wrap_policy 可調用對象的一個範例。

例如，要使用內部 FSDP 自動包裝所有 torch.nn.Conv2d 子模組，可以使用

from torch_xla.distributed.fsdp.wrap import transformer_auto_wrap_policy
auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={torch.nn.Conv2d})

此外，還可以指定 auto_wrapper_callable 參數，以便對子模組使用自定義可調用包裝器（默認包裝器只是 XlaFullyShardedDataParallel 類別本身）。例如，可以使用以下內容對每個自動包裝的子模組應用梯度檢查點（即激活檢查點/重新實例化）。

from torch_xla.distributed.fsdp import checkpoint_module
auto_wrapper_callable = lambda m, *args, **kwargs: XlaFullyShardedDataParallel(
    checkpoint_module(m), *args, **kwargs)

當更新優化器時，直接調用 optimizer.step，而不要調用 xm.optimizer_step。後者會減少跨等級的梯度，而 FSDP（其中參數已經分片）不需要這樣做。
在訓練期間保存模型和優化器檢查點時，每個訓練進程都需要保存其自己的（分片的）模型和優化器狀態字典的檢查點（使用 master_only=False 並為 xm.save 中的每個等級設置不同的路徑）。恢復時，它需要加載相應等級的檢查點。
還請將 model.get_shard_metadata() 與 model.state_dict() 一起保存，如下所示，並使用 consolidate_sharded_model_checkpoints 將分片的模型檢查點拼接在一起，形成一個完整的模型狀態字典。範例請參閱 test/test_train_mp_mnist_fsdp_with_ckpt.py。 .. code-block:: python3

ckpt = {
‘model’: model.state_dict(), ‘shard_metadata’: model.get_shard_metadata(), ‘optimizer’: optimizer.state_dict(),

} ckpt_path = f’/tmp/rank-{xm.get_ordinal()}-of-{xm.xrt_world_size()}.pth’ xm.save(ckpt, ckpt_path, master_only=False)
檢查點合併腳本也可以從命令行啟動，如下所示。 .. code-block:: bash

# 通過命令行工具合併已保存的檢查點 python3 -m torch_xla.distributed.fsdp.consolidate_sharded_ckpts –ckpt_prefix /path/to/your_sharded_checkpoint_files –ckpt_suffix “_rank--of-.pth”

這個類別的實現很大程度上受到了 https://fairscale.readthedocs.io/en/stable/api/nn/fsdp.html 中的 fairscale.nn.FullyShardedDataParallel 的啟發，並且大多遵循其結構。與 fairscale.nn.FullyShardedDataParallel 的最大區別之一是，在 XLA 中，我們沒有明確的參數存儲，因此我們採用了一種不同的方法來釋放 ZeRO-3 的完整參數。

MNIST 和 ImageNet 上的訓練腳本範例¶

MNIST：``test/test_train_mp_mnist_fsdp_with_ckpt.py` <https://github.com/pytorch/xla/blob/master/test/test_train_mp_mnist_fsdp_with_ckpt.py>`_（它還測試檢查點合併）
ImageNet：``test/test_train_mp_imagenet_fsdp.py` <https://github.com/pytorch/xla/blob/master/test/test_train_mp_imagenet_fsdp.py>`_

安裝¶

FSDP 可在 PyTorch/XLA 1.12 版和更新的每晚版本中使用。安裝指南請參閱 https://github.com/pytorch/xla#-available-images-and-wheels。

複製 PyTorch/XLA 儲存庫¶

git clone --recursive https://github.com/pytorch/pytorch
cd pytorch/
git clone --recursive https://github.com/pytorch/xla.git
cd ~/

在 v3-8 TPU 上訓練 MNIST¶

2 個時期後，準確率約為 98.9%

python3 ~/pytorch/xla/test/test_train_mp_mnist_fsdp_with_ckpt.py \
  --batch_size 16 --drop_last --num_epochs 2 \
  --use_nested_fsdp --use_gradient_checkpointing

此腳本會在最後自動測試檢查點合併。您也可以通過以下方式手動合併分片的檢查點

# consolidate the saved checkpoints via command line tool
python3 -m torch_xla.distributed.fsdp.consolidate_sharded_ckpts \
  --ckpt_prefix /tmp/mnist-fsdp/final_ckpt \
  --ckpt_suffix "_rank-*-of-*.pth"

在 v3-8 TPU 上使用 ResNet-50 訓練 ImageNet¶

100 個時期後，準確率約為 75.9%；將 ImageNet-1k 下載到 /datasets/imagenet-1k

python3 ~/pytorch/xla/test/test_train_mp_imagenet_fsdp.py \
  --datadir /datasets/imagenet-1k --drop_last \
  --model resnet50 --test_set_batch_size 64 --eval_interval 10 \
  --lr 0.4 --batch_size 128 --num_warmup_epochs 5 --lr_scheduler_divide_every_n_epochs 30 --lr_scheduler_divisor 10 --num_epochs 100 \
  --use_nested_fsdp

您還可以添加 --use_gradient_checkpointing（需要與 --use_nested_fsdp 或 --auto_wrap_policy 一起使用）以對殘差塊應用梯度檢查點。

TPU pod 上的訓練腳本範例（具有 100 億個參數）¶

要訓練無法放入單個 TPU 的大型模型，在構建整個模型以實現 ZeRO-3 演算法時，應應用自動包裝或使用內部 FSDP 手動包裝子模組。

有關使用此 XLA FSDP PR 對 Vision Transformer (ViT) 模型進行分片訓練的範例，請參閱 https://github.com/ronghanghu/vit_10b_fsdp_example。

如何執行 `DistributedDataParallel`¶

本文檔說明如何在 xla 中使用 torch.nn.parallel.DistributedDataParallel，並進一步說明其與原生 xla 數據並行方法的區別。

背景/動機¶

客戶長期以來一直要求能夠將 PyTorch 的 DistributedDataParallel API 與 xla 一起使用。在這裡，我們將其作為一個實驗性功能啟用。

如何使用 DistributedDataParallel¶

對於那些從 PyTorch 急切模式切換到 XLA 的人來說，以下是將急切 DDP 模型轉換為 XLA 模型所需進行的所有更改。我們假設您已經知道如何在單個設備上使用 XLA。

導入 xla 特定的分佈式套件

import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_backend

初始化 xla 進程組，類似於其他進程組，例如 nccl 和 gloo。

dist.init_process_group("xla", rank=rank, world_size=world_size)

如果需要，請使用 xla 特定的 API 獲取等級和 world_size。

new_rank = xm.get_ordinal()
world_size = xm.xrt_world_size()

將 gradient_as_bucket_view=True 傳遞給 DDP 包裝器。

ddp_model = DDP(model, gradient_as_bucket_view=True)

最後使用 xla 特定的啟動器啟動您的模型。

xmp.spawn(demo_fn)

在這裡，我們將所有內容放在一起（該範例實際上取自 DDP 教程）。您的編碼方式與急切體驗非常相似。只需在單個設備上進行 xla 特定的調整，再加上對腳本的上述五項更改即可。

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

# additional imports for xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_backend
import torch_xla.distributed.xla_multiprocessing as xmp

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the xla process group
    dist.init_process_group("xla", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 1000000)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(1000000, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

def demo_basic(rank):
    # xla specific APIs to get rank, world_size.
    new_rank = xm.get_ordinal()
    assert new_rank == rank
    world_size = xm.xrt_world_size()

    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to XLA device
    device = xm.xla_device()
    model = ToyModel().to(device)
    # currently, graident_as_bucket_view is needed to make DDP work for xla
    ddp_model = DDP(model, gradient_as_bucket_view=True)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10).to(device))
    labels = torch.randn(20, 5).to(device)
    loss_fn(outputs, labels).backward()
    optimizer.step()
    # xla specific API to execute the graph
    xm.mark_step()

    cleanup()


def run_demo(demo_fn):
    # xla specific launcher
    xmp.spawn(demo_fn)

if __name__ == "__main__":
    run_demo(demo_basic)

基準測試¶

Resnet50 與虛擬數據¶

以下結果是使用命令：python test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1 在具有 ToT PyTorch 和 PyTorch/XLA 的 TPU VM V3-8 環境中收集的。統計指標是通過使用此拉取請求中的腳本生成的。速率單位為每秒圖像數。

類型	平均值	中位數	第 90 個百分位數	標準差	變異係數
xm.optimizer_step	418.54	419.22	430.40	9.76	0.02
DDP	395.97	395.54	407.13	7.60	0.02

我們的原生分佈式數據並行方法與 DistributedDataParallel 包裝器之間的效能差異為：1 - 395.97 / 418.54 = 5.39%。鑑於 DDP 包裝器在追蹤 DDP 運行時引入了額外的開銷，因此該結果似乎是合理的。

MNIST 與虛擬數據¶

以下結果是使用指令： python test/test_train_mp_mnist.py --fake_data 在搭載最新版 PyTorch 和 PyTorch/XLA 的 TPU VM V3-8 環境中收集而得。統計指標則是用這個 pull request 中的腳本產生。速率單位為每秒影像數。

類型	平均值	中位數	第 90 個百分位數	標準差	變異係數
xm.optimizer_step	17864.19	20108.96	24351.74	5866.83	0.33
DDP	10701.39	11770.00	14313.78	3102.92	0.29

我們原生分散式資料平行處理方法和 DistributedDataParallel 包裝器之間的效能差異為：1 - 14313.78 / 24351.74 = 41.22%。這裡我們比較的是第 90 個百分位數，因為資料集很小，而且前幾輪會受到資料載入的嚴重影響。這種速度減慢幅度很大，但考慮到模型很小，這是合理的。額外的 DDP 執行時間追蹤成本難以攤銷。

使用真實資料的 MNIST¶

以下結果是使用指令： python test/test_train_mp_mnist.py --logdir mnist/ 在搭載最新版 PyTorch 和 PyTorch/XLA 的 TPU VM V3-8 環境中收集而得。

我們可以觀察到，即使 DDP 包裝器最終仍能達到 97.48% 的高準確率，但其收斂速度比原生 XLA 方法慢。（原生方法可達到 99%。）

免責聲明¶

此功能仍處於實驗階段，並且正在積極開發中。請謹慎使用，並歡迎將任何錯誤回報至 xla github 存放庫。對於那些對原生 xla 資料平行處理方法感興趣的人，這裡有一份教學。

以下是一些已知問題，目前正在調查中

需要強制執行 gradient_as_bucket_view=True。
與 torch.utils.data.DataLoader 一起使用時會出現一些問題。test_train_mp_mnist.py 在使用真實資料時會在結束前當機。

如何使用 PyTorch/XLA:GPU 執行¶

PyTorch/XLA 讓 PyTorch 使用者能夠利用支援 TPU、GPU、CPU 等加速器的 XLA 編譯器。本文件將說明在 nvidia gpu 執行個體上執行 PyTorch/XLA 的基本步驟

建立 GPU 執行個體¶

Pytorch/XLA 目前發布了使用 cuda11.2 和 python 3.7/3.8 的預先建置 Docker 映像檔和 wheel 檔案。我們建議使用者建立具有對應設定的 GPU 執行個體。如需 Docker 映像檔和 wheel 檔案的完整清單，請參閱本文件。

環境設定¶

Docker¶

sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_cuda_11.2
sudo apt-get install -y apt-transport-https ca-certificates curl gnupg-agent    software-properties-common
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker
sudo docker run --gpus all -it -d gcr.io/tpu-pytorch/xla:nightly_3.7\8_cuda_11.2 bin/bash
sudo docker exec -it $(sudo docker ps | awk 'NR==2 { print $1 }') /bin/bash

請注意，您需要重新啟動 Docker 才能在 Docker 容器中顯示 gpu 裝置。登入 Docker 後，您可以使用 nvidia-smi 來確認裝置是否已正確設定。

(pytorch) root@20ab2c7a2d06:/# nvidia-smi
Thu Dec  8 06:24:29 2022
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000000:00:04.0 Off |                    0 |
| N/A   36C    P0    38W / 300W |      0MiB / 16384MiB |      1%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

Wheel¶

pip3 install torch=1.13
pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/cuda/112/torch_xla-1.13-cp37-cp37m-linux_x86_64.whl

執行簡單模型¶

為了執行以下範例，您需要複製 pytorch/xla 存放庫才能存取 imagenet 範例（我們已經在 Docker 中複製了它）。

(pytorch) root@20ab2c7a2d06:/# export GPU_NUM_DEVICES=1
(pytorch) root@20ab2c7a2d06:/# python pytorch/xla/test/test_train_mp_imagenet.py --fake_data
==> Preparing data..
Epoch 1 train begin 06:12:38
2022-12-08 06:13:12.452874: W      79 tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc:729] None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.  Conv: (f32[128,256,28,28]{3,2,1,0}, u8[0]{0}) custom-call(f32[128,256,14,14]{3,2,1,0}, f32[3,3,256,256]{1,0,2,3}), window={size=3x3 stride=2x2 pad=1_1x1_1}, dim_labels=bf01_01io->bf01, custom_call_target="__cudnn$convBackwardInput", backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
2022-12-08 06:13:13.780992: W      79 tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc:729] None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.  Conv: (f32[128,128,56,56]{3,2,1,0}, u8[0]{0}) custom-call(f32[128,128,28,28]{3,2,1,0}, f32[3,3,128,128]{1,0,2,3}), window={size=3x3 stride=2x2 pad=1_1x1_1}, dim_labels=bf01_01io->bf01, custom_call_target="__cudnn$convBackwardInput", backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
| Training Device=xla:0/0 Epoch=1 Step=0 Loss=6.89059 Rate=2.82 GlobalRate=2.82 Time=06:13:23
| Training Device=xla:0/0 Epoch=1 Step=20 Loss=6.79297 Rate=117.16 GlobalRate=45.84 Time=06:13:36
| Training Device=xla:0/0 Epoch=1 Step=40 Loss=6.43628 Rate=281.16 GlobalRate=80.49 Time=06:13:43
| Training Device=xla:0/0 Epoch=1 Step=60 Loss=5.83108 Rate=346.88 GlobalRate=108.82 Time=06:13:49
| Training Device=xla:0/0 Epoch=1 Step=80 Loss=4.99023 Rate=373.62 GlobalRate=132.43 Time=06:13:56
| Training Device=xla:0/0 Epoch=1 Step=100 Loss=3.92699 Rate=384.33 GlobalRate=152.40 Time=06:14:02
| Training Device=xla:0/0 Epoch=1 Step=120 Loss=2.68816 Rate=388.35 GlobalRate=169.49 Time=06:14:09

AMP（自動混合精度）¶

AMP 在 GPU 訓練中非常有用，而 PyTorch/XLA 重用了 Cuda 的 AMP 規則。您可以查看我們的 mnist 範例和 imagenet 範例。請注意，我們還使用了修改版的優化器，以避免裝置和主機之間的額外同步。

XLA 設備上的 PyTorch¶

建立 XLA 張量¶

XLA 張量是 PyTorch 張量¶

在 XLA 設備上執行模型¶

在單一 XLA 設備上執行¶

使用多程序在多個 XLA 設備上執行¶

在 TPU Pod 上執行¶

XLA 張量深入探討¶

XLA 張量是惰性的¶

XLA 張量和 bFloat16¶

内存佈局¶

在 CPU 和 XLA 設備之間移動 XLA 張量¶

保存和加載 XLA 張量¶

進一步閱讀¶

PyTorch/XLA API¶

xla_model¶

distributed¶

utils¶

test¶

疑難排解¶

執行 A 自動指標分析¶

獲取指標報告¶

了解指標報告¶

清除指標報告¶

效能分析¶

已知的效能注意事項¶

XLA 張量怪癖¶

更多除錯工具¶

環境變數¶

擷取堆疊追蹤¶

使用 debug_run.py 收集除錯資訊¶

常見問題¶

PJRT 執行階段（測試版）¶

TL;DR¶

優點¶

快速入門¶

CPU¶

TPU¶

Pods¶

Docker¶

GPU¶

與 XRT 的差異¶

TPU v2/v3 上的多執行緒處理¶

xm.rendezvous 的變更¶

PJRT 和 torch.distributed¶

效能¶

新的 TPU 執行階段¶

PyTorch XLA 中的 TorchDynamo(torch.compile) 整合¶

推理¶

訓練¶

結論¶

PyTorch XLA 中的全分片數據並行 (FSDP)¶

MNIST 和 ImageNet 上的訓練腳本範例¶

安裝¶

複製 PyTorch/XLA 儲存庫¶

在 v3-8 TPU 上訓練 MNIST¶

在 v3-8 TPU 上使用 ResNet-50 訓練 ImageNet¶

TPU pod 上的訓練腳本範例（具有 100 億個參數）¶

如何執行 DistributedDataParallel¶

背景/動機¶

如何使用 DistributedDataParallel¶

基準測試¶

Resnet50 與虛擬數據¶

MNIST 與虛擬數據¶

使用真實資料的 MNIST¶

免責聲明¶

如何使用 PyTorch/XLA:GPU 執行¶

建立 GPU 執行個體¶

環境設定¶

Docker¶

Wheel¶

執行簡單模型¶

AMP（自動混合精度）¶

文件

教學課程

資源

如何執行 `DistributedDataParallel`¶