Vista geral do Cloud TPU Multislice

posteriores da TPU.

O Cloud TPU Multislice é uma tecnologia de escalamento de desempenho de pilha completa que permite que uma tarefa de preparação use várias divisões de TPU numa única divisão ou em divisões em vários Pods com paralelismo de dados padrão. Com os chips TPU v4, isto significa que as tarefas de preparação podem usar mais de 4096 chips numa única execução. Para tarefas de preparação que requerem menos de 4096 chips, uma única fatia pode oferecer o melhor desempenho. No entanto, estão mais facilmente disponíveis várias divisões mais pequenas, o que permite um tempo de início mais rápido quando a funcionalidade de várias divisões é usada com divisões mais pequenas.

Várias segmentações escalam o desempenho linearmente

Quando implementados em configurações de várias fatias, os chips de TPU em cada fatia comunicam através da interligação entre chips (ICI). Os chips de TPU em diferentes fatias comunicam transferindo dados para CPUs (anfitriões), que, por sua vez, transmitem os dados através da rede do centro de dados (DCN). Para mais informações sobre a escalabilidade com o Multislice, consulte o artigo Como dimensionar a preparação de IA para até dezenas de milhares de chips Cloud TPU com o Multislice.

Fluxo de dados com várias divisões

Os programadores não têm de escrever código para implementar a comunicação DCN entre segmentos. O compilador XLA gera esse código por si e sobrepõe a comunicação com os cálculos para um desempenho máximo.

Conceitos

Tipo de acelerador
O formato de cada slice de TPU que compõe um Multislice. Cada fatia num pedido de várias fatias é do mesmo tipo de acelerador. Um tipo de acelerador consiste num tipo de TPU (v4 ou posterior) seguido do número de TensorCores. Por exemplo, v5litepod-128 especifica uma TPU v5e com 128 TensorCores.
Reparação automática
Quando uma divisão encontra um evento de manutenção, uma preempção ou uma falha de hardware, o Cloud TPU cria uma nova divisão. Se existirem recursos insuficientes para criar uma nova fatia, a criação não é concluída até que o hardware fique disponível. Depois de criar a nova fatia, todas as outras fatias no ambiente Multislice são reiniciadas para que a preparação possa continuar. Com um script de início configurado corretamente, o script de preparação pode ser reiniciado automaticamente sem intervenção do utilizador, carregando e retomando a partir do ponto de verificação mais recente.
Redes de centros de dados (DCN)
Uma rede com latência mais elevada e débito mais baixo (quando comparada com a ICI) que liga fatias de TPU numa configuração de várias fatias.
Agendamento em grupo
Quando todas as divisões de TPU são aprovisionadas em conjunto, ao mesmo tempo, garantindo que todas ou nenhuma das divisões são aprovisionadas com êxito.
Interligação entre chips (ICI)
Ligações internas de alta velocidade e baixa latência que ligam TPUs num TPU Pod.
Multislice
Duas ou mais divisões do chip da TPU que podem comunicar através da DCN.
Node
No contexto de Multislice, o nó refere-se a uma única divisão da TPU. Cada fatia de TPU num Multislice recebe um ID do nó.
Script de arranque
Um script de arranque do Compute Engine padrão que é executado sempre que uma VM é iniciada ou reiniciada. Para o Multislice, é especificado no pedido de criação de QR. Para mais informações sobre scripts de arranque do Cloud TPU, consulte o artigo Gerir recursos de TPU.
Tensor
Uma estrutura de dados usada para representar dados multidimensionais num modelo de aprendizagem automática.
Tipos de capacidade do Cloud TPU

As TPUs podem ser criadas a partir de diferentes tipos de capacidade (consulte as opções de utilização em Como funciona o preço das TPUs):

  • Reserva: para consumir uma reserva, tem de ter um contrato de reserva com a Google. Use a flag --reserved quando criar os seus recursos.

  • Spot: segmenta a quota preemptiva com VMs do Spot. Os seus recursos podem ser interrompidos para dar lugar a pedidos de uma tarefa de prioridade mais elevada. Use a flag --spot ao criar os seus recursos.

  • A pedido: segmenta a quota a pedido, que não precisa de reserva e não é antecipada. O pedido de TPU é colocado numa fila de quotas a pedido oferecida pela Cloud TPU. A disponibilidade de recursos não é garantida. Selecionado por predefinição, não são necessárias etiquetas.

Começar

  1. Configure o seu ambiente de Cloud TPU.

  2. In the Google Cloud console, activate Cloud Shell.

    Activate Cloud Shell

    At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

    3. Para usar o Multislice, os seus recursos de TPU têm de ser geridos como recursos em fila.

    Exemplo introdutório

    Este tutorial usa código do repositório do GitHub do MaxText. O MaxText é um LLM básico de código aberto, arbitrariamente escalável, de elevado desempenho e bem testado, escrito em Python e Jax. O MaxText foi concebido para ser treinado de forma eficiente na TPU na nuvem.

    O código em shardings.py foi concebido para ajudar a começar a experimentar diferentes opções de paralelismo. Por exemplo, paralelismo de dados, paralelismo de dados totalmente fragmentados (FSDP) e paralelismo de tensores. O código é dimensionado de ambientes de fatia única para vários ambientes de fatias.

    Paralelismo de ICI

    ICI refere-se à interligação de alta velocidade que liga as TPUs numa única fatia. A divisão em partições de ICI corresponde à divisão em partições numa fatia. shardings.py oferece três parâmetros de paralelismo da ICI:

    • ici_data_parallelism
    • ici_fsdp_parallelism
    • ici_tensor_parallelism

    Os valores especificados para estes parâmetros determinam o número de fragmentos para cada método de paralelismo.

    Estas entradas têm de ser restritas para que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism seja igual ao número de chips na fatia.

    A tabela seguinte mostra exemplos de entradas do utilizador para o paralelismo da ICI para os quatro chips disponíveis na v4-8:

    ici_data_parallelism ici_fsdp_parallelism ici_tensor_parallelism
    FSDP de 4 vias 1 4 1
    Paralelismo de tensores quadridirecional 1 1 4
    FSDP bidirecional + paralelismo do Tensor bidirecional 1 2 2

    Tenha em atenção que ici_data_parallelism deve ser deixado como 1 na maioria dos casos, porque a rede ICI é suficientemente rápida para quase sempre preferir o FSDP ao paralelismo de dados.

    Este exemplo parte do princípio de que está familiarizado com a execução de código numa única fatia de TPU, como em Executar um cálculo numa VM de TPU na nuvem com o JAX. Este exemplo mostra como executar shardings.py numa única fatia.

    1. Configure o ambiente:

      $ gcloud auth login
