Entraîner DLRM sur Cloud TPU avec PyTorch

Ce tutoriel vous explique comment entraîner Facebook Research DLRM sur un appareil Cloud TPU.

Objectifs

  • Créer et configurer l'environnement PyTorch
  • Exécuter la tâche d'entraînement avec des données factices
  • (Facultatif) Entraîner sur un ensemble de données Criteo Kaggle

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. Lorsque vous avez terminé, veillez à nettoyer les ressources que vous avez créées, afin d'éviter des frais inutiles.

Configurer une instance Compute Engine

  1. Ouvrez une fenêtre Cloud Shell.

    Ouvrir Cloud Shell

  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. Remarque : Si vous effectuez un entraînement sur l'ensemble de données Criteo Kaggle, vous devrez utiliser une valeur machine-type n1-highmem-96.

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

    gcloud compute ssh dlrm-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 dlrm-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 dlrm-tutorial --zone=us-central1-a

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"

Exécuter la tâche d'entraînement avec des données factices

  1. Installez les dépendances.

    (vm) $ pip install onnx

  2. Exécutez le modèle sur des données aléatoires. Cette opération devrait prendre cinq à dix minutes.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-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

(Facultatif) Entraîner sur un ensemble de données Criteo Kaggle

Ces étapes sont facultatives. Vous ne devez les exécuter que si vous souhaitez effectuer l'entraînement sur l'ensemble de données Criteo Kaggle.

  1. Téléchargez l'ensemble de données.

    Téléchargez l'ensemble de données à partir de l'ensemble de données Criteo Kaggle en suivant ces instructions. Une fois le téléchargement terminé, copiez le fichier dac.tar.gz dans un répertoire nommé ./criteo-kaggle/. Utilisez la commande tar -xzvf pour extraire le contenu du fichier tar.gz dans le répertoire ./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. Prétraiter l'ensemble de données.

    Exécutez ce script pour prétraiter l'ensemble de données Criteo. Ce script produit un fichier nommé kaggleAdDisplayChallenge_processed.npz et met plus de trois heures à prétraiter l'ensemble de données.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_data_pytorch.py \
    --data-generation=dataset \
    --data-set=kaggle \
    --raw-data-file=criteo-kaggle/train.txt \
    --mini-batch-size=128 \
    --memory-map \
    --test-mini-batch-size=16384 \
    --test-num-workers=4

  3. Vérifiez que le prétraitement a bien été effectué.

    Le fichier kaggleAdDisplayChallenge_processed.npz devrait s'afficher dans le répertoire criteo-kaggle.

  4. Exécutez le script d'entraînement suivant sur l'ensemble de données prétraité Criteo Kaggle.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-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 \
    --memory-map \
    --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

    L'entraînement prend généralement plus de deux heures avec une précision de plus de 78,75 %.

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, supprimez l'instance Compute Engine à l'aide de Google Cloud CLI:

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

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

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

Étapes suivantes

Essayez les colabs PyTorch :