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.

Ziele

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

Kosten

In dieser Anleitung werden kostenpflichtige Komponenten von Google Cloud verwendet, darunter:

  • Compute Engine
  • Cloud TPU

Sie können mithilfe des Preisrechners die Kosten für Ihre voraussichtliche Nutzung kalkulieren. Neuen Google Cloud-Nutzern steht möglicherweise eine kostenlose Testversion zur Verfügung.

Hinweise

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

  1. Melden Sie sich bei Ihrem Google-Konto an.

    Wenn Sie noch kein Konto haben, melden Sie sich hier für ein neues Konto an.

  2. Wählen Sie in der Cloud Console auf der Seite für die Projektauswahl ein Cloud-Projekt aus oder erstellen Sie eines.

    Zur Projektauswahl

  3. Die Abrechnung für das Google Cloud-Projekt muss aktiviert sein. So prüfen Sie, ob die Abrechnung für Ihr Projekt aktiviert ist.

  4. 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.

    Cloud Shell öffnen

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

    export PROJECT_ID=project-id
    
  3. Konfigurieren Sie das gcloud-Befehlszeilentool für das Projekt, in dem Sie eine Cloud TPU erstellen möchten.

    gcloud config set project ${PROJECT_ID}
    
  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-1.6  \
    --accelerator-type=v3-8
    
  2. Ermitteln Sie die IP-Adresse für die Cloud TPU-Ressource.

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

PyTorch-Umgebung erstellen und konfigurieren

  1. Starten Sie eine conda-Umgebung.

    (vm) $ conda activate torch-xla-1.6
    
  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-1.6/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) $ 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
   

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. Sie gibt an, dass Sie sich in Cloud Shell befinden.

  2. Verwenden Sie in Ihrer Cloud Shell das gcloud-Befehlszeilentool, um die Compute Engine-Instanz zu löschen.

    $ gcloud compute instances delete roberta-tutorial --zone=us-central1-a
    
  3. Verwenden Sie zum Löschen der Cloud TPU-Ressource das gcloud-Befehlszeilentool.

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

Nächste Schritte

Testen Sie die PyTorch Colabs: