在 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 训练模型。请参阅我们的训练示例部分。
运行推理工作负载。请查看我们的服务指南。