Pré-formation FairSeq RoBERTa sur Cloud TPU avec PyTorch

Ce didacticiel vous montre comment pré-entraîner FairSeq RoBERTa sur un TPU Cloud. Plus précisément, il suit le tutoriel de FairSeq, qui consiste à pré-entraîner le modèle sur l'ensemble de données public wikitext -103.

Objectifs

  • Créer et configurer l'environnement PyTorch
  • Préparer l'ensemble de données
  • Exécuter la tâche d'entraînement
  • Vérifiez que vous pouvez afficher les résultats

Coûts

Dans ce document, vous utilisez les composants facturables suivants de Google Cloud :

  • Compute Engine
  • Cloud TPU

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût. Les nouveaux utilisateurs de Google Cloud peuvent bénéficier d'un essai gratuit.

Avant de commencer

Avant de commencer ce tutoriel, vérifiez que votre projet Google Cloud est correctement configuré.

  1. Connectez-vous à votre compte Google Cloud. Si vous débutez sur Google Cloud, créez un compte pour évaluer les performances de nos produits en conditions réelles. Les nouveaux clients bénéficient également de 300 $ de crédits gratuits pour exécuter, tester et déployer des charges de travail.

  2. Dans Google Cloud Console, sur la page de sélection du projet, sélectionnez ou créez un projet Google Cloud.

    Accéder au sélecteur de projet

  3. Vérifiez que la facturation est activée pour votre projet Google Cloud.

  4. Dans Google Cloud Console, sur la page de sélection du projet, sélectionnez ou créez un projet Google Cloud.

    Accéder au sélecteur de projet

  5. Vérifiez que la facturation est activée pour votre projet Google Cloud.

    6. Ce tutoriel utilise des composants facturables de Google Cloud. Consultez la grille tarifaire de Cloud TPU pour estimer vos coûts. Veillez à nettoyer les ressources que vous avez créées lorsque vous avez terminé, afin d'éviter des frais inutiles.

Configurer une instance Compute Engine

  1. Ouvrez une fenêtre Cloud Shell.

    Ouvrir v

  2. Créez une variable pour l'ID de votre projet.

    export PROJECT_ID=project-id

  3. Configurez Google Cloud CLI pour utiliser le projet dans lequel vous souhaitez créer Cloud TPU.

    gcloud config set project ${PROJECT_ID}

    La première fois que vous exécutez cette commande dans une nouvelle VM Cloud Shell, une page Authorize Cloud Shell s'affiche. Cliquez sur Authorize en bas de la page pour autoriser gcloud à effectuer des appels d'API Google Cloud avec vos identifiants.

  4. À partir de Cloud Shell, lancez la ressource Compute Engine requise pour ce tutoriel.

    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. Connectez-vous à la nouvelle instance Compute Engine

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

Lancer une ressource Cloud TPU

  1. À partir de la machine virtuelle Compute Engine, lancez une ressource Cloud TPU à l'aide de la commande suivante :

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

  2. Identifiez l'adresse IP de la ressource Cloud TPU.

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

Créer et configurer l'environnement PyTorch

  1. Démarrez un environnement conda.

    (vm) $ conda activate torch-xla-2.0

  2. Configurez des variables d'environnement pour la ressource Cloud TPU.

    (vm) $ export TPU_IP_ADDRESS=ip-address

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

Configurer les données

  1. Installez FairSeq en exécutant la commande suivante:

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

  2. Créez un répertoire, pytorch-tutorial-data pour stocker les données du modèle.

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

  3. Suivez les instructions de la section README "Prétraiter les données" de FairSeq RoBERTa. La préparation de l'ensemble de données prend environ 10 minutes.

Entraîner le modèle

Pour entraîner le modèle, commencez par configurer des variables d'environnement:

(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

Exécutez ensuite le script suivant :

(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

Le script d'entraînement s'exécute pendant environ 15 minutes et, lorsqu'il se termine, génère un message similaire à celui-ci :

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

Vérifier les résultats

Une fois la tâche d'entraînement terminée, vous pouvez trouver vos points de contrôle dans le répertoire suivant :

$HOME/checkpoints

Effectuer un nettoyage

Pour éviter une facturation inutile sur votre compte, effectuez un nettoyage des ressources que vous avez créées :

  1. Déconnectez-vous de l'instance Compute Engine, si vous ne l'avez pas déjà fait :

    (vm) $ exit

    Votre invite devrait maintenant être user@projectname, indiquant que vous êtes dans Cloud Shell.

  2. Dans Cloud Shell, utilisez Google Cloud CLI pour supprimer l'instance Compute Engine.

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

  3. Utilisez Google Cloud CLI pour supprimer la ressource Cloud TPU.

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

Étapes suivantes

Essayez les colabs PyTorch :