Inferência do Cloud TPU v5e [pré-lançamento público]

Começar

Visão geral e benefícios

O Cloud TPU v5e é o acelerador de IA de última geração do Google Cloud. Com um tamanho menor de 256 chips por pod, os pods v5e são otimizados para treinamento, ajuste e exibição baseados em transformadores, de texto para imagem e em CNN.

conceitos

Se você nunca usou Cloud TPUs, confira a página inicial de documentos da TPU.

Ícones

Há 256 chips em um único v5e com 8 chips por host. Consulte Arquitetura do sistema para saber mais detalhes.

Núcleos

Os chips de TPU têm um ou dois TensorCores para executar a multiplicação de matrizes. Semelhante aos pods da v2 e v3, a v5e tem um TensorCore por chip. Por outro lado, os pods v4 têm dois TensorCores por chip. Consulte Arquitetura do sistema para ver mais detalhes sobre o v5e TensorCores. Outras informações sobre o TensorCores podem ser encontradas neste artigo do ACM.

Host

Um host é um computador físico (CPU) que executa VMs. Um host pode executar várias VMs ao mesmo tempo.

Inferência em lote

Inferência em lote ou off-line refere-se a fazer inferência fora dos pipelines de produção, normalmente em uma massa de entradas. A inferência em lote é usada para tarefas off-line, como rotulagem de dados, e para avaliação do modelo treinado. Os SLOs de latência não são uma prioridade para inferência em lote.

Disponibilização

Exibição refere-se ao processo de implantação de um modelo de machine learning treinado em um ambiente de produção, onde ele pode ser usado para fazer previsões ou decisões. Os SLOs de latência são uma prioridade para a disponibilização.

Host único vs. vários hosts

Slices com menos de oito ícones usam no máximo um host. Slices com mais de oito ícones têm acesso a mais de um host e podem executar treinamentos distribuídos usando vários hosts. A visualização não é compatível com a exibição em vários hosts. O número máximo de ícones permitidos em um único job de disponibilização é de oito ícones.

Recurso na fila

Uma representação de recursos de TPU, usada para enfileirar e gerenciar uma solicitação para um ambiente de TPU de fatia única ou de várias fatias. Consulte o guia do usuário "Recursos na fila" para mais informações.

VM da Cloud TPU

Uma máquina virtual executando Linux que tem acesso às TPUs subjacentes. Para TPUs v5e, cada VM de TPU tem acesso direto a 1, 4 ou 8 chips, dependendo do tipo de acelerador especificado pelo usuário. Uma VM de TPU também é conhecida como worker.

Arquitetura do sistema

Cada chip v5e contém um TensorCore. Cada TensorCore tem quatro unidades de multiplicação de matriz (MXU, na sigla em inglês), uma unidade vetorial e uma unidade escalar. A tabela a seguir mostra as especificações das chaves e seus valores para uma v5e. As especificações do pod estão incluídas na tabela que segue as especificações do ícone.

Especificações do chip de chave Valores de v5e
Pico de computação por chip (bf16) 197 TFLOPs
Pico de computação por chip (Int8) 393 TFLOPs
Capacidade de HBM2 e largura de banda 16 GB, 819 GBps
Interchip Interconnect BW 1.600 Gbps
Largura total reduzida por pod 51,2 TB/s
Largura de banda de bisseção por pod 1,6 TB/s
Largura de banda da rede de data centers por pod 6,4 Tbps
Principais especificações do pod Valores de v5e
Tamanho do Pod de TPU 256 chips
Topologia de interconexão Torus 2D
Pico de computação por pod 100 PetaOps(Int8)
Largura total reduzida por pod 51,2 TB/s
Largura de banda de bisseção por pod 1,6 TB/s
Largura de banda da rede de data centers por pod 6,4 Tbps

Chip TPU v5e

O diagrama a seguir ilustra um chip TPU v5e.

desenho

Tipos de acelerador

O Cloud TPU v5e é um produto combinado de treinamento e inferência. A sinalização AcceleratorType é usada para diferenciar entre o ambiente de TPU provisionado para treinamento e o ambiente provisionado para exibição. Os jobs de treinamento são otimizados para capacidade e disponibilidade, enquanto os jobs de exibição são otimizados para latência. Portanto, um job de treinamento em TPUs provisionadas para disponibilização pode ter menor disponibilidade e, da mesma forma, um job de veiculação executado em TPUs provisionadas para treinamento pode ter maior latência.

AcceleratorType usa o número de TensorCores no nó para descrever o tamanho da fatia. AcceleratorType é uma string formatada "v$VERSION_NUMBER-$CORES_COUNT".

As seguintes formas de fatias 2D são compatíveis com a v5e:

topologia Número de chips do TPU Número de hosts
1x1 1 1/8
2x2 4 1/2
2x4 8 1
4x4 16 2
4x8 32 4
8x8 64 8
8x16 128 16
16x16 256 32

Tipos do Cloud TPU v5e para exibição

Há tipos de TPU v5e com melhor disponibilidade para veiculação. Esses tipos oferecem suporte a até oito chips v5e (host único). As seguintes configurações são aceitas: fatias 1 x 1, 2 x 2 e 2 x 4. Cada fatia tem 1, 4 e 8 ícones, respectivamente.

Configurações do \TPU v5e compatíveis com a disponibilização: 1x1, 2x2 e 2x4.

A veiculação em mais de 8 chips v5e, ou seja, exibição em vários hosts, está nos planos da Cloud TPU e estará disponível posteriormente.

Para provisionar TPUs para um job de exibição, use os seguintes tipos de acelerador na solicitação de criação de TPU de API ou CLI. Toda a capacidade de disponibilização do v5e está atualmente na zona us-west4-a:

  • v5e-1
  • v5e-4
  • v5e-8

Tipos de VM

Cada VM em uma fatia da TPU pode conter de 1, 4 ou 8 ícones. Fatias menores e com 4 chips terão o mesmo nó NUMA (para saber mais sobre nós NUMA, consulte Arquitetura do sistema). Para VMs de TPU com 8 chips, a comunicação entre CPU e TPU será mais eficiente nas partições NUMA. Por exemplo, na figura abaixo, a comunicação CPU0-Chip0 é mais rápida que a comunicação CPU0-Chip4.

VM de TPU com 8 chips contendo duas partições NUMA

Comparação de tipos de VM:

Para o pré-lançamento público, os clientes de inferência têm acesso a um chip (v5e-1), quatro chips (v5e-4) e 8 chips (v5e-8).

Tipo de VM n2d-48-24-v5lite-tpu (link em inglês) n2d-192-112-v5lite-tpu (link em inglês) n2d-384-224-v5lite-tpu (link em inglês)
No de ícones v5e 1 4 8
No de vCPUs 24 112 224
RAM (GB) 48 192 384
No de nós NUMA 1 1 2
Aplicável a v5e-1 v5e-4 v5e-8
Interrupção Alta Meio Baixa

Para liberar espaço para VMs com mais chips, os programadores podem forçar a interrupção de VMs com menos chips. Portanto, é provável que as VMs de 8 chips saiam de VMs de 1 e 4 chips.

Começar

Como proteger a capacidade

O Cloud TPU v5e agora está em Acesso antecipado. Entre em contato com a equipe de vendas do Cloud para começar a usar o Cloud TPU v5e nas suas cargas de trabalho de IA.

Preparar um projeto do Google Cloud

  1. Faça login na sua Conta do Google. Crie uma nova conta se ainda não tiver feito isso.

  2. No Console do Cloud, selecione ou crie um projeto do Cloud na página do seletor de projetos.

  3. A configuração de faturamento é necessária para todo o uso do Google Cloud. Portanto, verifique se o faturamento está ativado para seu projeto.

    O faturamento para uso do Acesso antecipado segue os preços regionais padrão mostrados na página de preços da Cloud TPU.

  4. Instale o gcloud alpha components.

  5. Ative a API TPU usando o comando gcloud a seguir no Cloud Shell. Também é possível ativá-lo no console do Google Cloud.

    gcloud services enable tpu.googleapis.com
    
  6. Ative a conta de serviço da TPU.

    As contas de serviço permitem que o serviço do Cloud TPU acesse outros serviços do Google Cloud. Uma conta de serviço gerenciado pelo usuário é uma prática recomendada do Google Cloud. Siga estes guias para criar e conceder os papéis a seguir à sua conta de serviço. Os seguintes papéis são necessários:

    • Administrador da TPU
    • Administrador de armazenamento
    • Gravador de registros
    • Gravador de métricas do Monitoring
  7. Configurar o projeto e a zona.

    O ID do projeto é o nome do projeto mostrado no Console do Cloud. A zona padrão do Cloud TPU v5e é us-west4-a.

    export PROJECT_ID=project-ID
    export ZONE=us-west4-a
    
    gcloud alpha compute tpus tpu-vm service-identity create --zone=${ZONE}
    
    gcloud auth login
    gcloud config set project ${PROJECT}
    gcloud config set compute/zone ${ZONE}
    
  8. Solicitar acesso aos artefatos de pré-lançamento público de inferência da Cloud TPU.

    Envie a conta de serviço que você planeja usar para a VM da TPU à equipe da TPU. A padrão é a conta de serviço padrão do Compute Engine (${PROJECT_ID}-compute@developer.gserviceaccount.com), que pode ser encontrada na página do IAM.

  9. Provisionar o ambiente do Cloud TPU v5e.

    Uma v5e é gerenciada como um recurso na fila. A capacidade pode ser provisionada usando o comando queued-resource create.

    Crie variáveis de ambiente para ID do projeto, tipo de acelerador, zona, versão de ambiente de execução e nome da TPU.

    export PROJECT_ID=project_ID
    export ACCELERATOR_TYPE=v5e-1
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=service_account
    export TPU_NAME=tpu-name
    export QUEUED_RESOURCE_ID=queued_resource_id
    

    Descrições de variáveis

    PROJECT_ID
    Nome do projeto. Use o nome do projeto do Google.
    ACCELERATOR_TYPE
    Consulte a seção Tipos de acelerador para ver os tipos compatíveis.
    ZONE
    Toda a capacidade do Acesso antecipado vai estar em us-west4-a.
    RUNTIME_VERSION
    v2-alpha-tpuv5-lite
    SERVICE_ACCOUNT
    Este é o endereço de e-mail da conta de serviço que pode ser encontrado no Console do Google Cloud -> IAM -> Contas de serviço Por exemplo: tpu-service-account@myprojectID.iam.gserviceaccount.com.
    TPU_NAME
    O ID atribuído pelo usuário da TPU criado quando a solicitação de recurso na fila é alocada.
    QUEUED_RESOURCE_ID
    O ID atribuído pelo usuário da solicitação de recurso na fila.
  10. Criar um recurso de TPU.

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT} \
    --reserved
    

    Se você quiser excluir o recurso reservado, será necessário excluir primeiro o TPU_NAME e depois a solicitação de recurso na fila.

    gcloud alpha compute tpus delete $TPU_NAME --zone ${ZONE} --project ${PROJECT_ID}
    gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
    --project ${PROJECT_ID} \
    --zone ${ZONE}
    
  11. SSH em VMs de TPU

    Para executar o código nas suas VMs de TPU, é preciso inserir SSH em cada uma delas. Neste exemplo, com uma v5e-1, há apenas uma VM de TPU.

    gcloud compute config-ssh
    gcloud compute tpus tpu-vm ssh $TPU_NAME --zone $ZONE --project $PROJECT_ID
    

Gerenciar VMs de TPU

Para todas as opções de gerenciamento de TPU para suas VMs de TPU, consulte Como gerenciar TPUs.

Desenvolver e executar

Nesta seção, descrevemos o processo geral de configuração para inferência de modelos personalizados usando JAX ou PyTorch no Cloud TPU v5e. O suporte ao TensorFlow será ativado em breve.

Para ver instruções de treinamento da v5e, consulte o guia de treinamento da v5e.

Como executar inferência na v5e

Pilha de software de inferência

Os detalhes da pilha de inferência SW são abordados nas seções a seguir. Este documento se concentra na exibição de host único para modelos treinados com JAX, TensorFlow (TF) e PyTorch.

Nesta seção, presumimos que você já tenha configurado seu projeto do Google Cloud de acordo com as instruções em Preparar um projeto do Google Cloud.

Inferência e exibição do modelo JAX

Na seção a seguir, explicamos o fluxo de trabalho da inferência do modelo JAX. Há dois caminhos para a inferência do JAX, como mostrado no diagrama. Esta seção aborda o caminho de produção de modelos JAX usando o jax2tf e o TensorFlow Serving.

  1. Usar jax2tf para converter o modelo para o TensorFlow 2 e salvá-lo
  2. Usar o conversor de inferência para converter o modelo salvo
  3. Usar o TensorFlow Serving para disponibilizar o modelo

Usar jax2tf para converter e salvar o modelo

Consulte Interoperação do JAX e do TensorFlow para converter e salvar seu modelo JAX no TensorFlow.

# Inference function
def model_jax(params, inputs):
  return params[0] + params[1] * inputs

# Wrap the parameter constants as tf.Variables; this will signal to the model
# saving code to save those constants as variables, separate from the
# computation graph.
params_vars = tf.nest.map_structure(tf.Variable, params)

# Build the prediction function by closing over the `params_vars`. If you
# instead were to close over `params` your SavedModel would have no variables
# and the parameters will be included in the function graph.
prediction_tf = lambda inputs: jax2tf.convert(model_jax)(params_vars, inputs)

