Inférence MaxDiffusion sur les TPU v6e

Ce tutoriel explique comment diffuser des modèles MaxDiffusion sur TPU v6e. Dans ce tutoriel, vous allez générer des images à l'aide du modèle Stable Diffusion XL.

Avant de commencer

Préparez-vous à provisionner un TPU v6e avec quatre puces:

  1. Connectez-vous à votre compte Google. Si vous ne l'avez pas déjà fait, créez un compte.
  2. Dans la console Google Cloud, sélectionnez ou créez un Google Cloud projet à partir de la page de sélection des projets.
  3. Activez la facturation pour votre Google Cloud projet. La facturation est obligatoire pour toute Google Cloud utilisation.
  4. Installez les composants gcloud alpha.
  5. Exécutez la commande suivante pour installer la dernière version des composants gcloud.

    gcloud components update
    
  6. Activez l'API TPU à l'aide de la commande gcloud suivante dans Cloud Shell. Vous pouvez également l'activer à partir de la console Google Cloud.

    gcloud services enable tpu.googleapis.com
    
  7. Créez une identité de service pour la VM TPU.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE
  8. Créez un compte de service TPU et accordez-lui l'accès aux services Google Cloud .

    Les comptes de service permettent au service Google Cloud TPU d'accéder à d'autres services Google Cloud. Un compte de service géré par l'utilisateur est recommandé. Suivez ces guides pour créer et accorder des rôles. Les rôles suivants sont nécessaires:

    • Administrateur TPU: rôle nécessaire pour créer un TPU
    • Administrateur de l'espace de stockage: rôle nécessaire pour accéder à Cloud Storage
    • Écrivain de journaux: nécessaire pour écrire des journaux avec l'API Logging
    • Rédacteur de métriques Monitoring: nécessaire pour écrire des métriques dans Cloud Monitoring
  9. Authentifiez-vous avec Google Cloud et configurez le projet et la zone par défaut pour Google Cloud CLI.

    gcloud auth login
    gcloud config set project PROJECT_ID
    gcloud config set compute/zone ZONE

Sécuriser la capacité

Contactez votre équipe commerciale ou votre équipe chargée des comptes Cloud TPU pour demander un quota de TPU et poser des questions sur la capacité.

Provisionner l'environnement Cloud TPU

Vous pouvez provisionner des TPU v6e avec GKE, avec GKE et XPK, ou en tant que ressources mises en file d'attente.

Prérequis

  • Vérifiez que votre projet dispose d'un quota TPUS_PER_TPU_FAMILY suffisant, qui spécifie le nombre maximal de chips auxquels vous pouvez accéder dans votre projetGoogle Cloud .
  • Ce tutoriel a été testé avec la configuration suivante :
    • Python 3.10 or later
    • Versions logicielles nocturnes :
      • 0.4.32.dev20240912 JAX
      • LibTPU 0.1.dev20240912+nightly nightly
    • Versions logicielles stables :
      • JAX + Bibliothèque JAX de v0.4.35
  • Vérifiez que votre projet dispose d'un quota TPU suffisant pour :
    • Quota de VM TPU
    • Quota d'adresses IP
    • Quota Hyperdisk équilibré
  • Autorisations de l'utilisateur sur le projet

Provisionner un TPU v6e

   gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
        --node-id TPU_NAME \
        --project PROJECT_ID \
        --zone ZONE \
        --accelerator-type v6e-4 \
        --runtime-version v2-alpha-tpuv6e \
        --service-account SERVICE_ACCOUNT

Utilisez les commandes list ou describe pour interroger l'état de votre ressource mise en file d'attente.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

Pour obtenir la liste complète des états des requêtes de ressources en file d'attente, consultez la documentation sur les ressources en file d'attente.

Se connecter au TPU à l'aide de SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Créer un environnement Conda

  1. Créez un répertoire pour Miniconda:

    mkdir -p ~/miniconda3
  2. Téléchargez le script d'installation de Miniconda:

    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
  3. Installez Miniconda:

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
  4. Supprimez le script d'installation Miniconda:

    rm -rf ~/miniconda3/miniconda.sh
  5. Ajoutez Miniconda à votre variable PATH:

    export PATH="$HOME/miniconda3/bin:$PATH"
  6. Actualisez ~/.bashrc pour appliquer les modifications à la variable PATH:

    source ~/.bashrc
  7. Créez un environnement Conda:

    conda create -n tpu python=3.10
  8. Activez l'environnement Conda:

    source activate tpu

Configurer MaxDiffusion

  1. Clonez le dépôt MaxDiffusion et accédez au répertoire MaxDiffusion:

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
  2. Passez à la branche mlperf-4.1:

    git checkout mlperf4.1
  3. Installez MaxDiffusion:

    pip install -e .
  4. Installez les dépendances :

    pip install -r requirements.txt
  5. Installez JAX:

    pip install -U --pre jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
  6. Installez des dépendances supplémentaires:

     pip install huggingface_hub==0.25 absl-py flax tensorboardX google-cloud-storage torch tensorflow transformers 

Générer des images

  1. Définissez des variables d'environnement pour configurer l'environnement d'exécution TPU:

    LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"
  2. Générez des images à l'aide de la requête et des configurations définies dans src/maxdiffusion/configs/base_xl.yml:

    python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"

    Une fois les images générées, veillez à nettoyer les ressources TPU.

Effectuer un nettoyage

Supprimez le TPU:

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async