$ export QR_ID=your-queued-resource-id
$ export TPU_NAME=your-tpu-name
$ export PROJECT=your-project-name
$ export ZONE=us-central1-a
$ export NETWORK_NAME=your-network-name
$ export SUBNETWORK_NAME=your-subnetwork-name
$ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
$ export ACCELERATOR_TYPE=v5litepod-16
$ export EXAMPLE_TAG_1=your-tag-1
$ export EXAMPLE_TAG_2=your-tag-2
$ export SLICE_COUNT=4
$ export STARTUP_SCRIPT='#!/bin/bash\n'

      Descrições de variáveis

      Entrada Descrição
      QR_ID O ID atribuído pelo utilizador do recurso em fila.
      TPU_NAME O nome atribuído pelo utilizador da sua TPU.
      PROJECT Google Cloud nome do projeto
      ZONE Especifica a zona na qual criar os recursos.
      NETWORK_NAME Nome das redes VPC.
      SUBNETWORK_NAME Nome da sub-rede nas redes VPC
      RUNTIME_VERSION A versão do software do Cloud TPU.
      ACCELERATOR_TYPE v4-16
      EXAMPLE_TAG_1, EXAMPLE_TAG_2 … Etiquetas usadas para identificar origens ou destinos válidos para firewalls de rede
      SLICE_COUNT Número de fatias. Limitado a um máximo de 256 fatias.
      STARTUP_SCRIPT Se especificar um script de arranque, o script é executado quando a fatia de TPU é aprovisionada ou reiniciada.

    2. Crie chaves SSH para gcloud. Recomendamos que deixe uma palavra-passe em branco (prima Enter duas vezes depois de executar o seguinte comando). Se lhe for pedido que o ficheiro google_compute_engine já existe, substitua a versão existente.

      $ ssh-keygen -f ~/.ssh/google_compute_engine

    3. Aprovisione as suas TPUs:

      gcloud

      $ gcloud compute tpus queued-resources \
    create ${QR_ID} \
    --accelerator-type=${ACCELERATOR_TYPE} \
    --runtime-version=${RUNTIME_VERSION} \
    --node-id=${TPU_NAME} \
    --zone=${ZONE} \
    [--reserved |--spot]

      A CLI Google Cloud não suporta todas as opções de criação de QR, como etiquetas. Para mais informações, consulte Crie QRs.

      Consola

      1. Na Google Cloud consola, aceda à página TPUs:

        Aceder às TPUs

      2. Clique em Criar TPU.

      3. No campo Nome, introduza um nome para a TPU.

      4. Na caixa Zona, selecione a zona onde quer criar a TPU.

      5. Na caixa Tipo de TPU, selecione um tipo de acelerador. O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.

      6. Na caixa Versão do software da TPU, selecione uma versão do software. Quando cria uma VM do Cloud TPU, a versão do software da TPU especifica a versão do tempo de execução da TPU a instalar. Para mais informações, consulte o artigo Versões de software da TPU.

      7. Clique no botão de ativar/desativar Ativar colocação em fila.

      8. No campo Nome do recurso em fila, introduza um nome para o seu pedido de recurso em fila.

      9. Clique em Criar para criar o pedido de recursos em fila.

    4. Aguarde até que o recurso em fila esteja no estado ACTIVE, o que significa que os nós de trabalho estão no estado READY. Assim que o aprovisionamento de recursos em fila de espera for iniciado, pode demorar entre 1 e 5 minutos a ser concluído, consoante o tamanho do recurso em fila de espera. Pode verificar o estado de um pedido de recurso em fila de espera através da CLI gcloud ou da Google Cloud consola:

      gcloud

      $ gcloud compute tpus queued-resources \
    list --filter=${QR_ID} --zone=${ZONE}

      Consola

      1. Na Google Cloud consola, aceda à página TPUs:

        Aceder às TPUs

      2. Clique no separador Recursos em fila.

      3. Clique no nome do seu pedido de recurso em fila.

    5. Estabeleça ligação à VM de TPU através de SSH:

      $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

    6. Clone o MaxText (que inclui o shardings.py) para a sua VM da TPU:

      $ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

    7. Instale o Python 3.10:

      $ sudo apt-get update
$ sudo apt install python3.10
$ sudo apt install python3.10-venv

    8. Crie e ative um ambiente virtual:

      $ python3 -m venv your-venv-name
