云TPU#
SkyPilot 支持在 Google 的 Cloud TPU 上运行任务，这是一个专门用于机器学习工作负载的硬件加速器。
通过TPU研究云（TRC）免费使用TPU#
鼓励ML研究人员和学生通过TPU研究云（TRC）项目申请免费TPU访问！
一键获取TPUs#
使用一个命令快速获取用于开发的TPU节点：
# Use latest TPU v6 (Trillium) VMs:
sky launch --gpus tpu-v6e-8
# Use TPU v4 (Titan) VMs:
sky launch --gpus tpu-v4-8
# Preemptible TPUs:
sky launch --gpus tpu-v6e-8 --use-spot
命令完成后，您将进入TPU主机虚拟机，并可以立即开始开发代码。
下面，我们展示了使用SkyPilot的示例：(1) 在TPU虚拟机/集群上训练LLMs，以及(2) 在TPU节点（旧版）上训练MNIST。
TPU架构#
GCP上有两种不同的TPU架构可用：
两者都受到SkyPilot的支持。我们推荐TPU VMs和Pods，它们是GCP鼓励的较新架构。
这两种架构的区别如下。
对于TPU虚拟机/集群，您可以直接SSH到物理连接到TPU设备的“TPU主机”虚拟机。
对于TPU节点，必须单独配置一个用户虚拟机（一个n1实例），以便通过gRPC与无法访问的TPU主机进行通信。
更多详情可以在GCP 文档中找到。
TPU虚拟机/容器#
谷歌最新的TPU v6（Trillium）虚拟机提供了卓越的性能，并且现在得到了SkyPilot的支持。
要使用TPU VMs/Pods，请在任务YAML的resources字段中设置以下内容：
resources:
accelerators: tpu-v6e-8
accelerator_args:
runtime_version: v2-alpha-tpuv6e # optional
accelerators 字段指定了 TPU 类型，而 accelerator_args 字典包含了可选的 tpu_vm 布尔值（默认为 true，表示使用 TPU VM），以及一个可选的 TPU runtime_version 字段。
要查看支持的 TPU 类型，请运行 sky show-gpus。
这是一个完整的任务YAML，它使用Torch XLA在TPU VM上训练一个Llama 3模型。
resources:
accelerators: tpu-v6e-8 # Fill in the accelerator type you want to use
envs:
HF_TOKEN: # fill in your huggingface token
workdir: .
setup: |
pip3 install huggingface_hub
python3 -c "import huggingface_hub; huggingface_hub.login('${HF_TOKEN}')"
# Setup TPU
pip3 install cloud-tpu-client
sudo apt update
sudo apt install -y libopenblas-base
pip3 install --pre torch==2.6.0.dev20240916+cpu torchvision==0.20.0.dev20240916+cpu \
--index-url https://download.pytorch.org/whl/nightly/cpu
pip install "torch_xla[tpu]@https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev20240916-cp310-cp310-linux_x86_64.whl" \
-f https://storage.googleapis.com/libtpu-releases/index.html
pip install torch_xla[pallas] \
-f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
-f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
# Setup runtime for training
git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git
cd transformers
pip3 install -e .
pip3 install datasets evaluate scikit-learn accelerate
run: |
unset LD_PRELOAD
PJRT_DEVICE=TPU XLA_USE_SPMD=1 ENABLE_PJRT_COMPATIBILITY=true \
python3 transformers/examples/pytorch/language-modeling/run_clm.py \
--dataset_name wikitext \
--dataset_config_name wikitext-2-raw-v1 \
--per_device_train_batch_size 16 \
--do_train \
--output_dir /home/$USER/tmp/test-clm \
--overwrite_output_dir \
--config_name /home/$USER/sky_workdir/config-8B.json \
--cache_dir /home/$USER/cache \
--tokenizer_name meta-llama/Meta-Llama-3-8B \
--block_size 8192 \
--optim adafactor \
--save_strategy no \
--logging_strategy no \
--fsdp "full_shard" \
--fsdp_config /home/$USER/sky_workdir/fsdp_config.json \
--torch_dtype bfloat16 \
--dataloader_drop_last yes \
--flash_attention \
--max_steps 20
这个YAML文件位于SkyPilot仓库，或者你可以将其粘贴到本地文件中。
启动它使用：
$ HF_TOKEN=<your-huggingface-token> sky launch train-llama3-8b.yaml -c llama-3-train --env HF_TOKEN
作业完成后，您应该会看到以下输出。
$ sky launch train-llama3-8b.yaml -c llama-3-train
(task, pid=17499) ***** train metrics *****
(task, pid=17499) epoch = 1.1765
(task, pid=17499) total_flos = 109935420GF
(task, pid=17499) train_loss = 10.6011
(task, pid=17499) train_runtime = 0:11:12.77
(task, pid=17499) train_samples = 282
(task, pid=17499) train_samples_per_second = 0.476
(task, pid=17499) train_steps_per_second = 0.03
多主机TPU Pods#
一个TPU Pod是由专用高速网络接口连接的TPU设备集合，用于高性能训练。
要使用TPU Pod，只需更改任务YAML中的accelerators字段（例如，tpu-v6e-8 -> tpu-v6e-32）。
resources:
accelerators: tpu-v6e-32 # Pods have > 8 cores (the last number)
注意
两种TPU架构，TPU VMs和TPU Nodes，都可以与TPU Pods一起使用。下面的示例基于TPU VMs。
要显示所有可用的TPU Pod类型，请运行 sky show-gpus（超过8个核心意味着Pods）：
GOOGLE_TPU AVAILABLE_QUANTITIES
tpu-v6e-8 1
tpu-v6e-32 1
tpu-v6e-128 1
tpu-v6e-256 1
...
创建TPU Pod后，会启动多个主机虚拟机（例如，tpu-v6e-32 带有4个主机虚拟机）。
通常，用户需要SSH到所有主机以准备文件和设置环境，然后在每个主机上启动作业，这是一个繁琐且容易出错的过程。
SkyPilot 自动处理这种复杂性。从您的笔记本电脑上，一个简单的 sky launch 命令将执行：
工作目录/文件挂载同步；以及
在pod的每个主机上执行setup/run命令。
我们可以在TPU Pod上运行相同的Llama 3训练任务，使用以下命令，并对YAML进行轻微修改（--per_device_train_batch_size 从16改为32）：
$ HF_TOKEN=<your-huggingface-token> sky launch -c tpu-pod --gpus tpu-v6e-32 train-llama3-8b.yaml --env HF_TOKEN
你应该看到以下输出。
(head, rank=0, pid=17894) ***** train metrics *****
(head, rank=0, pid=17894) epoch = 2.5
(head, rank=0, pid=17894) total_flos = 219870840GF
(head, rank=0, pid=17894) train_loss = 10.1527
(head, rank=0, pid=17894) train_runtime = 0:11:13.18
(head, rank=0, pid=17894) train_samples = 282
(head, rank=0, pid=17894) train_samples_per_second = 0.951
(head, rank=0, pid=17894) train_steps_per_second = 0.03
(worker1, rank=1, pid=15406, ip=10.164.0.57) ***** train metrics *****
(worker1, rank=1, pid=15406, ip=10.164.0.57) epoch = 2.5
(worker1, rank=1, pid=15406, ip=10.164.0.57) total_flos = 219870840GF
(worker1, rank=1, pid=15406, ip=10.164.0.57) train_loss = 10.1527
(worker1, rank=1, pid=15406, ip=10.164.0.57) train_runtime = 0:11:15.08
(worker1, rank=1, pid=15406, ip=10.164.0.57) train_samples = 282
(worker1, rank=1, pid=15406, ip=10.164.0.57) train_samples_per_second = 0.948
(worker1, rank=1, pid=15406, ip=10.164.0.57) train_steps_per_second = 0.03
(worker2, rank=2, pid=16552, ip=10.164.0.58) ***** train metrics *****
(worker2, rank=2, pid=16552, ip=10.164.0.58) epoch = 2.5
(worker2, rank=2, pid=16552, ip=10.164.0.58) total_flos = 219870840GF
(worker2, rank=2, pid=16552, ip=10.164.0.58) train_loss = 10.1527
(worker2, rank=2, pid=16552, ip=10.164.0.58) train_runtime = 0:11:15.61
(worker2, rank=2, pid=16552, ip=10.164.0.58) train_samples = 282
(worker2, rank=2, pid=16552, ip=10.164.0.58) train_samples_per_second = 0.947
(worker2, rank=2, pid=16552, ip=10.164.0.58) train_steps_per_second = 0.03
(worker3, rank=3, pid=17469, ip=10.164.0.59) ***** train metrics *****
(worker3, rank=3, pid=17469, ip=10.164.0.59) epoch = 2.5
(worker3, rank=3, pid=17469, ip=10.164.0.59) total_flos = 219870840GF
(worker3, rank=3, pid=17469, ip=10.164.0.59) train_loss = 10.1527
(worker3, rank=3, pid=17469, ip=10.164.0.59) train_runtime = 0:11:15.10
(worker3, rank=3, pid=17469, ip=10.164.0.59) train_samples = 282
(worker3, rank=3, pid=17469, ip=10.164.0.59) train_samples_per_second = 0.948
(worker3, rank=3, pid=17469, ip=10.164.0.59) train_steps_per_second = 0.03
要向同一个TPU Pod提交更多作业，请使用sky exec：
$ HF_TOKEN=<your-huggingface-token> sky exec tpu-pod train-llama3-8b.yaml --env HF_TOKEN
你可以在 SkyPilot repo 中找到更多关于在TPUs上服务LLMs的有用示例。
TPU 节点（旧版）#
在TPU节点中，需要配置一个普通的CPU虚拟机（一个n1实例）来与TPU主机/设备进行通信。
要使用TPU节点，请在任务的YAML文件的resources字段中设置以下内容：
resources:
instance_type: n1-highmem-8
accelerators: tpu-v2-8
accelerator_args:
runtime_version: 2.12.0 # optional, TPU runtime version.
tpu_vm: False
上述YAML将n1-highmem-8视为主机，将tpu-v2-8视为TPU节点资源。
您可以修改主机实例类型或TPU类型。
这是一个完整的任务YAML，它使用TensorFlow在TPU节点上运行MNIST训练。
name: mnist-tpu-node
resources:
accelerators: tpu-v2-8
accelerator_args:
runtime_version: 2.12.0 # optional, TPU runtime version.
tpu_vm: False
# TPU node requires loading data from a GCS bucket.
# We use SkyPilot bucket mounting to mount a GCS bucket to /dataset.
file_mounts:
/dataset:
name: mnist-tpu-node
store: gcs
mode: MOUNT
setup: |
git clone https://github.com/tensorflow/models.git
conda activate mnist
if [ $? -eq 0 ]; then
echo 'conda env exists'
else
conda create -n mnist python=3.8 -y
conda activate mnist
pip install tensorflow==2.12.0 tensorflow-datasets tensorflow-model-optimization cloud-tpu-client
fi
run: |
conda activate mnist
cd models/official/legacy/image_classification/
export STORAGE_BUCKET=gs://mnist-tpu-node
export MODEL_DIR=${STORAGE_BUCKET}/mnist
export DATA_DIR=${STORAGE_BUCKET}/data
export PYTHONPATH=/home/gcpuser/sky_workdir/models
python3 mnist_main.py \
--tpu=${TPU_NAME} \
--model_dir=${MODEL_DIR} \
--data_dir=${DATA_DIR} \
--train_epochs=10 \
--distribution_strategy=tpu \
--download
注意
TPU节点需要从GCS存储桶加载数据。上面的file_mounts规范通过使用SkyPilot存储桶挂载来创建一个新存储桶或挂载现有存储桶，从而简化了这一过程。
如果您遇到存储桶Permission denied错误，
请确保存储桶是在与主机虚拟机/TPU节点相同的区域创建的，并且Cloud TPU的IAM权限已正确设置（按照这里的说明操作）。
注意
特殊环境变量 $TPU_NAME 在运行时由 SkyPilot 自动设置，因此可以在 run 命令中使用。
这个YAML文件位于SkyPilot 仓库 (examples/tpu/tpu_node_mnist.yaml)。使用以下命令启动它：
$ sky launch examples/tpu/tpu_node_mnist.yaml -c mycluster
...
(mnist-tpu-node pid=28961) Epoch 9/10
(mnist-tpu-node pid=28961) 58/58 [==============================] - 1s 19ms/step - loss: 0.1181 - sparse_categorical_accuracy: 0.9646 - val_loss: 0.0921 - val_sparse_categorical_accuracy: 0.9719
(mnist-tpu-node pid=28961) Epoch 10/10
(mnist-tpu-node pid=28961) 58/58 [==============================] - 1s 20ms/step - loss: 0.1139 - sparse_categorical_accuracy: 0.9655 - val_loss: 0.0831 - val_sparse_categorical_accuracy: 0.9742
...
(mnist-tpu-node pid=28961) {'accuracy_top_1': 0.9741753339767456, 'eval_loss': 0.0831054300069809, 'loss': 0.11388632655143738, 'training_accuracy_top_1': 0.9654667377471924}