my_model = tf.Module()
# Tell the model saver what the variables are.
my_model._variables = tf.nest.flatten(params_vars)
my_model.f = tf.function(prediction_tf, jit_compile=True, autograph=False)
tf.saved_model.save(my_model)

Usar o conversor de inferência para converter o modelo salvo

As etapas para conversores de inferência são descritas no Guia de conversores de inferência.

Usar o TensorFlow Serving

As etapas do TensorFlow Serving são descritas em TensorFlow Serving.

Exemplo de exibição do modelo JAX E2E

Pré-requisito:

Você precisa configurar suas credenciais do Docker e extrair a imagem do Docker "Conversor de inferência" e "TensorFlow Serving". Execute os comandos a seguir caso ainda não tenha feito isso:

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
    us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Faça o download do código de demonstração: SSH na VM da TPU e instale o código de demonstração de inferência.

gsutil -m cp -r \
  "gs://cloud-tpu-inference-public/demo" \
  .

Instale as dependências de demonstração do JAX na VM da TPU, instale requirements.txt.

pip install -r ./demo/jax/requirements.txt

Executar a demonstração da exibição do JAX BERT E2E

O modelo BERT pré-treinado (em inglês) é da Hugging Face.

  1. Exporte um modelo salvo do TF2 compatível com TPU de um modelo Flax BERT:

    cd demo/jax/bert
    
    python3 export_bert_model.py
    
  2. Inicie o contêiner do servidor do modelo do TensorFlow para o modelo:

    docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=/tmp/jax/bert_tpu,target=/models/bert \
     -e MODEL_NAME=bert \
     us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
    

    Verifique o registro do contêiner do servidor de modelo e se o gRPC e o servidor HTTP estão ativos:

    CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
    docker logs ${CONTAINER_ID}
    

    Se o registro terminando com as informações a seguir, significa que o servidor está pronto para atender às solicitações. Isso leva cerca de 30 segundos.

    2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
    [warn] getaddrinfo: address family for nodename not supported
    2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
    [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
    
  3. Envie a solicitação ao servidor do modelo.

    python3 bert_request.py
    

    A saída será semelhante a esta:

    For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
    For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."
    
  4. Fazer a limpeza.

    Limpe o contêiner do Docker antes de executar outras demonstrações.

    CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
    docker stop ${CONTAINER_ID}
    

    Limpe os artefatos do modelo:

    sudo rm -rf /tmp/jax/
    

Executar a demonstração do JAX Stable Diffusion E2E Serving

O modelo de difusão estável pré-treinado é da Hugging Face.

  1. Exporte o modelo salvo do TF2 compatível com TPU a partir do modelo de difusão estável do Flax:

    cd demo/jax/stable_diffusion
    
    python3 export_stable_diffusion_model.py
    
  2. Inicie o contêiner do servidor do modelo do TensorFlow para o modelo:

    docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=/tmp/jax/stable_diffusion_tpu,target=/models/stable_diffusion \
     -e MODEL_NAME=stable_diffusion \
     us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
    

    Verifique o registro do contêiner do servidor do modelo e confira se o gRPC e o servidor HTTP estão ativos:

    CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
    docker logs ${CONTAINER_ID}
    

    Se o registro terminando com as informações a seguir, significa que o servidor está pronto para atender às solicitações. O processo leva cerca de dois minutos.

    2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
    [warn] getaddrinfo: address family for nodename not supported
    2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
    [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
    
  3. Envie a solicitação ao servidor do modelo.

    python3 stable_diffusion_request.py
    

    A solicitação é "Painting of a squirrel skating in New York" e a imagem de saída será salva como stable_diffusion_images.jpg no diretório atual.

  4. Fazer a limpeza.

    Limpe o contêiner do Docker antes de executar outras demonstrações.

    CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
    docker stop ${CONTAINER_ID}
    

    Limpar os artefatos do modelo

    sudo rm -rf /tmp/jax/
    

Inferência e disponibilização de modelos do TensorFlow

Nas seções a seguir, explicamos o fluxo de trabalho da inferência de modelo do TensorFlow.

  1. Usar o conversor de inferência para converter o modelo
  2. Usar o TensorFlow Serving para disponibilizar o modelo

Conversor de inferência

O conversor de inferência do Cloud TPU prepara e otimiza um modelo exportado do TensorFlow ou do JAX para inferência de TPU. O conversor é executado em um shell local ou no shell da VM da TPU. O shell da VM da TPU é recomendado porque vem pré-instalado com as ferramentas de linha de comando necessárias para o conversor. Para mais detalhes sobre o conversor de inferência, consulte o Guia do usuário sobre conversores de inferência.

Pré-requisitos

  1. O modelo precisa ser exportado do TensorFlow ou do JAX no formato LoadBalancer.

  2. O modelo precisa ter um alias de função para a função da TPU. Consulte os exemplos de código no Guia do usuário do conversor de inferência para instruções sobre como fazer isso. Os exemplos a seguir usam tpu_func como alias da função da TPU.

  3. Verifique se a CPU da sua máquina é compatível com as instruções Advanced Vector eXtensions (AVX), já que a biblioteca do TensorFlow (a dependência do Cloud TPU Inference Converter) é compilada para usar as instruções do AVX. A maioria das CPUs tem suporte a AVX.

    • Você pode executar lscpu | grep avx para verificar se o conjunto de instruções AVX tem suporte.

Primeiros passos

  • Configurar o ambiente de VM da TPU Configure o ambiente seguindo as etapas a seguir, dependendo do shell que estiver usando:

    Shell da VM da TPU

    • No shell da VM da TPU, execute os seguintes comandos para permitir o uso do Docker não raiz:
    sudo usermod -a -G docker ${USER}
    newgrp docker
    
    • Inicialize os auxiliares de credenciais do Docker:
    gcloud auth configure-docker \
      us-docker.pkg.dev
    

    Shell local

    No shell local, configure o ambiente seguindo estas etapas:

    • Instale o SDK Cloud, que inclui a ferramenta de linha de comando gcloud.

    • Instale o Docker:

    • Permitir o uso do Docker não raiz:

      sudo usermod -a -G docker ${USER}
      newgrp docker
      
    • Faça login no seu ambiente:

      gcloud auth login
      
    • Inicialize os auxiliares de credenciais do Docker:

      gcloud auth configure-docker \
      us-docker.pkg.dev
      
  • Extraia a imagem do Docker do conversor de inferência:

      CONVERTER_IMAGE=us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
      docker pull ${CONVERTER_IMAGE}
      

Imagem do conversor

A imagem serve para fazer conversões únicas de modelo. Defina os caminhos do modelo e ajuste as opções do conversor de acordo com suas necessidades. A seção "Exemplos de uso" no Guia do usuário sobre conversores de inferência oferece vários casos de uso comuns.

docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      function_alias: "tpu_func"
    }
    batch_options {
      num_batch_threads: 2
      max_batch_size: 8
      batch_timeout_micros: 5000
      allowed_batch_sizes: 2
      allowed_batch_sizes: 4
      allowed_batch_sizes: 8
      max_enqueued_batches: 10
    }
'

Veja na seção a seguir como executar esse modelo com um servidor de modelos do TensorFlow.

TensorFlow Serving

As instruções a seguir mostram como disponibilizar seu modelo do TensorFlow em VMs de TPU.

Pré-requisitos:

Configure suas credenciais do Docker se ainda não tiver feito isso:

  1. Faça o download da imagem do Docker do TensorFlow Serving para sua VM de TPU.

    Definir as variáveis de ambiente de amostra

    export YOUR_LOCAL_MODEL_PATH=model-path
    export MODEL_NAME=model-name
    # Note: this image name may change later.
    export IMAGE_NAME=us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
    

    Fazer o download da imagem do Docker

    docker pull ${IMAGE_NAME}
    
  2. Disponibilize seu modelo do TensorFlow usando a imagem do Docker do TensorFlow Serving na sua VM de TPU.

    # PORT 8500 is for gRPC model server and 8501 is for HTTP/REST model server.
    docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=${YOUR_LOCAL_MODEL_PATH},target=/models/${MODEL_NAME} \
     -e MODEL_NAME=${MODEL_NAME} \
     ${IMAGE_NAME}
    
  3. Siga a API Serving Client para consultar seu modelo.

Exemplo de exibição de modelo do TensorFlow E2E:

Pré-requisito: verifique se você já configurou as credenciais do Docker e extraiu a imagem do Docker do TensorFlow Serving e do conversor de inferência. Caso contrário, execute os seguintes comandos:

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
    us-docker.pkg.dev
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
docker pull us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0

Faça o download do código de demonstração:

gsutil -m cp -r \
  "gs://cloud-tpu-inference-public/demo" \
  .

Instale as dependências de demonstração do TensorFlow:

pip install -r ./demo/tf/requirements.txt

Executar a demonstração do TensorFlow ResNet-50 E2E Serving

  1. Exporte um modelo salvo do TF2 compatível com TPU a partir do modelo Keras ResNet-50.

    cd demo/tf/resnet-50
    
    python3 export_resnet_model.py
    
  2. Iniciar o contêiner do servidor do modelo do TensorFlow para o modelo.

    docker run -t --rm --privileged -d \
     -p 8500:8500 -p 8501:8501 \
     --mount type=bind,source=/tmp/tf/resnet_tpu,target=/models/resnet \
     -e MODEL_NAME=resnet \
     us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
    

    Verifique o registro do contêiner do servidor do modelo e confira se o gRPC e o servidor HTTP estão ativos:

    CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
    docker logs ${CONTAINER_ID}
    

    Se o registro terminando com as seguintes informações, o servidor está pronto para atender a solicitações. Isso leva cerca de 30 segundos.

    2023-04-08 00:43:10.481682: I tensorflow_serving/model_servers/server.cc:409] Running gRPC ModelServer at 0.0.0.0:8500 ...
    [warn] getaddrinfo: address family for nodename not supported
    2023-04-08 00:43:10.520578: I tensorflow_serving/model_servers/server.cc:430] Exporting HTTP/REST API at:localhost:8501 ...
    [evhttp_server.cc : 245] NET_LOG: Entering the event loop ...
    
  3. Envie a solicitação ao servidor do modelo.

    A imagem da solicitação é uma banana de https://i.imgur.com/j9xCCzn.jpeg .

    python3 resnet_request.py
    

    A saída será semelhante a esta:

    Predict result: [[('n07753592', 'banana', 0.94921875), ('n03532672', 'hook', 0.022338867), ('n07749582', 'lemon', 0.005126953)]]
    
  4. Fazer a limpeza.

    Limpe o contêiner do Docker antes de executar outras demonstrações.

    CONTAINER_ID=$(docker ps | grep "tf-serving-tpu" | awk '{print $1}')
    docker stop ${CONTAINER_ID}
    

    Limpe os artefatos do modelo:

    sudo rm -rf /tmp/tf/
    

