PyTorch on XLA Devices¶

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 張量¶

可以在 XLA 張量上執行 PyTorch 運算，就像 CPU 或 CUDA 張量一樣。

例如，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
import torch_xla.core.xla_model as xm
import torch_xla.distributed.parallel_loader as pl

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__':
  torch_xla.launch(_mp_fn, args=())

此多裝置程式碼片段與先前的單一裝置程式碼片段之間有三個差異。讓我們逐一檢視。

torch_xla.launch()
- 建立每個執行 XLA 裝置的程序。
- 此函式是多執行緒衍生 (multithreading spawn) 的包裝函式，可讓使用者也使用 torchrun 命令列執行指令碼。每個程序將只能存取指派給目前程序的裝置。例如，在 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，並傳回正在減少的損失。

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

注意： 請務必注意，當使用多重處理時，使用者只能從 torch_xla.launch() 的目標函式 (或任何將 torch_xla.launch() 作為呼叫堆疊中父項的函式) 內開始擷取和存取 XLA 裝置。

如需關於透過多重處理在多個 XLA 裝置上訓練網路的詳細資訊，請參閱完整的多重處理範例。

在 TPU Pod 上執行¶

不同加速器的多主機設定可能差異很大。本文件將討論多主機訓練中與裝置無關的部分，並以 TPU + PJRT 執行階段 (目前在 1.13 和 2.x 版本中提供) 作為範例。

在您開始之前，請先查看我們的使用者指南這裡，其中將說明一些 Google Cloud 基礎知識，例如如何使用 gcloud 命令以及如何設定您的專案。您也可以查看這裡，以取得所有 Cloud TPU Howto。本文件將著重於 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 上使用 fakedata 執行 resnet50 的完整指南，請參閱此指南。

XLA 張量深入探討¶

使用 XLA 張量和裝置只需要變更幾行程式碼。但即使 XLA 張量的行為與 CPU 和 CUDA 張量非常相似，它們的內部結構卻有所不同。本節說明 XLA 張量的獨特之處。

XLA 張量是延遲的¶

CPU 和 CUDA 張量會立即或主動啟動運算。另一方面，XLA 張量是延遲的。它們會在圖形中記錄運算，直到需要結果為止。像這樣延遲執行可讓 XLA 最佳化它。例如，多個個別運算的圖形可能會融合為單一最佳化運算。

延遲執行通常對呼叫者而言是不可見的。PyTorch/XLA 會自動建構圖形、將其傳送至 XLA 裝置，並在 XLA 裝置與 CPU 之間複製資料時進行同步處理。在執行最佳化器步驟時插入屏障會明確地同步處理 CPU 和 XLA 裝置。如需關於我們的延遲張量設計的更多資訊，您可以閱讀本論文。

記憶體配置¶

XLA 張量的內部資料表示對使用者而言是不透明的。它們不會公開其儲存空間，而且與 CPU 和 CUDA 張量不同，它們始終顯示為連續的。這讓 XLA 可以調整張量的記憶體配置，以獲得更佳的效能。

將 XLA 張量移至 CPU 和從 CPU 移出¶

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 一樣，正處於積極開發中，此行為未來可能會變更。

編譯快取¶

XLA 編譯器會將追蹤的 HLO 轉換為可在裝置上執行的可執行檔。編譯可能很耗時，而且在 HLO 在執行之間沒有變更的情況下，編譯結果可以保存到磁碟以供重複使用，從而大幅縮短開發迭代時間。

請注意，如果 HLO 在執行之間變更，仍會發生重新編譯。

這目前是一個實驗性的選擇加入 API，必須在執行任何運算之前啟用。初始化是透過 initialize_cache API 完成的

import torch_xla.runtime as xr
xr.initialize_cache('YOUR_CACHE_PATH', readonly=False)

這將在指定的路徑初始化持續性編譯快取。readonly 參數可用於控制工作站是否能夠寫入快取，這在共用快取掛載用於 SPMD 工作負載時可能很有用。

如果您想要在多程序訓練 (使用 torch_xla.launch 或 xmp.spawn) 中使用持續性編譯快取，您應該為不同的程序使用不同的路徑。

def _mp_fn(index):
  # cache init needs to happens inside the mp_fn.
  xr.initialize_cache(f'/tmp/xla_cache_{index}', readonly=False)
  ....

if __name__ == '__main__':
  torch_xla.launch(_mp_fn, args=())

如果您無法存取 index，您可以使用 xr.global_ordinal()。請查看這裡中的可執行範例。

延伸閱讀¶

其他文件可在 PyTorch/XLA 儲存庫中取得。如需在 TPU 上執行網路的更多範例，請參閱這裡。