$ source your-venv-name/bin/activate

    9. No diretório do repositório MaxText, execute o script de configuração para instalar o JAX e outras dependências na sua fatia de TPU. O script de configuração demora alguns minutos a ser executado.

      $ bash setup.sh

    10. Execute o seguinte comando para executar o shardings.py na sua fatia de TPU.

      $ python3 -m pedagogical_examples.shardings \
  --ici_fsdp_parallelism 4 \
  --batch_size 131072 \
  --embedding_dimension 2048

      Pode ver os resultados nos registos. As suas UTPs devem atingir cerca de 260 TFLOPs por segundo ou uma utilização de FLOPs impressionante de mais de 90%! Neste caso, selecionámos aproximadamente o lote máximo que se ajusta à memória de largura de banda elevada (HBM) da TPU.

    11. Não hesite em explorar outras estratégias de divisão na ICI. Por exemplo, pode experimentar a seguinte combinação:

      $ python3 -m pedagogical_examples.shardings \
  --ici_tensor_parallelism 4 \
  --batch_size 131072 \
  --embedding_dimension 2048

    12. Elimine o recurso em fila e a fatia de TPU quando terminar. Deve executar estes passos de limpeza no ambiente onde configurou a fatia (primeiro, execute exit para sair da sessão de SSH). A eliminação demora entre dois e cinco minutos. Se estiver a usar a CLI gcloud, pode executar este comando em segundo plano com a flag --async opcional.

      gcloud

      $ gcloud compute tpus queued-resources \
    delete ${QR_ID} --force (--async)

      Consola

      1. Na Google Cloud consola, aceda à página TPUs:

        Aceder às TPUs

      2. Clique no separador Recursos em fila.

      3. Selecione a caixa de verificação junto ao pedido de recurso em fila.

      4. Clique em Eliminar.

    Particionamento em vários fragmentos com paralelismo de DCN

    O script shardings.py usa três parâmetros que especificam o paralelismo da DCN, correspondente ao número de fragmentos de cada tipo de paralelismo de dados:

    • dcn_data_parallelism
    • dcn_fsdp_parallelism
    • dcn_tensor_parallelism

    Os valores destes parâmetros têm de ser restritos para que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism seja igual ao número de fatias.

    Como exemplo para duas fatias, use --dcn_data_parallelism = 2.

    dcn_data_parallelism dcn_fsdp_parallelism dcn_tensor_parallelism N.º de fatias
    Paralelismo de dados bidirecional 2 1 1 2

    dcn_tensor_parallelism deve ser sempre definido como 1 porque a DCN não é adequada para essa divisão. Para cargas de trabalho típicas de LLM em chips v4, dcn_fsdp_parallelism também deve ser definido como 1 e, por conseguinte, dcn_data_parallelism deve ser definido como o número de fatias, mas isto depende da aplicação.

    À medida que aumenta o número de divisões (partindo do princípio de que mantém o tamanho da divisão e o lote por divisão constantes), aumenta a quantidade de paralelismo de dados.

    Executar o shardings.py num ambiente de várias divisões

    Pode executar shardings.py num ambiente de vários fragmentos usando multihost_runner.py ou executando shardings.py em cada VM de TPU. Aqui, usamos multihost_runner.py. Os passos seguintes são muito semelhantes aos descritos no artigo Introdução: experiências rápidas em várias divisões do repositório MaxText, exceto que aqui executamos shardings.py em vez do LLM mais complexo em train.py.

    A ferramenta multihost_runner.py está otimizada para experiências rápidas, reutilizando repetidamente as mesmas TPUs. Uma vez que o script multihost_runner.py depende de ligações SSH de longa duração, não o recomendamos para tarefas de longa duração. Se quiser executar uma tarefa mais longa (por exemplo, durante horas ou dias), recomendamos que use o comando multihost_job.py.

    Neste tutorial, usamos o termo runner para indicar a máquina na qual executa o script multihost_runner.py. Usamos o termo trabalhadores para indicar as VMs de TPU que compõem as suas divisões. Pode executar o multihost_runner.py numa máquina local ou em qualquer VM do Compute Engine no mesmo projeto que as suas fatias. A execução de multihost_runner.py num trabalhador não é suportada.

    O multihost_runner.py liga-se automaticamente aos trabalhadores da TPU através de SSH.

    Neste exemplo, executa shardings.py em duas fatias v5e-16, um total de quatro VMs e 16 chips de TPU. Pode modificar o exemplo para ser executado em mais TPUs.

    Configure o seu ambiente

    1. Clone o MaxText na máquina do executor:

      $ git clone https://github.com/AI-Hypercomputer/maxtext

    2. Aceda ao diretório do repositório.

      $ cd maxtext

    3. Para criar chaves SSH para gcloud, recomendamos que deixe uma palavra-passe em branco (prima Enter duas vezes depois de executar o seguinte comando). Se lhe for pedido que o ficheiro google_compute_engine já existe, selecione a opção para não manter a sua versão existente.

        $ ssh-keygen -f ~/.ssh/google_compute_engine

    4. Adicione uma variável de ambiente para definir a contagem de fatias de TPU como 2.

        $ export SLICE_COUNT=2

    5. Crie um ambiente de várias divisões com o comando queued-resources create ou a consola Google Cloud .

      gcloud

      O comando seguinte mostra como criar uma TPU Multislice v5e. Para usar uma versão diferente da TPU, especifique um accelerator-type diferente e runtime-version.

      $ gcloud compute tpus queued-resources \
    create ${QR_ID} \
    --accelerator-type=${ACCELERATOR_TYPE} \
    --runtime-version=${RUNTIME_VERSION} \
    --node-count=${SLICE_COUNT} \
    --node-prefix=${TPU_NAME} \
    --zone=${ZONE} \
    [--reserved|--spot]

      Consola

      1. Na Google Cloud consola, aceda à página TPUs:

        Aceder às TPUs

      2. Clique em Criar TPU.

      3. No campo Nome, introduza um nome para a TPU.

      4. Na caixa Zona, selecione a zona onde quer criar a TPU.

      5. Na caixa Tipo de TPU, selecione um tipo de acelerador. O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que quer criar. O multislice só é suportado na TPU Cloud v4 e em versões posteriores da TPU. Para mais informações sobre as versões de TPU, consulte o artigo Versões de TPU.

      6. Na caixa Versão do software da TPU, selecione uma versão do software. Quando cria uma VM do Cloud TPU, a versão do software da TPU especifica a versão do tempo de execução da TPU a instalar nas VMs da TPU. Para mais informações, consulte o artigo Versões de software da TPU.

      7. Clique no botão de ativar/desativar Ativar colocação em fila.

      8. No campo Nome do recurso em fila, introduza um nome para o seu pedido de recurso em fila.

      9. Clique na caixa de verificação Make this a Multislice TPU (Transformar isto num TPU de vários segmentos).

      10. No campo Número de divisões, introduza o número de divisões que quer criar.

      11. Clique em Criar para criar o pedido de recursos em fila.

    6. Quando o aprovisionamento de recursos em fila começa, pode demorar até cinco minutos a ser concluído, consoante o tamanho do recurso em fila. Aguarde até que o recurso em fila esteja no estado ACTIVE. Pode verificar o estado de um pedido de recurso em fila de espera através da CLI gcloud ou da Google Cloud consola:

      gcloud

      $ gcloud compute tpus queued-resources list \
    --filter=${QR_ID} --zone=${ZONE} --project=${PROJECT}

      Isto deve gerar um resultado com o seguinte aspeto:

      NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central2-b  4           v5litepod-16             ACTIVE