Inferência e exibição do modelo PyTorch

As seções a seguir explicam o fluxo de trabalho da inferência de modelo PyTorch:

  1. Criar um gerenciador de modelo Python para carregar e inferência usando TorchDynamo e PyTorch/XLA
  2. Use o TorchModelArchiver para criar um arquivo de modelo
  3. Use o TorchServe para exibir o modelo

TorchDynamo e PyTorch/XLA

O TorchDynamo (Dynamo) é um compilador JIT no nível do Python projetado para tornar os programas PyTorch não modificados mais rápidos. Fornece uma API limpa para conexão de back-ends do compilador. O maior recurso é modificar dinamicamente o bytecode Python antes da execução. Na versão PyTorch/XLA 2.0, um back-end experimental para o Dynamo é fornecido para inferência e treinamento.

O Dynamo fornece um gráfico Torch FX (FX) quando reconhece um padrão de modelo, e o PyTorch/XLA usa uma abordagem do Lazy Tensor para compilar o gráfico de FX e retornar a função compilada. Para mais informações sobre os detalhes técnicos da implementação do dínamo do PyTorch/XLA, consulte a postagem do desenvolvedor sobre o PyTorch/XLA e a documentação do TorchDynamo. Veja mais detalhes neste blog.

Veja um pequeno exemplo de código de execução da inferência densenet161 com torch.compile.

import torch
import torchvision
import torch_xla.core.xla_model as xm

def eval_model(loader):
  device = xm.xla_device()
  xla_densenet161 = torchvision.models.densenet161().to(device)
  xla_densenet161.eval()
  dynamo_densenet161 = torch.compile(
      xla_densenet161, backend='torchxla_trace_once')
  for data, _ in loader:
    output = dynamo_densenet161(data)

Tocha

A imagem do Docker TorchServe do Cloud TPU permite exibir o modelo de modo antecipado do PyTorch usando o TorchServe em uma VM do Cloud TPU.

É possível usar a imagem do Docker torchserve-tpu fornecida, pronta para disponibilizar seu modelo pytorch arquivado em uma VM do Cloud TPU.

Configure a autenticação do Docker:

sudo usermod -a -G docker ${USER}
newgrp docker
gcloud auth configure-docker \
    us-docker.pkg.dev

Extraia a imagem do Docker do Cloud TPU TorchServe para a VM da TPU:

CLOUD_TPU_TORCHSERVE_IMAGE_URL=us-docker.pkg.dev/cloud-tpu-images/inference/torchserve-tpu:v0.8.2-20230829
docker pull ${CLOUD_TPU_TORCHSERVE_IMAGE_URL}

Coletar artefatos do modelo

