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.
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.
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.
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
Faça login na sua Conta do Google. Crie uma nova conta se ainda não tiver feito isso.
No Console do Cloud, selecione ou crie um projeto do Cloud na página do seletor de projetos.
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.
Instale o gcloud alpha components.
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
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
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}
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.
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
- Nome do projeto. Use o nome do projeto do Google.
- Consulte a seção Tipos de acelerador para ver os tipos compatíveis.
- Toda a capacidade do Acesso antecipado vai estar em us-west4-a.
- v2-alpha-tpuv5-lite
- 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.
- O ID atribuído pelo usuário da TPU criado quando a solicitação de recurso na fila é alocada.
- O ID atribuído pelo usuário da solicitação de recurso na fila.
PROJECT_ID
ACCELERATOR_TYPE
ZONE
RUNTIME_VERSION
SERVICE_ACCOUNT
TPU_NAME
QUEUED_RESOURCE_ID
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}
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.
- Usar
jax2tf
para converter o modelo para o TensorFlow 2 e salvá-lo - Usar o conversor de inferência para converter o modelo salvo
- 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.
Exporte um modelo salvo do TF2 compatível com TPU de um modelo Flax BERT:
cd demo/jax/bert
python3 export_bert_model.py
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 ...
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?."
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.
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
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 ...
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.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.
- Usar o conversor de inferência para converter o modelo
- 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
O modelo precisa ser exportado do TensorFlow ou do JAX no formato LoadBalancer.
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.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.
- Você pode executar
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:
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}
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}
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
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
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 ...
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)]]
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:
- Criar um gerenciador de modelo Python para carregar e inferência usando TorchDynamo e PyTorch/XLA
- Use o TorchModelArchiver para criar um arquivo de modelo
- 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:
- model.py (link em inglês)
- index_to_name.json (link em inglês)
- image_classifier.py.
O diretório de trabalho é mostrado abaixo.
CWD="$(pwd)" WORKDIR="${CWD}/densenet_161" mkdir -p ${WORKDIR}/model-store mkdir -p ${WORKDIR}/logs
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}
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}
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.
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
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 }
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.
Criar um cluster SAX
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}
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
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}
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}
, sigahttps://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/
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.
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.'
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.
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.
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"
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}
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.
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:
- Como criar um perfil no Cloud TPU
- Criação de perfis do TensorFlow
- Criação de perfil do PyTorch
- Criação de perfil do JAX
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).