Esta página foi traduzida pela API Cloud Translation.

Treinar um modelo usando a TPU7x (Ironwood)

Este documento descreve como provisionar recursos TPU7x e dá um exemplo de implantação de uma carga de trabalho de treinamento usando MaxText e XPK.

A TPU7x é a primeira versão da família Ironwood,a TPU de sétima geração do Google Cloud. A geração Ironwood foi projetada para treinamento e inferência de IA em grande escala. Para mais informações, consulte TPU7x.

Para mais exemplos otimizados para TPU7x, consulte Receitas de treinamento para TPU Ironwood no GitHub.

Provisionar TPUs

É possível provisionar e gerenciar TPUs v7x usando os seguintes métodos:

GKE: é possível usar o GKE para provisionar e gerenciar TPUs como um pool de aceleradores para cargas de trabalho conteinerizadas de machine learning. Use a Google Cloud CLI para criar manualmente a instância do cluster do GKE para personalização precisa ou expansão de ambientes de produção do GKE. Para mais informações, consulte Sobre TPUs no GKE.
GKE e XPK: o XPK é uma ferramenta de linha de comando que simplifica a criação de clusters e a execução de cargas de trabalho no GKE. Com ele, os profissionais de ML podem provisionar TPUs e executar jobs de treinamento sem precisar ter grande familiaridade com o Kubernetes. Use o XPK para criar rapidamente clusters do GKE e executar cargas de trabalho para prova de conceito e testes. Para mais informações, consulte o repositório do GitHub do XPK.
GKE e TPU Cluster Director: o TPU Cluster Director está disponível por uma reserva no modo "Toda a capacidade", que oferece acesso total a toda a capacidade reservada (sem retenções) e visibilidade completa da topologia de hardware, do status de utilização e do status de integridade da TPU. Para mais informações, consulte Visão geral do modo de capacidade total.

Implantar uma carga de trabalho de treinamento com o MaxText e o XPK

Use o kit de processamento acelerado (XPK) para criar clusters do GKE para prova de conceito e testes. O XPK é uma ferramenta de linha de comando projetada para simplificar o provisionamento, o gerenciamento e a execução de cargas de trabalho de machine learning.

As seções a seguir mostram como implantar uma carga de trabalho de treinamento usando MaxText e XPK.

Antes de começar

Antes de começar, siga estas etapas:

Verifique se você tem um projeto do Google Cloud com o faturamento ativado.
Acessar o TPU7x. Para mais informações, entre em contato com sua equipe de conta.
Verifique se a conta que você está usando com o XPK tem as funções listadas no repositório do XPK no GitHub.

Instalar o XPK e as dependências

Instale o XPK. Siga as instruções no repositório do GitHub do XPK.
Instale o Docker seguindo as instruções fornecidas pelo administrador ou siga as instruções oficiais de instalação. Depois de instalado, execute os seguintes comandos para configurar o Docker e testar a instalação:
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
Configure as variáveis de ambiente a seguir:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export RESERVATION_NAME=YOUR_RESERVATION_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
Substitua:
- YOUR_PROJECT_ID: o ID do projeto do Google Cloud .
- YOUR_ZONE: a zona em que o cluster será criado. Na prévia, somente us-central1-c é compatível.
- YOUR_CLUSTER_NAME: o nome do novo cluster.
- YOUR_ACCELERATOR_TYPE: a versão e a topologia da TPU. Por exemplo, tpu7x-4x4x8. Para ver uma lista de topologias compatíveis, consulte Configurações aceitas.
- YOUR_RESERVATION_NAME: o nome da sua reserva. Para reservas compartilhadas, use projects/YOUR_PROJECT_NUMBER/reservations/YOUR_RESERVATION_NAME.
- YOUR_BUCKET_NAME: o nome do seu bucket do Cloud Storage, que será o diretório de saída para treinamento de modelo.

Se você não tiver um bucket do Cloud Storage, crie um usando o comando a seguir:

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

Criar um cluster de NIC única e fração única

Siga as instruções na seção Configurar MTU para otimizar a configuração de rede.
Preencha a variável ${CLUSTER_ARGUMENTS}, que você vai usar no comando xpk cluster create:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${SUBNET_NAME}"
```
Crie seu cluster do GKE com pools de nós TPU7x usando o comando xpk cluster create:
```
xpk cluster create \
    --project=${PROJECT_ID} \
    --zone=${ZONE} \
    --cluster ${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --tpu-type=${ACCELERATOR_TYPE} \
    --reservation=${RESERVATION_NAME} \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
```
Definir a flag --cluster-cpu-machine-type como n1-standard-8 (ou maior) garante que o pool de nós padrão tenha CPU suficiente para pods do sistema, por exemplo, o webhook JobSet, evitando erros. Por padrão, o XPK usa e2-standard-16. Algumas zonas só aceitam tipos específicos de CPU. Portanto, talvez seja necessário alternar entre os tipos n1, n2 e e2. Caso contrário, você poderá encontrar erros de cota.

Adicione uma exclusão de manutenção para evitar upgrades no cluster:

gcloud container clusters update ${CLUSTER_NAME} \
    --zone=${ZONE} \
    --add-maintenance-exclusion-name="no-upgrade-next-month" \
    --add-maintenance-exclusion-start="EXCLUSION_START_TIME" \
    --add-maintenance-exclusion-end="EXCLUSION_END_TIME" \
    --add-maintenance-exclusion-scope="no_upgrades"

Substitua:

EXCLUSION_START_TIME: o horário de início selecionado para a exclusão de manutenção no formato YYYY-MM-DDTHH:MM:SSZ.
EXCLUSION_END_TIME: o horário de término selecionado para a exclusão de manutenção no formato YYYY-MM-DDTHH:MM:SSZ.

Crie ou faça upload da imagem Docker do MaxText

É possível criar uma imagem do Docker localmente usando scripts fornecidos pelo MaxText ou usar uma imagem pré-criada.

Criar localmente

Os comandos a seguir copiam seu diretório local para o contêiner:

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Custom Jax and LibTPU wheels
pip download libtpu==0.0.28.dev20251104+nightly -f "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
pip download --pre jax==0.8.1.dev20251104 jaxlib==0.8.1.dev20251104 --index https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

# Build the Docker image
bash docker_build_dependency_image.sh MODE=custom_wheels

Após a execução dos comandos, uma imagem chamada maxtext_base_image será criada localmente. É possível usar a imagem local diretamente no comando xpk workload.

Fazer upload de uma imagem (opcional)

Depois de criar a imagem do Docker localmente usando as instruções na seção anterior, faça upload dela para o registro com o seguinte comando:

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

Após a execução bem-sucedida desse comando, a imagem MaxText vai aparecer em gcr.io com o nome gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

Definir o comando de treinamento do MaxText

Prepare o comando para executar o script de treinamento no contêiner do Docker.

O modelo MaxText 1B é uma configuração no framework MaxText projetada para treinar um modelo de linguagem com aproximadamente 1 bilhão de parâmetros. Use este modelo para testar escalas pequenas de chips. O desempenho não é otimizado.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

Implantar a carga de trabalho de treinamento

Execute o comando xpk workload create para implantar o job de treinamento. Você precisa especificar a flag --base-docker-image para usar a imagem de base do MaxText ou especificar a flag --docker-image e a imagem que você quer usar. Você pode incluir a flag --enable-debug-logs para ativar o registro de depuração.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

Os nomes das cargas de trabalho precisam ser exclusivos no cluster. Neste exemplo, $(date +%H%M) é anexado ao nome da carga de trabalho para garantir a exclusividade.