Para começar, é necessário fornecer um gerenciador de modelo, que instrui o worker do servidor de modelo do TochServe a carregar seu modelo, processar os dados de entrada e executar a inferência. É possível usar os gerenciadores de inferência padrão do TochServe (origem) ou desenvolver seu próprio gerenciador de modelo personalizado seguindo o base_handler.py. Talvez também seja necessário fornecer o modelo treinado e o arquivo de definição do modelo.

No exemplo do Densenet 161 abaixo, usamos artefatos de modelo e o gerenciador de classificador de imagem padrão fornecido pelo TorchServe:

O diretório de trabalho é mostrado abaixo.

CWD="$(pwd)"

WORKDIR="${CWD}/densenet_161"

mkdir -p ${WORKDIR}/model-store
mkdir -p ${WORKDIR}/logs
  1. Faça o download e copie os artefatos de modelo do exemplo de classificador de imagens do TorchServe:

    git clone https://github.com/pytorch/serve.git
    
    cp ${CWD}/serve/examples/image_classifier/densenet_161/model.py ${WORKDIR}
    cp ${CWD}/serve/examples/image_classifier/index_to_name.json ${WORKDIR}
    
  2. Faça o download dos pesos do modelo:

    wget https://download.pytorch.org/models/densenet161-8d451a50.pth -O densenet161-8d451a50.pth
    
    mv densenet161-8d451a50.pth ${WORKDIR}
    
  3. Crie um arquivo de configuração de modelo do TorchServe para usar o back-end do Dynamo:

    echo 'pt2: "torchxla_trace_once"' >> ${WORKDIR}/model_config.yaml
    

    Você vai encontrar os arquivos e diretórios mostrados abaixo:

    >> ls ${WORKDIR}
    model_config.yaml
    index_to_name.json
    logs
    model.py
    densenet161-8d451a50.pth
    model-store
    

Gerar um arquivo modelo

Para disponibilizar seu modelo PyTorch com o Cloud TPU TorchServe, é preciso empacotar o gerenciador de modelos e todos os artefatos de modelo em um arquivo de modelo (*.mar) usando o Torch Model Archiver.

Gere um arquivo modelo com torch-model-archiver:

MODEL_NAME=Densenet161

docker run \
    --privileged  \
    --shm-size 16G \
    --name torch-model-archiver \
    -it \
    -d \
    --rm \
    --mount type=bind,source=${WORKDIR},target=/home/model-server/ \
    ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
    torch-model-archiver \
        --model-name ${MODEL_NAME} \
        --version 1.0 \
        --model-file model.py \
        --serialized-file densenet161-8d451a50.pth \
        --handler image_classifier \
        --export-path model-store \
        --extra-files index_to_name.json \
        --config-file model_config.yaml

Você verá o arquivo modelo gerado no diretório model-store:

>> ls ${WORKDIR}/model-store
Densenet161.mar

Exibir solicitações de inferência

Agora que você tem o arquivo de modelo, inicie o servidor de modelos do TorchServe e exiba solicitações de inferência.

  1. Inicie o servidor de modelo do TorchServe:

    docker run \
       --privileged  \
       --shm-size 16G \
       --name torchserve-tpu \
       -it \
       -d \
       --rm \
       -p 7070:7070 \
       -p 7071:7071 \
       -p 8080:8080 \
       -p 8081:8081 \
       -p 8082:8082 \
       -p 9001:9001 \
       -p 9012:9012 \
       --mount type=bind,source=${WORKDIR}/model-store,target=/home/model-server/model-store \
       --mount type=bind,source=${WORKDIR}/logs,target=/home/model-server/logs \
       ${CLOUD_TPU_TORCHSERVE_IMAGE_URL} \
       torchserve \
           --start \
           --ncs \
           --models ${MODEL_NAME}.mar \
           --ts-config /home/model-server/config.properties
    
  2. Integridade do servidor do modelo de consulta:

    curl http://localhost:8080/ping
    

    Se o servidor de modelo estiver em execução, você vai ver:

    {
     "status": "Healthy"
    }
    

    Para consultar as versões padrão do modelo registrado atual, use:

    curl http://localhost:8081/models
    

    Você verá o modelo registrado:

    {
     "models": [
       {
         "modelName": "Densenet161",
         "modelUrl": "Densenet161.mar"
       }
     ]
    }
    

    Para fazer o download de uma imagem para inferência, use:

    curl -O https://raw.githubusercontent.com/pytorch/serve/master/docs/images/kitten_small.jpg
    
    mv kitten_small.jpg ${WORKDIR}
    

    Para enviar uma solicitação de inferência ao servidor do modelo, use:

    curl http://localhost:8080/predictions/${MODEL_NAME} -T ${WORKDIR}/kitten_small.jpg
    

    Uma resposta semelhante a esta será exibida:

    {
     "tabby": 0.47878125309944153,
     "lynx": 0.20393909513950348,
     "tiger_cat": 0.16572578251361847,
     "tiger": 0.061157409101724625,
     "Egyptian_cat": 0.04997897148132324
    }
    
  3. Registros do servidor do modelo

    Use os seguintes comandos para acessar os registros:

    ls ${WORKDIR}/logs/
    cat ${WORKDIR}/logs/model_log.log
    

    Você verá a seguinte mensagem no registro:

    "Modelo compilado com back-end torchxla_trace_once"

Limpeza

Pare o contêiner do Docker:

rm -rf serve
rm -rf ${WORKDIR}

docker stop torch-model-archiver
docker stop torchserve-tpu

Exibição de modelos de linguagem grande

O SAX é um framework de exibição compatível com a exibição de modelos grandes que podem exigir TPUs em vários hosts para serem executadas com o GSPMD, como modelos de linguagem grandes baseados em PAX. O PAX é um framework baseado no JAX usado para treinar modelos em grande escala que permite experimentação e carregamento em paralelo avançados e totalmente configuráveis.

A seção de cluster SAX descreve os principais elementos para entender como o SAX funciona. A seção de exibição do modelo SAX explica um exemplo de exibição de modelo de host único com um modelo GPTJ6B. O SAX também fornece exibição de vários hosts em Cloud TPUs e os usuários podem executar modelos em topologias de TPU maiores para visualização experimental de disponibilização em vários hosts. O exemplo abaixo com um modelo de teste de 175B mostra como testar essa configuração.

Cluster SAX (célula SAX)

O servidor administrador SAX e o servidor do modelo SAX são dois componentes essenciais que executam um cluster SAX.

Servidor de administração do SAX

O servidor de administração SAX monitora e coordena todos os servidores de modelo SAX em um cluster SAX. Em um cluster SAX, é possível iniciar vários servidores de administrador SAX, em que apenas um dos servidores administrador SAX está ativo por meio da eleição de líder, os outros são servidores em espera. Quando o servidor de administração ativo falhar, um servidor de administração em espera ficará ativo. O servidor de administração SAX ativo atribui réplicas de modelo e solicitações de inferência aos servidores de modelos SAX disponíveis.

