Esta página foi traduzida pela API Cloud Translation.

Inferência do MaxDiffusion em TPUs v6e

Este tutorial mostra como publicar modelos MaxDiffusion na TPU v6e. Neste tutorial, vai gerar imagens com o modelo Stable Diffusion XL.

Antes de começar

Prepare-se para aprovisionar uma TPU v6e com 4 chips:

Siga o guia Configure o ambiente do Cloud TPU para configurar um Google Cloud projeto, configurar a CLI do Google Cloud, ativar a API Cloud TPU e garantir que tem acesso à utilização de Cloud TPUs.

Autentique-se com o Google Cloud e configure o projeto e a zona predefinidos para a CLI Google Cloud.

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

Capacidade segura

Quando estiver pronto para proteger a capacidade da TPU, consulte as cotas da Cloud TPU para mais informações sobre as cotas da Cloud TPU. Se tiver mais perguntas sobre a proteção da capacidade, contacte a equipa de vendas ou da conta do Cloud TPU.

Aprovisione o ambiente do Cloud TPU

Pode aprovisionar VMs de TPU com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

Verifique se o seu projeto tem quota suficiente, que especifica o número máximo de chips aos quais pode aceder no seu projeto.TPUS_PER_TPU_FAMILYGoogle Cloud
Verifique se o seu projeto tem quota de TPUs suficiente para:
- Quota de VMs da TPU
- Quota de endereços IP
- Quota de Hyperdisk Balanced
Autorizações do projeto do utilizador
- Se estiver a usar o GKE com o XPK, consulte as autorizações da Cloud Console na conta de utilizador ou de serviço para ver as autorizações necessárias para executar o XPK.

Aprovisione uma 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

Use os comandos list ou describe para consultar o estado do recurso em fila.

   gcloud alpha compute tpus queued-resources describe QUEUED_RESOURCE_ID  \
      --project=PROJECT_ID --zone=ZONE

Para ver uma lista completa dos estados dos pedidos de recursos em fila, consulte a documentação Recursos em fila.

Estabeleça ligação à TPU através do SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Crie um ambiente Conda

Crie um diretório para o Miniconda:
```
mkdir -p ~/miniconda3
```

Transfira o script do instalador do Miniconda:

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

Instale o Miniconda:

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

Remova o script do instalador do Miniconda:
```
rm -rf ~/miniconda3/miniconda.sh
```
Adicione o Miniconda à sua variável PATH:
```
export PATH="$HOME/miniconda3/bin:$PATH"
```
Atualize ~/.bashrc para aplicar as alterações à variável PATH:
```
source ~/.bashrc
```
Crie um novo ambiente Conda:
```
conda create -n tpu python=3.10
```
Ative o ambiente Conda:
```
source activate tpu
```

Configure o MaxDiffusion

Clone o repositório do GitHub do MaxDiffusion e navegue para o diretório MaxDiffusion:
```
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
```
Mude para o ramo mlperf-4.1:
```
git checkout mlperf4.1
```
Instale o MaxDiffusion:
```
pip install -e .
```
Instalar dependências:
```
pip install -r requirements.txt
```

Instale o 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

Instale dependências adicionais:

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

Gerar imagens

Defina variáveis de ambiente para configurar o tempo de execução da TPU:

LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"

Gere imagens através do comando e das configurações definidas em src/maxdiffusion/configs/base_xl.yml:
```
python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"
```
Quando as imagens forem geradas, certifique-se de que limpa os recursos da TPU.

Limpar

Elimine a TPU:

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