Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

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.

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.

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 as instruções oficiais de instalação. Depois de instalado, execute os comandos a seguir 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 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.
- 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 compatíveis.
- YOUR_BUCKET_NAME: o nome do seu bucket do Cloud Storage, que será o diretório de saída para o 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

Escolha uma das seguintes opções para criar o cluster. Para ter uma performance ideal,recomendamos usar uma rede personalizada com MTU de 8.896 bytes.

Rede personalizada

Para criar uma rede personalizada com MTU de 8.896 e usá-la no cluster, siga estas etapas:

Defina variáveis de ambiente para os nomes da rede e do firewall:
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
Substitua:
- NETWORK_NAME: um nome para a rede.
- FIREWALL_NAME: um nome para a regra de firewall de rede.

Crie uma rede personalizada com uma MTU de 8.896:

gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

Crie uma regra de firewall que permita o tráfego TCP, ICMP e UDP na sua rede:

gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

Defina uma variável de ambiente para os argumentos do cluster XPK para usar a rede que você criou:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

Crie o cluster do XPK. O comando a seguir provisiona a capacidade sob demanda:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

Para usar a capacidade reservada, substitua --on-demand por --reservation=RESERVATION_NAME. Para usar VMs spot de TPU, substitua --on-demand por --spot.

Rede padrão

Se você não precisar de uma rede de MTU alta, crie um cluster que use a rede VPC padrão. O comando a seguir provisiona a capacidade sob demanda:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

Para usar a capacidade reservada, substitua --on-demand por --reservation=RESERVATION_NAME. Para usar VMs spot de TPU, substitua --on-demand por --spot.

Criar ou fazer upload da imagem Docker do MaxText

É possível criar uma imagem 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

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

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 da carga de trabalho xpk.

Fazer upload de uma imagem (opcional)

Depois de criar a imagem Docker localmente usando as instruções da seção anterior, faça upload da imagem Docker do MaxText 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 pequenas escalas de chips. O desempenho não está 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"

Implante 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.