Bucket de armazenamento do administrador do SAX

Cada cluster SAX requer um bucket do Cloud Storage para armazenar as configurações e os locais dos servidores de administração SAX e de servidores do modelo SAX no cluster SAX.

Servidor do modelo SAX

O servidor do modelo SAX carrega um checkpoint do modelo e executa a inferência com o GSPMD. Um servidor de modelo SAX é executado em um único worker de VM da TPU. A exibição de modelos de TPU de host único requer um único servidor de modelos SAX em uma VM de TPU de host único. A exibição de modelos de TPU com vários hosts requer um grupo de servidores de modelos SAX em uma fração de TPU de vários hosts.

Disponibilização de modelo SAX

Na seção a seguir, apresentamos o fluxo de trabalho para exibir modelos de idioma usando SAX. Ele usa o modelo GPT-J 6B como exemplo para disponibilização de modelos de host único e um modelo de teste 175B para a exibição de modelos de vários hosts.

Antes de começar, instale as imagens do Docker SAX do Cloud TPU na VM da TPU:

sudo usermod -a -G docker ${USER}
newgrp docker

gcloud auth configure-docker us-docker.pkg.dev

SAX_ADMIN_SERVER_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server"
SAX_MODEL_SERVER_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server"
SAX_UTIL_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/inference/sax-util"

SAX_VERSION=v1.0.0

export SAX_ADMIN_SERVER_IMAGE_URL=${SAX_ADMIN_SERVER_IMAGE_NAME}:${SAX_VERSION}
export SAX_MODEL_SERVER_IMAGE_URL=${SAX_MODEL_SERVER_IMAGE_NAME}:${SAX_VERSION}
export SAX_UTIL_IMAGE_URL="${SAX_UTIL_IMAGE_NAME}:${SAX_VERSION}"

docker pull ${SAX_ADMIN_SERVER_IMAGE_URL}
docker pull ${SAX_MODEL_SERVER_IMAGE_URL}
docker pull ${SAX_UTIL_IMAGE_URL}

Defina algumas outras variáveis que você usará mais tarde:

export SAX_ADMIN_SERVER_DOCKER_NAME="sax-admin-server"
export SAX_MODEL_SERVER_DOCKER_NAME="sax-model-server"
export SAX_CELL="/sax/test"

Exemplo de veiculação do modelo de host único da GPT-J 6B

