Inferência do MaxDiffusion em TPUs v6e

Este tutorial mostra como disponibilizar modelos do MaxDiffusion na TPU v6e. Durante o processo, você vai gerar imagens usando o modelo Stable Diffusion XL.

Antes de começar

Prepare-se para provisionar uma TPU v6e com quatro chips:

Siga o guia Configurar o ambiente do Cloud TPU para definir um projeto do Google Cloud , configurar a CLI do Google Cloud, ativar a API Cloud TPU e verificar se você tem acesso para usar Cloud TPUs.
Faça a autenticação com o Google Cloud e configure a zona e o projeto padrão para a CLI do Google Cloud.
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

Reservar capacidade

Quando estiver tudo pronto para reservar a capacidade da TPU, consulte Cotas do Cloud TPU para saber mais sobre as cotas do Cloud TPU. Se você tiver outras dúvidas sobre como reservar a capacidade, entre em contato com a equipe de vendas ou de conta do Cloud TPU.

Provisionar o ambiente do Cloud TPU

É possível provisionar VMs de TPU com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

Verifique se o projeto tem uma cota de TPUS_PER_TPU_FAMILY suficiente, porque ela especifica o número máximo de chips que podem ser acessados no projeto doGoogle Cloud .
Verifique se o projeto tem cota de TPU suficiente para:
- Cota de VM de TPU
- Cota de endereços IP
- Cota do Hyperdisk Balanced
Permissões do projeto do usuário
- Ao usar o GKE com o XPK, consulte Permissões do console do Cloud na conta de usuário ou de serviço para saber quais permissões são necessárias para executar o XPK.

Provisionar 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 status do recurso em fila.

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

Para uma lista completa de status de solicitação de recursos em fila, consulte a documentação Recursos em fila.

Conectar-se à TPU usando SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Criar um ambiente do Conda

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

Faça o download do script de instalação 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 de instalação do Miniconda:
```
rm -rf ~/miniconda3/miniconda.sh
```

Adicione o Miniconda à variável PATH:

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

Atualize ~/.bashrc para aplicar as mudanças à variável PATH:
```
source ~/.bashrc
```
Crie um ambiente do Conda:
```
conda create -n tpu python=3.10
```
Ative o ambiente do Conda:
```
source activate tpu
```

Configurar o MaxDiffusion

Clone o repositório do GitHub do MaxDiffusion e acesse o diretório do MaxDiffusion:
```
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
```
Mude para a ramificação mlperf-4.1:
```
git checkout mlperf4.1
```
Instale o MaxDiffusion:
```
pip install -e .
```
Instale as 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 extras:

 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 ambiente 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 usando o comando e as 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, limpe os recursos da TPU.

Limpeza

Exclua a TPU:

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