在 Cloud TPU (TF 2.x) 上训练 HuggingFace GPT2

---

如果您不熟悉 Cloud TPU，我们建议您学习 快速入门，了解如何 创建一个 TPU 虚拟机。

本教程介绍了如何在 Cloud TPU 上训练 HuggingFace GPT2 模型。

警告：本教程使用第三方数据集。Google 不对此数据集的有效性或任何其他方面提供任何代理、担保或其他保证。

目标

• 创建 Cloud TPU
• 安装依赖项
• 运行训练作业

费用

在本文档中，您将使用 Google Cloud 的以下收费组件：

• Compute Engine
• Cloud TPU

您可使用价格计算器根据您的预计使用情况来估算费用。 Google Cloud 新用户可能有资格申请免费试用。

准备工作

在开始学习本教程之前，请检查您的 Google Cloud 项目是否已正确设置。

1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

2. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

3. Make sure that billing is enabled for your Google Cloud project.

4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

   Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

   Go to project selector

5. Make sure that billing is enabled for your Google Cloud project.

本演示使用 Google Cloud 的收费组件。请查看 Cloud TPU 价格页面估算您的费用。请务必在使用完您创建的资源以后清理这些资源，以免产生不必要的费用。

重要提示：请在同一区域/可用区设置所有资源（Cloud TPU 虚拟机和 Cloud Storage 存储桶），以减少网络延迟和网络费用。TPU 虚拟机位于特定可用区，即区域内的细分。

使用 Cloud TPU 训练 HuggingFace GPT2

1. 打开一个 Cloud Shell 窗口。

   打开 Cloud Shell

2. 为项目 ID 创建一个环境变量。

   export PROJECT_ID=your-project-id

3. 将 Google Cloud CLI 配置为使用您的 Google Cloud 项目 创建 Cloud TPU 的位置。

   gcloud config set project ${PROJECT_ID}

   当您第一次在新的 Cloud Shell 虚拟机中运行此命令时，系统会显示 Authorize Cloud Shell 页面。点击页面底部的 Authorize 以允许 gcloud 使用您的 Google Cloud 凭据进行 Google Cloud API 调用。

4. 为 Cloud TPU 项目创建服务账号。

   服务账号允许 Cloud TPU 服务访问其他 Google Cloud 服务。

   $ gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID

   该命令将返回以下格式的 Cloud TPU 服务账号：

   [e-mail]
   

创建 Cloud TPU

1. 使用 gcloud 命令创建 Cloud TPU 虚拟机。以下命令 会创建一个 v4-8 TPU。您还可以通过将 --accelerator-type 标志设置为 Pod 切片类型（例如 v4-32）来创建 TPU Podslice。

   $ gcloud compute tpus tpu-vm create hf-gpt2 \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-tf-2.17.0-pjrt

   命令标志说明

   zone

   拟在其中创建 Cloud TPU 的区域。

   accelerator-type

   加速器类型指定要创建的 Cloud TPU 的版本和大小。 如需详细了解每个 TPU 版本支持的加速器类型，请参阅 TPU 版本。

   version

   Cloud TPU 软件版本。

   注意：如果您有多个 Google Cloud 项目，则必须使用 --project 标志指定您要创建 Cloud TPU 的 Google Cloud 项目的 ID。

2. 使用 SSH 连接到 TPU 虚拟机。连接到网络后 您的 Shell 提示符会从 username@projectname 更改为 username@vm-name:

   gcloud compute tpus tpu-vm ssh hf-gpt2 --zone=us-central2-b

   要点 ：从现在开始，前缀 (vm) $ 表示您应该 请在 Compute Engine 虚拟机实例上运行该命令。

安装依赖项

1. 克隆 HuggingFace Transformers 代码库：

   (vm)$ cd /tmp
   (vm)$ git clone https://github.com/huggingface/transformers.git
   (vm)$ cd transformers

2. 安装依赖项：

   (vm)$ pip install .
   (vm)$ pip install -r examples/tensorflow/_tests_requirements.txt
   (vm)$ cd /tmp/transformers/examples/tensorflow/language-modeling
   (vm)$ pip install -r requirements.txt

3. 创建临时目录：

   (vm)$ mkdir /tmp/gpt2-wikitext

4. 创建 TPU 时，如果您将 --version 参数设置为以 -pjrt 时，请设置以下环境变量以启用 PJRT 运行时：

     (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
     (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

运行训练脚本

(vm)$ python3 run_clm.py \
  --model_name_or_path distilgpt2 \
  --max_train_samples 1000 \
  --max_eval_samples 100 \
  --num_train_epochs 1 \
  --output_dir /tmp/gpt2-wikitext \
  --dataset_name wikitext \
  --dataset_config_name wikitext-103-raw-v1

命令标志说明

model_name_or_path

要训练的模型的名称。

max_train_samples

用于训练的样本数量上限。

max_eval_samples

用于评估的最大样本数。

num_train_epochs

训练模型的周期数。

output_dir

训练脚本的输出目录。

dataset_name

要使用的相应数据集的名称。

dataset_config_name

数据集配置名称

训练完成后，系统会显示如下所示的消息：

  125/125 [============================] - ETA: 0s - loss: 3.61762023-07-07 21:38:17.902850: W tensorflow/core/framework/dataset.cc:956] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
  125/125 [============================] - 763s 6s/step - loss: 3.6176 - val_loss: 3.4233
  Configuration saved in /tmp/gpt2-wikitext/config.json
  Configuration saved in /tmp/gpt2-wikitext/generation_config.json
  Model weights saved in /tmp/gpt2-wikitext/tf_model.h5
  D0707 21:38:45.640973681   12027 init.cc:191]                          grpc_shutdown starts clean-up now
  

清理

1. 断开与 TPU 虚拟机实例的连接：

   (vm)$ exit

   您的提示符现在应为 username@projectname，表明您位于 Cloud Shell 中。

2. 删除 TPU 资源。

   $ gcloud compute tpus tpu-vm delete hf-gpt2 \
   --zone=us-central2-b

后续步骤

请尝试使用其他受支持的参考模型之一。