...

      Consola

      1. Na Google Cloud consola, aceda à página TPUs:

        Aceder às TPUs

      2. Clique no separador Recursos em fila.

      3. Clique no nome do seu pedido de recurso em fila.

      Contacte o seu Google Cloud representante da conta se o estado do QR estiver no estado WAITING_FOR_RESOURCES ou PROVISIONING durante mais de 15 minutos.

    7. Instale dependências.

      $ python3 multihost_runner.py \
    --TPU_PREFIX=${TPU_NAME} \
    --ZONE=${ZONE} \
    --COMMAND="bash setup.sh"

    8. Execute shardings.py em cada trabalhador através do multihost_runner.py.

      $ python3 multihost_runner.py \
    --TPU_PREFIX=${TPU_NAME} \
    --ZONE=${ZONE} \
    --COMMAND="python3 -m pedagogical_examples.shardings \
    --dcn_data_parallelism ${SLICE_COUNT} \
    --ici_fsdp_parallelism 16 \
    --batch_size 131072 \
    --embedding_dimension 2048"

      Vai ver aproximadamente 230 TFLOPs por segundo de desempenho nos ficheiros de registo.

      Para mais informações sobre a configuração do paralelismo, consulte os artigos Fragmentação em vários segmentos com paralelismo da DCN e shardings.py.

    9. Limpe as UTPs e os recursos em fila quando terminar. A eliminação demora entre dois e cinco minutos. Se estiver a usar a CLI gcloud, pode executar este comando em segundo plano com a flag --async opcional.

    Escalar uma carga de trabalho para o Multislice

    Antes de executar o modelo num ambiente de vários segmentos, faça as seguintes alterações ao código:

    Estas devem ser as únicas alterações de código necessárias quando se muda para o Multislice. Para alcançar um elevado desempenho, a DCN tem de ser mapeada em eixos paralelos de dados, paralelos de dados totalmente fragmentados ou paralelos de pipeline. As considerações de desempenho e as estratégias de divisão são abordadas mais detalhadamente no artigo Divisão com várias fatias para máximo desempenho.

    Para validar que o seu código pode aceder a todos os dispositivos, pode afirmar que len(jax.devices()) é igual ao número de chips no seu ambiente Multislice. Por exemplo, se estiver a usar quatro fatias de v4-16, tem oito chips por fatia * 4 fatias, pelo que len(jax.devices()) deve devolver 32.

    Escolher os tamanhos das divisões para ambientes com várias divisões

    Para obter um aumento linear da velocidade, adicione novas divisões do mesmo tamanho que a divisão existente. Por exemplo, se usar uma fatia v4-512, o Multislice alcança aproximadamente o dobro do desempenho adicionando uma segunda fatia v4-512 e duplicando o tamanho do lote global. Para mais informações, consulte o artigo Partição com várias fatias para máximo desempenho.

    Executar o seu trabalho em várias divisões

    Existem três abordagens diferentes para executar a sua carga de trabalho personalizada num ambiente de vários segmentos:

    1. Usando o script do executor de experiências, multihost_runner.py
    2. Usando o script do executor de produção, multihost_job.py
    3. Usar uma abordagem manual

    Script do executor de experiências

    O script multihost_runner.py distribui código a um ambiente Multislice existente e executa o seu comando em cada anfitrião, copia os seus registos e acompanha o estado de erro de cada comando. O script multihost_runner.py está documentado no MaxText README.

    Uma vez que o multihost_runner.py mantém ligações SSH persistentes, só é adequado para experiências de tamanho modesto e duração relativamente curta. Pode adaptar os passos no multihost_runner.pytutorial à sua carga de trabalho e configuração de hardware.

    Script do executor de produção

    Para tarefas de produção que precisam de resiliência contra falhas de hardware e outras interrupções, é melhor integrar diretamente com a API Create Queued Resource (criar recurso em fila). Use multihost_job.py como exemplo funcional que aciona a chamada da API de recursos em fila com o script de arranque adequado para executar o seu treino e retomá-lo em caso de preemptção. O script está documentado no multihost_job.py MaxText README.

    Uma vez que o multihost_job.py tem de aprovisionar recursos para cada execução, não oferece um ciclo de iteração tão rápido quanto o multihost_runner.py.

    Abordagem manual

    Recomendamos que use ou adapte multihost_runner.py ou multihost_job.py para executar a sua carga de trabalho personalizada na configuração de várias fatias. No entanto, se preferir aprovisionar e gerir o seu ambiente através de comandos QR diretamente, consulte o artigo Gerir um ambiente com várias divisões.

    Faça a gestão de um ambiente Multislice

    Para aprovisionar e gerir manualmente QRs sem usar as ferramentas disponibilizadas no repositório MaxText, leia as secções seguintes.

    Crie recursos em fila

    gcloud

    1. Crie um pedido de recursos em fila usando o seguinte comando:

      $ gcloud compute tpus queued-resources \
    create ${QR_ID} \
    --project=${PROJECT} \
    --zone=${ZONE} \
    --node-count=${SLICE_COUNT} \
    --accelerator-type=${ACCELERATOR_TYPE} \
    --runtime-version=${RUNTIME_VERSION} \
    --network=${NETWORK_NAME} \
    --subnetwork=${SUBNETWORK_NAME} \
    --tags=${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \
    --metadata=startup-script="${STARTUP_SCRIPT}" \
    [--reserved|--spot]

    Certifique-se de que tem a quota respetiva antes de selecionar --reserved, --spot ou a quota a pedido predefinida. Para obter informações sobre os tipos de quotas, consulte a Política de Quotas.

    curl

    1. Crie um ficheiro denominado queued-resource-req.json e copie o seguinte JSON para o mesmo.

      {
"guaranteed": { "reserved": true },
"tpu": {
    "node_spec": [
    {
    "parent": "projects/your-project-number/locations/your-zone",
        "node": {
        "accelerator_type": "accelerator-type",
        "runtime_version": "tpu-vm-runtime-version",
        "network_config": {
            "network": "your-network-name",
            "subnetwork": "your-subnetwork-name",
            "enable_external_ips": true
        },
        "tags" : ["example-tag-1"]
        "metadata": {
            "startup-script": "your-startup-script"
        }
    },
    "multi_node_params": {
        "node_count": slice-count,
        "node_id_prefix": "your-queued-resource-id"
    }
    }
    ]
}
}

      Substitua os seguintes valores:

      • your-project-number - O seu Google Cloud número do projeto
      • your-zone - A zona na qual quer criar o recurso em fila
      • accelerator-type - A versão e o tamanho de uma única fatia. O Multislice só é suportado na Cloud TPU v4 e em versões posteriores da TPU.
      • tpu-vm-runtime-version: a versão de tempo de execução da VM da TPU que quer usar.
      • your-network-name - Opcional, uma rede à qual o recurso em fila vai ser anexado
      • your-subnetwork-name - Opcional, uma sub-rede à qual o recurso em fila vai ser anexado
      • example-tag-1: opcional, uma string de etiqueta arbitrária
      • your-startup-script: um script de arranque que é executado quando o recurso em fila é atribuído
      • slice-count: o número de fatias de TPU no seu ambiente de várias fatias
      • your-queued-resource-id - O ID fornecido pelo utilizador para o recurso em fila

      Para mais informações, consulte a documentação da API REST de recursos em fila para ver todas as opções disponíveis.

      Para usar a capacidade de spot, substitua:

      "guaranteed": { "reserved": true } com "spot": {}

      Remova a linha para usar a capacidade a pedido predefinida.

    2. Envie o pedido de criação de recursos em fila com o payload JSON:

      $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
-H "Content-Type: application/json" \
-d @queuedresourcereq.json \
https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-queued-resource-id

      Substitua os seguintes valores:

      • your-project-id - O ID do seu Google Cloud projeto
      • your-zone - A zona na qual quer criar o recurso em fila
      • your-queued-resource-id - O ID fornecido pelo utilizador para o recurso em fila

    A resposta deve ter o seguinte aspeto:

    {
"name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
"metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
},
"done": false
}

    Use o valor GUID no final do valor da string do atributo name para obter informações sobre o pedido de recurso em fila.

    Consola

    1. Na Google Cloud consola, aceda à página TPUs:

      Aceder às TPUs

    2. Clique em Criar TPU.

    3. No campo Nome, introduza um nome para a TPU.

    4. Na caixa Zona, selecione a zona onde quer criar a TPU.

    5. Na caixa Tipo de TPU, selecione um tipo de acelerador. O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que quer criar. O Multislice só é suportado na Cloud TPU v4 e em versões posteriores da TPU. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.

    6. Na caixa Versão do software da TPU, selecione uma versão do software. Quando cria uma VM do Cloud TPU, a versão do software da TPU especifica a versão do tempo de execução da TPU a instalar. Para mais informações, consulte o artigo Versões de software da TPU.

    7. Clique no botão de ativar/desativar Ativar colocação em fila.

    8. No campo Nome do recurso em fila, introduza um nome para o seu pedido de recurso em fila.

    9. Clique na caixa de verificação Make this a Multislice TPU (Transformar isto num TPU de vários segmentos).

    10. No campo Número de divisões, introduza o número de divisões que quer criar.

    11. Clique em Criar para criar o pedido de recursos em fila.

    Obtenha o estado de um recurso em fila

    gcloud

    $ gcloud compute tpus queued-resources describe ${QR_ID} --zone=${ZONE}

    Para um recurso em fila que se encontra no estado ACTIVE, o resultado tem o seguinte aspeto:

    ...
