Google TPUs

在 Docker 容器內安裝 Optimum-TPU

本指南說明如何使用官方 PyTorch/XLA 映像檔在 Docker 容器內執行 Optimum-TPU。

先修條件

開始之前，請確保您具備以下條件：

您的系統已安裝 Docker
具備 TPU 執行個體的存取權限
具備執行特權容器 (privileged containers) 的足夠權限

使用 PyTorch/XLA 基礎映像檔

1. 拉取 Docker 映像檔

首先，設定映像檔 URL 與版本的環境變數

export TPUVM_IMAGE_URL=us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla
export TPUVM_IMAGE_VERSION=r2.5.0_3.10_tpuvm

# Pull the image
docker pull ${TPUVM_IMAGE_URL}:${TPUVM_IMAGE_VERSION}

2. 執行容器

使用存取 TPU 所需的旗標啟動容器

docker run -ti \
    --rm \
    --shm-size 16GB
    --privileged \
    --net=host \
    ${TPUVM_IMAGE_URL}@sha256:${TPUVM_IMAGE_VERSION} \
    bash

Docker 若要存取 TPU，需要 --shm-size 16GB --privileged --net=host 設定

3. 安裝 Optimum-TPU

進入容器後，安裝 Optimum-TPU

pip install optimum-tpu -f https://storage.googleapis.com/libtpu-releases/index.html

驗證

若要驗證您的設定，請執行此簡單測試

python3 -c "import torch_xla.core.xla_model as xm; print(xm.xla_device())"

您應該會看到輸出顯示 XLA 裝置已可用（例如 xla:0）。

下一步

設定好容器後，您可以：

開始使用 Optimum-TPU 訓練模型。請參閱我們的訓練範例章節。
執行推論工作負載。請查看我們的服務指南。