A exibição do modelo de host único é aplicável à fração de TPU de host único, ou seja, v5litepod-1, v5litepod-4 e v5litepod-8.

  1. Criar um cluster SAX

    1. Crie um bucket de armazenamento do Cloud Storage para o cluster SAX:

      SAX_ADMIN_STORAGE_BUCKET=${your_admin_storage_bucket}
      
      gcloud storage buckets create gs://${SAX_ADMIN_STORAGE_BUCKET} \
      --project=${PROJECT_ID}
      

      Talvez seja necessário outro bucket de armazenamento do Cloud Storage para armazenar o checkpoint.

      SAX_DATA_STORAGE_BUCKET=${your_data_storage_bucket}
      
    2. Conecte-se à sua VM da TPU usando SSH em um terminal para iniciar o servidor de administração do SAX:

      docker run \
      --name ${SAX_ADMIN_SERVER_DOCKER_NAME} \
      -it \
      -d \
      --rm \
      --network host \
      --env GSBUCKET=${SAX_ADMIN_STORAGE_BUCKET} \
      ${SAX_ADMIN_SERVER_IMAGE_URL}
      

      É possível verificar o registro do Docker:

      docker logs -f ${SAX_ADMIN_SERVER_DOCKER_NAME}
      

      A saída no registro será semelhante a esta:

      I0829 01:22:31.184198       7 config.go:111] Creating config fs_root: "gs://test_sax_admin/sax-fs-root"
      I0829 01:22:31.347883       7 config.go:115] Created config fs_root: "gs://test_sax_admin/sax-fs-root"
      I0829 01:22:31.360837      24 admin_server.go:44] Starting the server
      I0829 01:22:31.361420      24 ipaddr.go:39] Skipping non-global IP address 127.0.0.1/8.
      I0829 01:22:31.361455      24 ipaddr.go:39] Skipping non-global IP address ::1/128.
      I0829 01:22:31.361462      24 ipaddr.go:39] Skipping non-global IP address fe80::4001:aff:fe8e:fc8/64.
      I0829 01:22:31.361469      24 ipaddr.go:39] Skipping non-global IP address fe80::42:bfff:fef9:1bd3/64.
      I0829 01:22:31.361474      24 ipaddr.go:39] Skipping non-global IP address fe80::20fb:c3ff:fe5b:baac/64.
      I0829 01:22:31.361482      24 ipaddr.go:56] IPNet address 10.142.15.200
      I0829 01:22:31.361488      24 ipaddr.go:56] IPNet address 172.17.0.1
      I0829 01:22:31.456952      24 admin.go:305] Loaded config: fs_root: "gs://test_sax_admin/sax-fs-root"
      I0829 01:22:31.609323      24 addr.go:105] SetAddr /gcs/test_sax_admin/sax-root/sax/test/location.proto "10.142.15.200:10000"
      I0829 01:22:31.656021      24 admin.go:325] Updated config: fs_root: "gs://test_sax_admin/sax-fs-root"
      I0829 01:22:31.773245      24 mgr.go:781] Loaded manager state
      I0829 01:22:31.773260      24 mgr.go:784] Refreshing manager state every 10s
      I0829 01:22:31.773285      24 admin.go:350] Starting the server on port 10000
      I0829 01:22:31.773292      24 cloud.go:506] Starting the HTTP server on port 8080
      
  2. Inicie um servidor do modelo SAX de host único no cluster SAX:

    Neste ponto, o cluster SAX contém apenas o servidor de administração SAX. É possível se conectar à sua VM de TPU por SSH em um segundo terminal para iniciar um servidor de modelo SAX no cluster SAX:

    docker run \
        --privileged  \
        -it \
        -d \
        --rm \
        --network host \
        --name ${SAX_MODEL_SERVER_DOCKER_NAME} \
        --env SAX_ROOT=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
        ${SAX_MODEL_SERVER_IMAGE_URL} \
           --sax_cell=${SAX_CELL} \
           --port=10001 \
           --platform_chip=tpuv4 \
           --platform_topology=1x1
    

    É possível verificar o registro do Docker:

    docker logs -f ${SAX_MODEL_SERVER_DOCKER_NAME}
    

  3. Converter checkpoint do modelo:

    É necessário instalar o PyTorch e o Transformers para fazer o download do checkpoint do GPT-J da EleutherAI:

    pip3 install accelerate
    pip3 install torch
    pip3 install transformers
    

    Para converter o checkpoint em SAX, é preciso instalar o paxml:

    pip3 install paxml==1.1.0
    

    Em seguida, defina a seguinte variável:

    >>PT_CHECKPOINT_PATH=./fine_tuned_pt_checkpoint
    

    Faça o download do checkpoint ajustado do PyTorch para ${PT_CHECKPOINT_PATH}, siga https://github.com/mlcommons/inference/blob/master/language/gpt-j/README.md#download-gpt-j-model e execute os seguintes comandos:

    ls ${PT_CHECKPOINT_PATH}
    

    Isso deve listar o seguinte:

    added_tokens.json  generation_config.json
    pytorch_model.bin.index.json
    pytorch_model-00001-of-00003.bin
    pytorch_model-00002-of-00003.bin
    pytorch_model-00003-of-00003.bin
    special_tokens_map.json
    trainer_state.json
    config.json
    merges.txt
    tokenizer_config.json
    vocab.json
    

    O script a seguir converte o checkpoint da GPT-J em SAX:

    python3 -m convert_gptj_ckpt --base EleutherAI/gpt-j-6b --pax pax_6b
    
    wget https://raw.githubusercontent.com/google/saxml/main/saxml/tools/convert_gptj_ckpt.py
    python3 -m convert_gptj_ckpt --base ${PT_CHECKPOINT_PATH} --pax .
    
    Isso gera um conteúdo semelhante ao seguinte:
    transformer.wte.weight (50401, 4096)
    transformer.h.0.ln_1.weight (4096,)
    transformer.h.0.ln_1.bias (4096,)
    transformer.h.0.attn.k_proj.weight (4096, 4096)
    .
    .
    .
    transformer.ln_f.weight (4096,)
    transformer.ln_f.bias (4096,)
    lm_head.weight (50401, 4096)
    lm_head.bias (50401,)
    Saving the pax model to .
    done
    

    Depois que a conversão for concluída, digite o seguinte comando:

    ls checkpoint_00000000/
    

    Isso deve listar o seguinte:

    metadate
    state
    

    Você precisa criar um arquivo commit_success e colocar nos subdiretórios:

    CHECKPOINT_PATH=gs://${SAX_DATA_STORAGE_BUCKET}/path/to/checkpoint_00000000
    gsutil -m cp -r checkpoint_00000000 ${CHECKPOINT_PATH}
    
    touch commit_success.txt
    gsutil cp commit_success.txt ${CHECKPOINT_PATH}/
    gsutil cp commit_success.txt ${CHECKPOINT_PATH}/metadata/
    gsutil cp commit_success.txt ${CHECKPOINT_PATH}/state/
    
  4. Publicar o modelo no cluster SAX

    Agora é possível publicar a GPT-J com o checkpoint convertido na etapa anterior.

    MODEL_NAME=gptjtokenizedbf16bs32
    MODEL_CONFIG_PATH=saxml.server.pax.lm.params.gptj.GPTJ4TokenizedBF16BS32
    REPLICA=1
    

    Para publicar o GPT-J (e as etapas posteriores), use o SSH para se conectar à VM da TPU em um terceiro terminal:

    docker run \
     ${SAX_UTIL_IMAGE_URL} \
       --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
       publish \
         ${SAX_CELL}/${MODEL_NAME} \
         ${MODEL_CONFIG_PATH} \
         ${CHECKPOINT_PATH} \
         ${REPLICA}
    

    Você verá muita atividade do registro do Docker do servidor de modelo até que algo parecido com isto seja exibido para indicar que o modelo foi carregado com êxito:

    I0829 01:33:49.287459 139865140229696 servable_model.py:697] loading completed.
    
  5. Gerar resultados de inferência

    Para a GPT-J, a entrada e a saída precisam ser formatadas como uma string de ID de token separada por vírgulas. Você vai precisar tokenizar a entrada de texto.

    TEXT = ("Below is an instruction that describes a task, paired with "
    "an input that provides further context. Write a response that "
    "appropriately completes the request.\n\n### Instruction\:\nSummarize the "
    "following news article\:\n\n### Input\:\nMarch 10, 2015 . We're truly "
    "international in scope on Tuesday. We're visiting Italy, Russia, the "
    "United Arab Emirates, and the Himalayan Mountains. Find out who's "
    "attempting to circumnavigate the globe in a plane powered partially by the "
    "sun, and explore the mysterious appearance of craters in northern Asia. "
    "You'll also get a view of Mount Everest that was previously reserved for "
    "climbers. On this page you will find today's show Transcript and a place "
    "for you to request to be on the CNN Student News Roll Call. TRANSCRIPT . "
    "Click here to access the transcript of today's CNN Student News program. "
    "Please note that there may be a delay between the time when the video is "
    "available and when the transcript is published. CNN Student News is "
    "created by a team of journalists who consider the Common Core State "
    "Standards, national standards in different subject areas, and state "
    "standards when producing the show. ROLL CALL . For a chance to be "
    "mentioned on the next CNN Student News, comment on the bottom of this page "
    "with your school name, mascot, city and state. We will be selecting "
    "schools from the comments of the previous show. You must be a teacher or a "
    "student age 13 or older to request a mention on the CNN Student News Roll "
    "Call! Thank you for using CNN Student News!\n\n### Response\:")
    

    É possível conseguir a string dos IDs de token usando o tokenizer EleutherAI/gpt-j-6b:

    from transformers import GPT2Tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained(EleutherAI/gpt-j-6b)
    

    Tokenizar o texto de entrada:

    encoded_example = tokenizer(TEXT)
    input_ids = encoded_example.input_ids
    INPUT_STR = ",".join([str(input_id) for input_id in input_ids])
    

    Você vai ver uma string de ID do token semelhante a esta:

    >>> INPUT_STR
    '21106,318,281,12064,326,8477,257,4876,11,20312,351,281,5128,326,3769,2252,4732,13,19430,257,2882,326,20431,32543,262,2581,13,198,198,21017,46486,25,198,13065,3876,1096,262,1708,1705,2708,25,198,198,21017,23412,25,198,16192,838,11,1853,764,775,821,4988,3230,287,8354,319,3431,13,775,821,10013,8031,11,3284,11,262,1578,4498,24880,11,290,262,42438,22931,21124,13,9938,503,508,338,9361,284,2498,4182,615,10055,262,13342,287,257,6614,13232,12387,416,262,4252,11,290,7301,262,11428,5585,286,1067,8605,287,7840,7229,13,921,1183,635,651,257,1570,286,5628,41336,326,373,4271,10395,329,39311,13,1550,428,2443,345,481,1064,1909,338,905,42978,290,257,1295,329,345,284,2581,284,307,319,262,8100,13613,3000,8299,4889,13,48213,6173,46023,764,6914,994,284,1895,262,14687,286,1909,338,8100,13613,3000,1430,13,4222,3465,326,612,743,307,257,5711,1022,262,640,618,262,2008,318,1695,290,618,262,14687,318,3199,13,8100,13613,3000,318,2727,416,257,1074,286,9046,508,2074,262,8070,7231,1812,20130,11,2260,5423,287,1180,2426,3006,11,290,1181,5423,618,9194,262,905,13,15107,3069,42815,764,1114,257,2863,284,307,4750,319,262,1306,8100,13613,3000,11,2912,319,262,4220,286,428,2443,351,534,1524,1438,11,37358,11,1748,290,1181,13,775,481,307,17246,4266,422,262,3651,286,262,2180,905,13,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,0,6952,345,329,1262,8100,13613,3000,0,198,198,21017,18261,25'
    

    Para gerar um resumo da sua matéria, faça o seguinte:

    docker run \
      ${SAX_UTIL_IMAGE_URL} \
        --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
        lm.generate \
          ${SAX_CELL}/${MODEL_NAME} \
          ${INPUT_STR}
    

    Você vai notar algo semelhante a:

    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
    |                                                                                                                                                    GENERATE                                                                                                                                                    |    SCORE     |
    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
    | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256 | -0.023136413 |
    | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,0,50256  |  -0.91842502 |
    | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,921,1276,307,257,4701,393,257,3710,2479,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256     |   -1.1726116 |
    | 1212,2443,3407,262,905,42978,764,198,11041,262,42978,284,1037,2444,351,3555,35915,290,25818,764,198,2953,262,4220,286,262,2443,11,2912,329,257,2863,284,307,4750,319,8100,13613,3000,13,220,921,1276,307,1511,393,4697,284,2581,257,3068,319,262,8100,13613,3000,8299,4889,13,50256                            |   -1.2472695 |
    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
    

    Para remover a tokenização da string de IDs do token de saída:

    output_token_ids = [int(token_id) for token_id in OUTPUT_STR.split(',')]
    OUTPUT_TEXT = tokenizer.decode(output_token_ids, skip_special_tokens=True)
    

    O texto sem token poderá ser mostrado como:

    >>> OUTPUT_TEXT
    'This page includes the show Transcript.\nUse the Transcript to help
    students with reading comprehension and vocabulary.\nAt the bottom of
    the page, comment for a chance to be mentioned on CNN Student News.
    You must be a teacher or a student age 13 or older to request a mention on the CNN Student News Roll Call.'
    
  6. Limpe os contêineres do Docker e os buckets de armazenamento do Cloud Storage.

