使用 PyTorch 在 Cloud TPU 上训练 DLRM

本教程介绍如何在 Cloud TPU 上训练 Facebook Research DLRM

目标

  • 创建并配置 PyTorch 环境
  • 使用虚构数据运行训练作业
  • (可选)在 Criteo Kaggle 数据集上训练

费用

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

  • Compute Engine
  • Cloud TPU

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

准备工作

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

  1. 登录您的 Google 帐号。

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

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

    转到项目选择器页面

  3. 确保您的 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}
    

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

  4. 从 Cloud Shell 启动本教程所需的 Compute Engine 资源。 注意:如果您在 Criteo Kaggle 数据集上进行训练,则需要使用 n1-highmem-96 machine-type

    gcloud compute instances create dlrm-tutorial \
    --zone=us-central1-a \
    --machine-type=n1-standard-64 \
    --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 dlrm-tutorial --zone=us-central1-a
    

启动 Cloud TPU 资源

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

    (vm) $ gcloud compute tpus create dlrm-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-1.7  \
    --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.7
    
  2. 为 Cloud TPU 资源配置环境变量。

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

使用虚构数据运行训练作业

  1. 运行以下命令以获取与 TPU 兼容的 DLRM:

    (vm) $ git clone --recursive https://github.com/pytorch-tpu/examples.git
    
  2. 安装依赖项。

    (vm) $ pip install onnx
    
  3. 对随机数据运行模型。

    (vm) $ python examples/deps/dlrm/dlrm_tpu_runner.py \
        --arch-embedding-size=1000000-1000000-1000000-1000000-1000000-1000000-1000000-1000000 \
        --arch-sparse-feature-size=64 \
        --arch-mlp-bot=512-512-64 \
        --arch-mlp-top=1024-1024-1024-1 \
        --arch-interaction-op=dot \
        --lr-num-warmup-steps=10 \
        --lr-decay-start-step=10 \
        --mini-batch-size=2048 \
        --num-batches=1000 \
        --data-generation='random' \
        --numpy-rand-seed=727 \
        --print-time \
        --print-freq=100 \
        --num-indices-per-lookup=100 \
        --use-tpu \
        --num-indices-per-lookup-fixed \
        --tpu-model-parallel-group-len=8 \
        --tpu-metrics-debug \
        --tpu-cores=8
    

(可选)在 Criteo Kaggle 数据集上训练

以下是可选步骤。仅当您希望在 Criteo Kaggle 数据集上训练时才应运行这些步骤。

  1. 下载数据集。

    按照此处的说明从 Criteo Kaggle 数据集中下载数据集。下载完成后,将 dac.tar.gz 文件复制到名为 ./criteo-kaggle/ 的目录中。使用 tar -xzvf 命令将 tar.gz 文件的内容解压缩到 ./critero-kaggle 目录中。

     (vm) $ mkdir criteo-kaggle
     (vm) $ cd criteo-kaggle
     (vm) $ # Download dataset from above link here.
     (vm) $ tar -xzvf dac.tar.gz
     (vm) $ cd ..
    
  2. 预处理数据集。

    启动此脚本以预处理 Criteo 数据集。此脚本会生成一个名为 kaggleAdDisplayChallenge_processed.npz 的文件,并且需要 3 个小时以上来预处理数据集。

    (vm) $ python examples/deps/dlrm/dlrm_data_pytorch.py \
        --data-generation=dataset \
        --data-set=kaggle \
        --raw-data-file=criteo-kaggle/train.txt \
        --mini-batch-size=128 \
        --test-mini-batch-size=16384 \
        --test-num-workers=4
    
  3. 验证预处理是否成功。

    您应该会在 criteo-kaggle 目录中看到 kaggleAdDisplayChallenge_processed.npz 文件。

  4. 在预处理过的 Criteo Kaggle 数据集上运行训练脚本。

    (vm) $ python examples/deps/dlrm/dlrm_tpu_runner.py \
        --arch-sparse-feature-size=16 \
        --arch-mlp-bot="13-512-256-64-16" \
        --arch-mlp-top="512-256-1" \
        --data-generation=dataset \
        --data-set=kaggle \
        --raw-data-file=criteo-kaggle/train.txt \
        --processed-data-file=criteo-kaggle/kaggleAdDisplayChallenge_processed.npz \
        --loss-function=bce \
        --round-targets=True \
        --learning-rate=0.1 \
        --mini-batch-size=128 \
        --print-freq=1024 \
        --print-time \
        --test-mini-batch-size=16384 \
        --test-num-workers=4 \
        --test-freq=101376 \
        --use-tpu \
        --num-indices-per-lookup=1 \
        --num-indices-per-lookup-fixed \
        --tpu-model-parallel-group-len 8 \
        --tpu-metrics-debug \
        --tpu-cores=8
    

    训练应在 2 个多小时内完成,且准确率高于 78.75%。

清理

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

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

    (vm) $ exit
    

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

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

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

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

后续步骤

试用 PyTorch Colab: