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

Ce tutoriel utilise des composants facturables de Google Cloud, dont :

  • Instance
  • 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.

    Si vous n'en possédez pas déjà un, vous devez en créer un.

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

    Accéder à la page de sélection du projet

  3. Vérifiez que la facturation est activée pour votre projet Google Cloud. Découvrez comment vérifier que la facturation est activée pour votre projet.

  4. 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 l'outil de ligne de commande gcloud pour utiliser le projet dans lequel vous souhaitez créer Cloud TPU.

    gcloud config set project ${PROJECT_ID}
    
  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-1.6  \
    --accelerator-type=v3-8
    
  2. Identifiez l'adresse IP de la ressource Cloud TPU.

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

Créer et configurer l'environnement PyTorch

  1. Démarrez un environnement conda.

    (vm) $ conda activate torch-xla-1.6
    
  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. Obtenez le DLRM compatible avec TPU en exécutant la commande suivante :

    (vm) $ git clone --recursive https://github.com/pytorch-tpu/examples.git
    
  2. Installez les dépendances.

    (vm) $ pip install onnx
    
  3. Exécutez le modèle sur des données aléatoires.

    (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
    

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

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

    Téléchargez l'ensemble de données Criteo Kaggle en suivant ces instructions. Une fois le téléchargement du fichier dac.tar.gz terminé, placez-le dans le répertoire ./criteo-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.

    Prétraitez l'ensemble de données Criteo en exécutant la commande suivante. Notez qu'il s'agit essentiellement de la même commande que la commande d'entraînement. Cependant, nous ne l'exécutons que sur un seul cœur, afin de terminer l'exécution de l'entraînement une fois le prétraitement terminé et le fichier kaggleAdDisplayChallenge_processed.npz généré. Le prétraitement peut prendre beaucoup de temps (plusieurs heures).

    (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 \
        --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 1 \
        --tpu-metrics-debug \
        --tpu-cores=1
    

    Lorsque les journaux s'affichent comme suit, l'ensemble de données a été prétraité et vous pouvez terminer (Ctrl^C) l'exécution et exécuter pkill python pour vous assurer qu'il n'y a pas de processus en fuite.

    Load 6548659/6548660  Split: 0  Label True: 0  Stored: 0
    Saved /home/jysohn/test/train_day_0.npz!
    Load 6548659/6548660  Split: 1  Label True: 1  Stored: 1
    Saved /home/jysohn/test/train_day_1.npz!
    Load 3742283/6548660  Split: 2  Label True: 0  Stored: 0
    ...
    time/loss/accuracy (if enabled): 2020-08-26 13:17:18.139566
    Finished training it 1024/306968 of epoch 0, -1.00 ms/it, loss 0.518172, accuracy 75.720 %, 131072 samples, @ 2020-08-26 13:18:04.515570
    
  3. Vérifiez que le prétraitement a bien été effectué.

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

  4. Exécuter l'entraînement sur l'ensemble de données Criteo Kaggle prétraité.

    (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
    

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

Nettoyer

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, ce qui indique que vous êtes dans Cloud Shell.

  2. Dans Cloud Shell, utilisez l'outil de ligne de commande gcloud pour supprimer l'instance Compute Engine.

    $ gcloud compute instances delete dlrm-tutorial --zone=us-central1-a
    
  3. Utilisez l'outil de ligne de commande gcloud pour supprimer la ressource Cloud TPU.

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

Étape suivante

Essayez les colabs PyTorch :