Prévia de disponibilização do modelo de vários hosts para 175 bilhões

Alguns dos modelos de linguagem grandes exigem uma fração de TPU de vários hosts, ou seja, a v5litepod-16 e versões mais recentes. Nesses casos, todos os hosts de TPU de vários hosts precisarão ter uma cópia de um servidor de modelo SAX, e todos os servidores de modelo funcionam como um grupo de servidores de modelos SAX para exibir o modelo grande em uma fração de TPU de vários hosts.

  1. Criar um novo cluster SAX

    Você pode seguir a mesma etapa de "Criar um cluster SAX" no tutorial sobre GPT-J para criar um novo cluster SAX e um servidor de administração SAX.

    Ou, se você já tiver um cluster SAX, poderá iniciar um servidor de modelo de vários hosts nele.

  2. Iniciar um servidor do modelo SAX de vários hosts em um cluster SAX

    Para criar uma fração de TPU de vários hosts, use o mesmo comando usado para uma fração de TPU de host único. Basta especificar o tipo de acelerador de vários hosts apropriado:

    ACCELERATOR_TYPE=v5litepod-32
    ZONE=us-east1-c
    
    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id ${TPU_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --accelerator-type ${ACCELERATOR_TYPE} \
      --runtime-version ${RUNTIME_VERSION} \
      --service-account ${SERVICE_ACCOUNT} \
      --reserved
    

    Para extrair a imagem do servidor do modelo SAX para todos os hosts/trabalhadores da TPU e iniciá-los:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --worker=all \
      --command="
        gcloud auth configure-docker \
          us-docker.pkg.dev
        # Pull SAX model server image
        docker pull ${SAX_MODEL_SERVER_IMAGE_URL}
        # Run model server
        docker run \
          --privileged  \
          -it \
          -d \
          --rm \
          --network host \
          --name ${SAX_MODEL_SERVER_DOCKER_NAME} \
          --env SAX_ROOT=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
          ${SAX_MODEL_SERVER_IMAGE_URL} \
            --sax_cell=${SAX_CELL} \
            --port=10001 \
            --platform_chip=tpuv4 \
            --platform_topology=1x1"
    
  3. Publicar o modelo no cluster SAX

    Neste exemplo, usamos um modelo LmCloudSpmd175B32Test:

    MODEL_NAME=lmcloudspmd175b32test
    MODEL_CONFIG_PATH=saxml.server.pax.lm.params.lm_cloud.LmCloudSpmd175B32Test
    CHECKPOINT_PATH=None
    REPLICA=1
    

    Para publicar o modelo de teste, faça o seguinte:

    docker run \
     ${SAX_UTIL_IMAGE_URL} \
       --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
       publish \
         ${SAX_CELL}/${MODEL_NAME} \
         ${MODEL_CONFIG_PATH} \
         ${CHECKPOINT_PATH} \
         ${REPLICA}
    
  4. Gerar resultados de inferência

    docker run \
      ${SAX_UTIL_IMAGE_URL} \
        --sax_root=gs://${SAX_ADMIN_STORAGE_BUCKET}/sax-root \
        lm.generate \
          ${SAX_CELL}/${MODEL_NAME} \
          "Q:  Who is Harry Porter's mother? A\: "
    

    Observe que, como este exemplo usa um modelo de teste com pesos aleatórios, a saída pode não ser significativa.

  5. Limpeza

    Pare os contêineres do Docker:

    docker stop ${SAX_ADMIN_SERVER_DOCKER_NAME}
    docker stop ${SAX_MODEL_SERVER_DOCKER_NAME}
    

    Exclua seu bucket de armazenamento de administrador do Cloud Storage e todos os buckets de armazenamento de dados usando gsutil, conforme mostrado abaixo.

    gsutil rm -rf gs://${SAX_ADMIN_STORAGE_BUCKET}
    gsutil rm -rf gs://${SAX_DATA_STORAGE_BUCKET}
    

Caracterização de perfil

Depois de configurar a inferência, os criadores de perfil podem ser usados para analisar o desempenho e o uso da TPU. Referências a alguns documentos relacionados à criação de perfis são mostrados abaixo:

Suporte e feedback

Seu feedback é muito bem-vindo! Para compartilhar feedback ou solicitar suporte, entre em contato neste link ou envie um e-mail para cloudtpu-preview-support@google.com

Termos

Todas as informações que o Google forneceu a você sobre este Acesso antecipado são informações confidenciais do Google e estão sujeitas às disposições de confidencialidade nos Termos de Serviço do Google Cloud Platform (ou outro contrato que regule seu uso do Google Cloud Platform).