在 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 需要 --shm-size 16GB --privileged --net=host 才能访问 TPU

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