Inférence MaxDiffusion sur des TPU v6e

Ce tutoriel explique comment mettre en service des modèles MaxDiffusion sur des 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 :

Suivez le guide Configurer l'environnement Cloud TPU pour configurer un projet Google Cloud , configurer la Google Cloud CLI, activer l'API Cloud TPU et vous assurer d'avoir accès à Cloud TPU.
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é

Lorsque vous êtes prêt à sécuriser la capacité des TPU, consultez Quotas Cloud TPU pour en savoir plus sur les quotas Cloud TPU. Si vous avez d'autres questions sur la sécurisation de la capacité, contactez votre équipe commerciale ou l'équipe chargée du compte Cloud TPU.

Provisionner l'environnement Cloud TPU

Vous pouvez provisionner des VM TPU avec GKE, avec GKE et XPK, ou en tant que ressources 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 puces auxquelles vous pouvez accéder dans votre projetGoogle Cloud .
Vérifiez que votre projet dispose d'un quota TPU suffisant pour :
- Quota de VM TPU
- Quota d'adresses IP
- Quota Hyperdisk Balanced
Autorisations liées au projet utilisateur
- Si vous utilisez GKE avec XPK, consultez Autorisations de la console Cloud sur le compte utilisateur ou de service pour connaître les autorisations nécessaires à l'exécution de XPK.

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 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 demandes de ressources en file d'attente, consultez la documentation 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

Créez un répertoire pour Miniconda :
```
mkdir -p ~/miniconda3
```

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

Installez Miniconda :

bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

Supprimez le script d'installation de Miniconda :
```
rm -rf ~/miniconda3/miniconda.sh
```

Ajoutez Miniconda à votre variable PATH :

export PATH="$HOME/miniconda3/bin:$PATH"

Actualisez ~/.bashrc pour appliquer les modifications à la variable PATH :
```
source ~/.bashrc
```
Créez un nouvel environnement Conda :
```
conda create -n tpu python=3.10
```
Activez l'environnement Conda :
```
source activate tpu
```

Configurer MaxDiffusion

Clonez le dépôt GitHub MaxDiffusion et accédez au répertoire MaxDiffusion :
```
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
```
Passez à la branche mlperf-4.1 :
```
git checkout mlperf4.1
```
Installez MaxDiffusion :
```
pip install -e .
```
Installez les dépendances :
```
pip install -r requirements.txt
```

Installez JAX :

pip install jax[tpu]==0.4.34 jaxlib==0.4.34 ml-dtypes==0.2.0 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Installer 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

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"

Générez des images à l'aide de l'invite de commande 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