state:
    state: ACTIVE
...

    curl

    $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${YOUR_QR_ID}

    Para um recurso em fila que se encontra no estado ACTIVE, o resultado tem o seguinte aspeto:

    {
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

    Consola

    1. Na Google Cloud consola, aceda à página TPUs:

      Aceder às TPUs

    2. Clique no separador Recursos em fila.

    3. Clique no nome do seu pedido de recurso em fila.

    Depois de o TPU ter sido aprovisionado, também pode ver detalhes sobre o pedido de recursos em fila de espera acedendo à página TPUs, encontrando o seu TPU e clicando no nome do pedido de recursos em fila de espera correspondente.

    Num cenário raro, pode encontrar o recurso em fila no estado FAILED enquanto algumas fatias estão ACTIVE. Se isto acontecer, elimine os recursos criados e tente novamente dentro de alguns minutos ou contacte o Google Cloud apoio técnico.

    SSH e instalar dependências

    Execute código JAX em fatias de TPU descreve como estabelecer ligação às suas VMs de TPU através de SSH numa única fatia. Para estabelecer ligação a todas as VMs de TPU no seu ambiente de vários fragmentos através de SSH e instalar dependências, use o seguinte comando gcloud:

      $ gcloud compute tpus queued-resources ssh ${QR_ID} \
        --zone=${ZONE} \
        --node=all \
        --worker=all \
        --command="command-to-run" \
        --batch-size=4

    Este comando gcloud envia o comando especificado a todos os trabalhadores e nós no QR através de SSH. O comando é processado em lotes de quatro e enviado em simultâneo. O lote de comandos seguinte é enviado quando o lote atual conclui a execução. Se ocorrer uma falha com um dos comandos, o processamento é interrompido e não são enviados mais lotes. Para mais informações, consulte a referência da API de recursos em fila. Se o número de divisões que está a usar exceder o limite de processamento simultâneo do seu computador local (também denominado limite de processamento em lote), vai deparar-se com um impasse. Por exemplo, suponha que o limite de processamento em lote na sua máquina local é 64. Se tentar executar um script de preparação em mais de 64 fatias, por exemplo, 100, o comando SSH divide as fatias em lotes. Executa o script de preparação no primeiro lote de 64 fatias e aguarda a conclusão dos scripts antes de executar o script no lote restante de 36 fatias. No entanto, o primeiro lote de 64 fatias não pode ser concluído até que as 36 fatias restantes comecem a executar o script, o que causa um impasse.

    Para evitar este cenário, pode executar o script de preparação em segundo plano em cada VM anexando um e comercial (&) ao comando do script que especificar com a sinalização --command. Quando o faz, depois de iniciar o script de preparação no primeiro lote de fatias, o controlo é imediatamente devolvido ao comando SSH. O comando SSH pode, então, começar a executar o script de preparação no lote restante de 36 fatias. Tem de encaminhar os streams stdout e stderr adequadamente quando executar os comandos em segundo plano. Para aumentar o paralelismo no mesmo QR, pode selecionar fatias específicas através do parâmetro --node.

    Configuração de rede

    Certifique-se de que as fatias de TPU conseguem comunicar entre si seguindo os passos abaixo. Instale o JAX em cada uma das divisões. Para mais informações, consulte o artigo Execute código JAX em fatias de TPU. Afirme que len(jax.devices()) é igual ao número de chips no seu ambiente Multislice. Para o fazer, em cada fatia, execute:

      $ python3 -c 'import jax; print(jax.devices())'

    Se executar este código em quatro fatias de v4-16, existem oito chips por fatia e quatro fatias, devendo ser devolvidos um total de 32 chips (dispositivos) por jax.devices().

    Liste os recursos em fila

    gcloud

    Pode ver o estado dos recursos em fila de espera através do comando queued-resources list:

    $ gcloud compute tpus queued-resources list --zone=${ZONE}

    O resultado tem um aspeto semelhante ao seguinte:

    NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central1-a  4           v5litepod-16             ACTIVE
...

    Consola

    1. Na Google Cloud consola, aceda à página TPUs:

      Aceder às TPUs

    2. Clique no separador Recursos em fila.

    Inicie a tarefa num ambiente aprovisionado

    Pode executar manualmente cargas de trabalho ligando-se a todos os anfitriões em cada fatia através de SSH e executando o seguinte comando em todos os anfitriões.

    $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --worker=all \
    --command="command-to-run"

    Repor respostas rápidas

    A API ResetQueuedResource pode ser usada para repôr todas as VMs num QR ACTIVE. A reposição das VMs apaga à força a memória da máquina e repõe a VM para o respetivo estado inicial. Todos os dados armazenados localmente vão permanecer intactos, e o script de arranque vai ser invocado após uma reposição. A API ResetQueuedResource pode ser útil quando quer reiniciar todas as TPUs. Por exemplo, quando a preparação está bloqueada e repor todas as VMs é mais fácil do que depurar.

    As reposições de todas as VMs são realizadas em paralelo e uma operação demora 1 a 2 minutos a ser concluída.ResetQueuedResource Para invocar a API, use o seguinte comando:

    $ gcloud compute tpus queued-resources reset ${QR_ID} --zone=${ZONE}

    A eliminar recursos em fila

    Para libertar recursos no final da sessão de formação, elimine o recurso em fila. A eliminação demora entre dois e cinco minutos. Se estiver a usar a CLI gcloud, pode executar este comando em segundo plano com a flag --async opcional.

    gcloud

    $ gcloud compute tpus queued-resources \
    delete ${QR_ID} --zone=${ZONE} --force [--async]

    Consola

    1. Na Google Cloud consola, aceda à página TPUs:

      Aceder às TPUs

    2. Clique no separador Recursos em fila.

    3. Selecione a caixa de verificação junto ao pedido de recurso em fila.

    4. Clique em Eliminar.

    Recuperação automática de falhas

    Em caso de interrupção, o Multislice oferece uma reparação sem intervenção da fatia afetada e a reposição de todas as fatias posteriormente. A fatia afetada é substituída por uma nova fatia e as fatias restantes, que estão de outra forma em bom estado, são repostas. Se não estiver disponível capacidade para atribuir uma fatia de substituição, a preparação para.

    Para retomar a preparação automaticamente após uma interrupção, tem de especificar um script de arranque que verifique e carregue os pontos de verificação guardados mais recentemente. O seu script de arranque é executado automaticamente sempre que uma fatia é reatribuída ou uma VM é reposta. Especifica um script de arranque no payload JSON que envia para a API de pedido de criação de QR.

    O seguinte script de arranque (usado em Criar QRs) permite-lhe recuperar automaticamente de falhas e retomar a preparação a partir de pontos de verificação armazenados num contentor do Cloud Storage durante a preparação do MaxText:

    {
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 -m MaxText.train MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

    Clone o repositório MaxText antes de experimentar esta funcionalidade.

    Criação de perfis e depuração

    A criação de perfis é igual em ambientes de fatia única e de várias fatias. Para mais informações, consulte o artigo Criar perfis de programas JAX.

    Otimizar a formação

    As secções seguintes descrevem como pode otimizar o treino de Multislice.

    Partição com várias fatias para máximo desempenho

    Para alcançar o máximo desempenho em ambientes de divisão múltipla, tem de considerar como dividir os dados em vários segmentos. Normalmente, existem três opções (paralelismo de dados, paralelismo de dados totalmente fragmentados e paralelismo de pipeline). Não recomendamos a divisão de ativações pelas dimensões do modelo (por vezes, denominada paralelismo de tensores) porque requer demasiada largura de banda entre fatias. Para todas estas estratégias, pode manter a mesma estratégia de divisão em fragmentos numa fatia que funcionou para si no passado.

    Recomendamos que comece com o paralelismo de dados puro. A paralelização de dados totalmente fragmentados é útil para libertar a utilização de memória. A desvantagem é que a comunicação entre as fatias usa a rede DCN e vai tornar a sua carga de trabalho mais lenta. Use o paralelismo de pipeline apenas quando necessário com base no tamanho do lote (conforme analisado abaixo).

    Quando usar o paralelismo de dados

    O paralelismo de dados puro funciona bem nos casos em que tem uma carga de trabalho que está a ser executada bem, mas quer melhorar o respetivo desempenho ao fazer o escalamento em várias divisões.

    Para alcançar uma forte escalabilidade em várias divisões, o tempo necessário para realizar a redução total no DCN tem de ser inferior ao tempo necessário para realizar uma passagem para trás. A DCN é usada para a comunicação entre fatias e é um fator limitativo no débito da carga de trabalho.

    Cada chip da TPU v4 tem um desempenho máximo de 275 * 1012 FLOPS por segundo.

    Existem quatro chips por anfitrião de TPU e cada anfitrião tem uma largura de banda de rede máxima de 50 Gbps.

    Isto significa que a intensidade aritmética é 4 * 275 * 1012 FLOPS / 50 Gbps = 22 000 FLOPS / bit.

    O seu modelo vai usar 32 a 64 bits de largura de banda da DCN para cada parâmetro por passo. Se usar duas divisões, o modelo usa 32 bits de largura de banda da DCN. Se usar mais de duas divisões, o compilador executa uma operação de redução total de mistura completa e usa até 64 bits de largura de banda da DCN para cada parâmetro por passo. A quantidade de FLOPS necessária para cada parâmetro varia consoante o seu modelo. Especificamente, para modelos de linguagem baseados em transformadores, o número de FLOPS necessários para uma passagem para a frente e para trás é de aproximadamente 6 * B * P, onde:

    • B é o tamanho do lote em tokens
    • P é o número de parâmetros

    O número de FLOPS por parâmetro é 6 * B e o número de FLOPS por parâmetro durante a passagem para trás é 4 * B.

    Para garantir uma forte escalabilidade em várias divisões, certifique-se de que a intensidade operacional excede a intensidade aritmética do hardware da TPU. Para calcular a intensidade operacional, divida o número de FLOPS por parâmetro durante a passagem para trás pela largura de banda da rede (em bits) por parâmetro por passo: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

    Por conseguinte, para um modelo de linguagem baseado em transformadores, se estiver a usar duas divisões: Operational intensity = 4 * B / 32

    Se estiver a usar mais de duas divisões: Operational intensity = 4 * B/64

    Isto sugere um tamanho mínimo do lote entre 176 mil e 352 mil para modelos de linguagem baseados em transformadores. Uma vez que a rede DCN pode perder pacotes brevemente, é melhor manter uma margem de erro significativa, implementando o paralelismo de dados apenas se o tamanho do lote por Pod for, pelo menos, de 350 mil (dois Pods) a 700 mil (muitos Pods).

    Para outras arquiteturas de modelos, tem de estimar o tempo de execução da passagem para trás por fatia (medindo o tempo com um perfilador ou contando FLOPS). Em seguida, pode comparar esse valor com o tempo de execução esperado para reduzir tudo no DCN e ter uma boa estimativa de se o paralelismo de dados faz sentido para si.

    Quando usar o paralelismo de dados totalmente fragmentado (FSDP)

    O paralelismo de dados totalmente fragmentado (FSDP) combina o paralelismo de dados (fragmentando os dados em vários nós) com a fragmentação dos pesos em vários nós. Para cada operação nas passagens para a frente e para trás, os pesos são reunidos para que cada fatia tenha os pesos de que precisa. Em vez de sincronizar os gradientes através da redução total, os gradientes são dispersos à medida que são produzidos. Desta forma, cada fatia só recebe os gradientes dos pesos pelos quais é responsável.

    Semelhante ao paralelismo de dados, o FSDP requer o dimensionamento do tamanho do lote global de forma linear com o número de divisões. A FSDP diminui a pressão da memória à medida que aumenta o número de divisões. Isto deve-se ao facto de o número de ponderações e o estado do otimizador por fatia diminuírem, mas ao preço de um aumento do tráfego de rede e de uma maior possibilidade de bloqueio devido a um coletivo atrasado.

    Na prática, o FSDP em fatias é melhor se estiver a aumentar o lote por fatia, a armazenar mais ativações para minimizar a rematerialização durante a passagem para trás ou a aumentar o número de parâmetros na sua rede neural.

    As operações all-gather e all-reduce no FSDP funcionam de forma semelhante às do DP, pelo que pode determinar se a sua carga de trabalho do FSDP está limitada pelo desempenho da DCN da mesma forma descrita na secção anterior.

    Quando usar o paralelismo de pipeline

    O paralelismo de pipeline torna-se relevante quando alcança um elevado desempenho com outras estratégias de paralelismo que requerem um tamanho do lote global superior ao tamanho do lote máximo preferido. O paralelismo de pipeline permite que as divisões que compõem um pipeline "partilhem" um lote. No entanto, o paralelismo de pipeline tem duas desvantagens significativas:

    1. Incorre na "bolha de pipeline", em que os chips estão inativos porque estão à espera de dados.
    2. Requer micro-lotes, o que diminui o tamanho do lote efetivo, a intensidade aritmética e, em última análise, a utilização de FLOPs do modelo.

    O paralelismo de pipeline só deve ser usado se as outras estratégias de paralelismo exigirem um tamanho do lote global demasiado grande. Antes de experimentar o paralelismo de pipelines, vale a pena experimentar para ver empiricamente se a convergência por amostra abranda no tamanho do lote necessário para alcançar um FSDP de alto desempenho. A FSDP tende a alcançar uma utilização de FLOPs do modelo mais elevada, mas se a convergência por amostra abrandar à medida que o tamanho do lote aumenta, o paralelismo de pipeline pode continuar a ser a melhor escolha. A maioria das cargas de trabalho pode tolerar tamanhos de lotes suficientemente grandes para não beneficiar do paralelismo de pipeline, mas a sua carga de trabalho pode ser diferente.

    Se o paralelismo de pipelines for necessário, recomendamos que o combine com o paralelismo de dados ou o FSDP. Isto permite-lhe minimizar a profundidade do pipeline enquanto aumenta o tamanho do lote por pipeline até que a latência da DCN se torne menos relevante no débito. Concretamente, se tiver N divisões, considere pipelines de profundidade 2 e N/2 réplicas de paralelismo de dados, depois pipelines de profundidade 4 e N/4 réplicas de paralelismo de dados, continuando da mesma forma, até que o lote por pipeline seja suficientemente grande para que os coletivos DCN possam ser ocultados atrás da aritmética na passagem para trás. Isto minimiza a desaceleração introduzida pelo paralelismo de pipeline, ao mesmo tempo que lhe permite expandir para além do limite do tamanho do lote global.

    Práticas recomendadas para os multislice

    As secções seguintes descrevem as práticas recomendadas para a preparação com várias fatias.

    Carregamento de dados

    Durante a preparação, carregamos repetidamente lotes de um conjunto de dados para introduzir no modelo. Ter um carregador de dados assíncrono eficiente que divida o lote em vários anfitriões é importante para evitar que as TPUs fiquem sem trabalho. O carregador de dados atual no MaxText faz com que cada anfitrião carregue um subconjunto igual dos exemplos. Esta solução é adequada para texto, mas requer uma nova divisão no modelo. Além disso, o MaxText ainda não oferece a criação de instantâneos determinísticos, o que permitiria ao iterador de dados carregar os mesmos dados antes e depois da preempção.

    Criação de pontos de restauro

    A biblioteca de pontos de verificação Orbax fornece primitivas para pontos de verificação de JAX PyTrees para armazenamento local ou Google Cloud armazenamento. Disponibilizamos uma integração de referência com a criação de pontos de verificação síncronos no MaxText em checkpointing.py.

    Configurações suportadas

    As secções seguintes descrevem as formas de divisão suportadas, a orquestração, as frameworks e o paralelismo para o Multislice.

    Formas

    Todas as fatias têm de ter a mesma forma (por exemplo, o mesmo AcceleratorType). As formas de fatias heterogéneas não são suportadas.

    Orquestração

    A orquestração é suportada com o GKE. Para mais informações, consulte TPUs no GKE.

    Frameworks

    O Multislice só suporta cargas de trabalho JAX e PyTorch.

    Paralelismo

    Recomendamos que os utilizadores testem o Multislice com paralelismo de dados. Para saber mais sobre a implementação do paralelismo de pipeline com o Multislice, contacte o seu Google Cloud representante da conta.

    Apoio técnico e feedback

    Agradecemos todo o feedback! Para partilhar feedback ou pedir apoio técnico, contacte-nos através do formulário de apoio técnico ou feedback do Cloud TPU.