使用 PyTorch 在 Cloud TPU 上预训练 FairSeq RoBERTa

本教程介绍如何在 Cloud TPU 上训练 FairSeq 的 RoBERTa。具体来说,它遵循 FairSeq 教程,使用公共数据集 wikitext-103 来预训练模型。

目标

  • 创建并配置 PyTorch 环境
  • 准备数据集
  • 运行训练作业
  • 确认您可以查看输出结果

费用

本教程使用 Google Cloud 的以下收费组件:

  • Compute Engine
  • Cloud TPU

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

准备工作

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

  1. 登录您的 Google 帐号。

    如果您还没有 Google 帐号,请注册一个新帐号

  2. 在 Cloud Console 的项目选择器页面上,选择或创建 Cloud 项目。

    转到项目选择器页面

  3. 确保您的 Google Cloud 项目已启用结算功能。 了解如何确认您的项目已启用结算功能

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

设置 Compute Engine 实例

  1. 打开一个 Cloud Shell 窗口。

    打开 Cloud Shell

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

    export PROJECT_ID=project-id
    
  3. 配置 gcloud 命令行工具,以使用要在其中创建 Cloud TPU 的项目。

    gcloud config set project ${PROJECT_ID}
    
  4. 从 Cloud Shell 启动本教程所需的 Compute Engine 资源。

    gcloud compute instances create roberta-tutorial \
    --zone=us-central1-a \
    --machine-type=n1-standard-16  \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=200GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform
    
  5. 连接到新的 Compute Engine 实例。

    gcloud compute ssh roberta-tutorial --zone=us-central1-a
    

启动 Cloud TPU 资源

  1. 在 Compute Engine 虚拟机中,使用以下命令启动 Cloud TPU 资源:

    (vm) $ gcloud compute tpus create roberta-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.6  \
    --accelerator-type=v3-8
    
  2. 确定 Cloud TPU 资源的 IP 地址。

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

创建并配置 PyTorch 环境

  1. 启动 conda 环境。

    (vm) $ conda activate torch-xla-1.6
    
  2. 为 Cloud TPU 资源配置环境变量。

    (vm) $ export TPU_IP_ADDRESS=ip-address
    
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

设置数据

  1. 运行以下命令来安装 FairSeq:

    (vm) $ pip install --editable /usr/share/torch-xla-1.6/tpu-examples/deps/fairseq
    
  2. 创建目录 pytorch-tutorial-data 来存储模型数据。

    (vm) $ mkdir $HOME/pytorch-tutorial-data
    (vm) $ cd $HOME/pytorch-tutorial-data
    
  3. 按照 FairSeq RoBERTa 的 README“预处理数据”部分中的说明操作。 准备数据集大约需要 10 分钟。

训练模型

要训练模型,请先设置一些环境变量:

    (vm) $ export TOTAL_UPDATES=125000    # Total number of training steps
    (vm) $ export WARMUP_UPDATES=10000    # Warmup the learning rate over this many updates
    (vm) $ export PEAK_LR=0.0005          # Peak learning rate, adjust as needed
    (vm) $ export TOKENS_PER_SAMPLE=512   # Max sequence length
    (vm) $ export UPDATE_FREQ=16          # Increase the batch size 16x
   
    (vm) $ export DATA_DIR=${HOME}/pytorch-tutorial-data/data-bin/wikitext-103
   

然后,运行以下脚本:

(vm) $ python /usr/share/torch-xla-1.6/tpu-examples/deps/fairseq/train.py $DATA_DIR \
    --task masked_lm \
    --criterion masked_lm \
    --arch roberta_base \
    --sample-break-mode complete \
    --tokens-per-sample $TOKENS_PER_SAMPLE \
    --optimizer adam \
    --adam-betas '(0.9,0.98)' \
    --adam-eps 1e-6 \
    --clip-norm 0.0 \
    --lr-scheduler polynomial_decay \
    --lr $PEAK_LR \
    --warmup-updates $WARMUP_UPDATES \
    --total-num-update $TOTAL_UPDATES \
    --dropout 0.1 \
    --attention-dropout 0.1 \
    --weight-decay 0.01 \
    --update-freq $UPDATE_FREQ \
    --max-update $TOTAL_UPDATES \
    --log-format simple \
    --valid-subset=valid \
    --train-subset=train \
    --num_cores=8 \
    --metrics_debug \
    --input_shapes 16x512 18x480 21x384 \
    --save-dir=${HOME}/checkpoints \
    --log_steps=30 \
    --max-epoch=1 \
    --skip-invalid-size-inputs-valid-test
   

训练脚本会运行大约 15 分钟,并在完成后生成类似于以下内容的消息:

    saved checkpoint /home/user/checkpoints/checkpoint1.pt
    (epoch 1 @ 119 updates) (writing took 25.19265842437744 seconds)
    | done training in 923.8 seconds

验证输出结果

训练作业完成后,可以在以下目录中查找模型检查点:

$HOME/checkpoints

清理

使用您创建的资源后,请进行清理,以免您的帐号产生不必要的费用:

  1. 断开与 Compute Engine 实例的连接(如果您尚未这样做):

    (vm) $ exit
    

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

  2. 在 Cloud Shell 中,使用 gcloud 命令行工具删除 Compute Engine 实例。

    $ gcloud compute instances delete roberta-tutorial --zone=us-central1-a
    
  3. 使用 gcloud 命令行工具删除 Cloud TPU 资源。

    $ gcloud compute tpus delete roberta-tutorial --zone=us-central1-a
    

后续步骤

试用 PyTorch Colab: