Training im Voraus von FairSeq RoBERTa mit PyTorch auf Cloud TPU

In dieser Anleitung erfahren Sie, wie Sie FairSeq RoBERTa auf einer Cloud TPU im Voraus trainieren. Sie orientiert sich insbesondere an der FairSeq-Anleitung, in der das Modell für das öffentliche wikitext-103-Dataset im Voraus trainiert wird.

Lernziele

  • PyTorch-Umgebung erstellen und konfigurieren
  • Dataset vorbereiten
  • Trainingsjob ausführen
  • Überprüfen Sie, ob Sie die Ausgabeergebnisse ansehen können

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

  • Compute Engine
  • Cloud TPU

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen. Neuen Google Cloud-Nutzern steht möglicherweise eine kostenlose Testversion zur Verfügung.

Hinweis

Bevor Sie mit dieser Anleitung beginnen, prüfen Sie, ob Ihr Google Cloud-Projekt ordnungsgemäß eingerichtet ist.

  1. Melden Sie sich bei Ihrem Google Cloud-Konto an. Wenn Sie mit Google Cloud noch nicht vertraut sind, erstellen Sie ein Konto, um die Leistungsfähigkeit unserer Produkte in der Praxis sehen und bewerten zu können. Neukunden erhalten außerdem ein Guthaben von 300 $, um Arbeitslasten auszuführen, zu testen und bereitzustellen.

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

    Go to project selector

  3. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

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

    Go to project selector

  5. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein.

    6. In dieser Anleitung werden kostenpflichtige Komponenten der Google Cloud verwendet. Rufen Sie die Seite mit den Cloud TPU-Preisen auf, um Ihre Kosten abzuschätzen. Denken Sie daran, nicht mehr benötigte Ressourcen zu bereinigen, um unnötige Kosten zu vermeiden.

Compute Engine-Instanz einrichten

  1. Öffnen Sie ein Cloud Shell-Fenster.

    Geöffnet

  2. Erstellen Sie eine Variable für Ihre Projekt-ID.

    export PROJECT_ID=project-id

  3. Konfigurieren Sie die Google Cloud CLI so, dass sie das Projekt verwendet, in dem Sie Cloud TPU erstellen möchten.

    gcloud config set project ${PROJECT_ID}

    Wenn Sie diesen Befehl zum ersten Mal in einer neuen Cloud Shell-VM ausführen, wird die Seite Authorize Cloud Shell angezeigt. Klicken Sie unten auf der Seite auf Authorize, um gcloud zu erlauben, Google Cloud API-Aufrufe mit Ihren Anmeldedaten auszuführen.

  4. Starten Sie in Cloud Shell die für diese Anleitung erforderliche Compute Engine-Ressource.

    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. Stellen Sie eine Verbindung zur neuen Compute Engine-Instanz her.

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

Cloud TPU-Ressource starten

  1. Starten Sie über die virtuelle Compute Engine-Maschine mit dem folgenden Befehl eine Cloud TPU-Ressource:

    (vm) $ gcloud compute tpus create roberta-tutorial \
--zone=us-central1-a \
--network=default \
--version=pytorch-2.0  \
--accelerator-type=v3-8

  2. Ermitteln Sie die IP-Adresse für die Cloud TPU-Ressource.

    (vm) $ gcloud compute tpus describe --zone=us-central1-a roberta-tutorial

PyTorch-Umgebung erstellen und konfigurieren

  1. Starten Sie eine conda-Umgebung.

    (vm) $ conda activate torch-xla-2.0

  2. Konfigurieren Sie Umgebungsvariablen für die Cloud TPU-Ressource.

    (vm) $ export TPU_IP_ADDRESS=ip-address

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

Daten einrichten

  1. Führen Sie den folgenden Befehl aus, um FairSeq zu installieren:

    (vm) $ pip install --editable /usr/share/torch-xla-2.0/tpu-examples/deps/fairseq

  2. Erstellen Sie das Verzeichnis "pytorch-tutorial-data", um die Modelldaten zu speichern.

    (vm) $ mkdir $HOME/pytorch-tutorial-data
(vm) $ cd $HOME/pytorch-tutorial-data

  3. Folgen Sie der Anleitung im README-Abschnitt "Daten vorverarbeiten" von FairSeq RoBERTA. Die Vorbereitung des Datasets dauert etwa 10 Minuten.

Modell trainieren

Richten Sie zum Trainieren des Modells zuerst einige Umgebungsvariablen ein:

(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

Führen Sie dann das folgende Skript aus:

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

Das Trainingsskript wird etwa 15 Minuten lang ausgeführt. Wenn es abgeschlossen ist, wird eine Meldung wie die folgende generiert:

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

Ausgabeergebnisse prüfen

Sie finden Ihre Modellprüfpunkte nach Abschluss des Trainingsjobs im folgenden Verzeichnis:

$HOME/checkpoints

Bereinigen

Führen Sie eine Bereinigung durch, damit Ihr Konto nach der Verwendung der von Ihnen erstellten Ressourcen nicht unnötig belastet wird:

  1. Trennen Sie die Verbindung zur Compute Engine-Instanz, sofern noch nicht geschehen:

    (vm) $ exit

    Die Eingabeaufforderung sollte nun user@projectname lauten und angeben, dass Sie sich in Cloud Shell befinden.

  2. Löschen Sie die Compute Engine-Instanz in der Cloud Shell über die Google Cloud CLI.

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

  3. Google Cloud CLI zum Löschen der Cloud TPU-Ressource verwenden

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

Nächste Schritte

Testen Sie die PyTorch Colabs: