在 Docker 容器中安裝 Optimum-TPU

本指南介紹如何使用官方 PyTorch/XLA 映象在 Docker 容器中執行 Optimum-TPU。

先決條件

開始之前，請確保您已具備：

系統上安裝了 Docker
可以訪問 TPU 例項
足夠的許可權來執行特權容器

使用 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 訓練模型。請參閱我們的訓練示例部分。
執行推理工作負載。請檢視我們的服務指南。