Esta página foi traduzida pela API Cloud Translation.

Inferência do MaxDiffusion em TPUs v6e

Este tutorial mostra como exibir modelos MaxDiffusion na TPU v6e. Neste tutorial, você vai gerar imagens usando o modelo Stable Diffusion XL.

Antes de começar

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

Faça login na sua Conta do Google. Se ainda não tiver uma, crie uma nova conta.
No console do Google Cloud, selecione ou crie um projeto Google Cloud na página do seletor de projetos.
Ative o faturamento do projeto Google Cloud . O faturamento é obrigatório para todo o uso do Google Cloud .
Instale os componentes da gcloud alfa.
Execute o comando a seguir para instalar a versão mais recente dos componentes gcloud.
```
gcloud components update
```
Ative a API TPU usando o comando gcloud a seguir no Cloud Shell. Também é possível ativá-la no console do Google Cloud.
```
gcloud services enable tpu.googleapis.com
```

Crie uma identidade de serviço para a VM da TPU.

gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE

Crie uma conta de serviço do TPU e conceda acesso aos serviços Google Cloud .

As contas de serviço permitem que o serviço Google Cloud TPU acesse outros serviços Google Cloud. Recomendamos usar uma conta de serviço gerenciado pelo usuário. Siga estes guias para criar e conceder funções. Os seguintes papéis são necessários:
- Administrador de TPU: necessário para criar uma TPU
- Administrador de armazenamento: necessário para acessar o Cloud Storage
- Logs Writer: necessário para gravar registros com a API Logging.
- Gravador de métricas do Monitoring: necessário para gravar métricas no Cloud Monitoring
Faça a autenticação com Google Cloud e configure o projeto e a zona padrão para a Google Cloud CLI.
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

Capacidade segura

Entre em contato com a equipe de vendas ou de contas do Cloud TPU para solicitar a cota de TPU e fazer perguntas sobre a capacidade.

Provisionar o ambiente do Cloud TPU

É possível provisionar TPUs v6e com o GKE, com o GKE e o XPK ou como recursos em fila.

Pré-requisitos

Verifique se o projeto tem cota de TPUS_PER_TPU_FAMILY suficiente, que especifica o número máximo de chips que você pode acessar no projetoGoogle Cloud .
Este tutorial foi testado com a seguinte configuração:
- Python 3.10 or later
- Versões noturnas do software:
  - 0.4.32.dev20240912 noturno JAX
  - LibTPU noturno 0.1.dev20240912+nightly
- Versões estáveis do software:
  - JAX + JAX Lib de v0.4.35
Verifique se o projeto tem cota suficiente de TPU para:
- Cota da VM de TPU
- Quota de endereço IP
- Cota do Hyperdisk equilibrado
Permissões do projeto do usuário
- Se você estiver usando o GKE com XPK, consulte Permissões do console do Google Cloud na conta de usuário ou de serviço para conferir as permissões necessárias para executar o XPK.

Provisionar um 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ções de recursos em fila, consulte a documentação de Recursos em fila.

Conectar-se à TPU usando SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Criar um ambiente da 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 novo ambiente do Conda:
```
conda create -n tpu python=3.10
```
Ative o ambiente da Conda:
```
source activate tpu
```

Configurar o MaxDiffusion

Clone o repositório MaxDiffusion e navegue até o diretório MaxDiffusion:

git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion

Alterne 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 -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

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"

Gerar 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"
```
Depois que as imagens forem geradas, limpe os recursos da TPU.

Limpar

Exclua a TPU:

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