Implantar e disponibilizar modelos Llama 2 na TPU v5e usando SAX


Este tutorial vai ajudá-lo a:

  • Criar uma VM do Cloud TPU para implantar a família Llama 2 de modelos de linguagem grandes (LLMs), disponível em tamanhos diferentes (7B, 13B ou 70B).
  • Preparando checkpoints para os modelos e implantando-os no SAX
  • Interagir com o modelo usando um endpoint HTTP

A disponibilização para experimentos do AGI (SAX) é um sistema experimental que exibe modelos Paxml, JAX e PyTorch para inferência. O código e a documentação do SAX estão no repositório Git do SAX. A versão estável atual compatível com a TPU v5e é a v1.1.0.

Sobre células SAX

Uma célula (ou cluster) SAX é a unidade principal para exibir seus modelos. Ele consiste em duas partes principais:

  • Servidor administrador: monitora os servidores de modelo, atribui modelos a eles e ajuda os clientes a encontrar o servidor de modelo certo para interagir.
  • Servidores de modelo: esses servidores executam seu modelo. Eles são responsáveis por processar as solicitações recebidas e gerar respostas.

Veja no diagrama a seguir uma célula SAX:

Célula SAX com servidor de administrador e servidores de modelo

Figura 1. Célula SAX com servidor administrador e servidor modelo.

É possível interagir com uma célula SAX usando clientes escritos em Python, C++ ou Go ou diretamente por um servidor HTTP. O diagrama a seguir mostra como um cliente externo pode interagir com uma célula SAX:

Um cliente externo e um endpoint HTTP interagindo com uma célula SAX

Figura 2. Arquitetura de ambiente de execução de um cliente externo interagindo com uma célula SAX.

Objetivos

  • Configurar recursos da TPU para veiculação
  • Criar um cluster SAX
  • Publicar o modelo Llama 2
  • Interagir com o modelo

Custos

Neste documento, você usará os seguintes componentes faturáveis do Google Cloud:

  • Cloud TPU
  • Compute Engine
  • Cloud Storage

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços. Novos usuários do Google Cloud podem estar qualificados para uma avaliação gratuita.

Antes de começar

Configure seu projeto do Google Cloud, ative a API Cloud TPU e crie uma conta de serviço seguindo as instruções em Configurar o ambiente do Cloud TPU.

Criar uma TPU

As etapas a seguir mostram como criar uma VM da TPU que atenderá seu modelo.

  1. Crie variáveis de ambiente:

    export PROJECT_ID=PROJECT_ID
    export ACCELERATOR_TYPE=ACCELERATOR_TYPE
    export ZONE=ZONE
    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 de ambiente

    PROJECT_ID
    O ID do seu projeto do Google Cloud.
    ACCELERATOR_TYPE
    O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Tamanhos diferentes do modelo Llama 2 têm requisitos distintos de tamanho de TPU:
    • 7B: v5litepod-4 ou maior
    • 13B: v5litepod-8 ou maior
    • 70B: v5litepod-16 ou maior
    ZONE
    A zona em que você quer criar o Cloud TPU.
    SERVICE_ACCOUNT
    A conta de serviço que você quer anexar ao Cloud TPU.
    TPU_NAME
    O nome do Cloud TPU.
    QUEUED_RESOURCE_ID
    Um identificador para sua solicitação de recurso na fila.
  2. Defina o ID do projeto e a zona na configuração ativa da Google Cloud CLI:

    gcloud config set project $PROJECT_ID  && gcloud config set compute/zone $ZONE
    
  3. Crie a VM da 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}
    
  4. Verifique se a TPU está ativa:

    gcloud compute tpus queued-resources list --project $PROJECT_ID --zone $ZONE
    

Configurar nó de conversão do checkpoint

Para executar os modelos LLama em um cluster SAX, você precisa converter os checkpoints originais do Llama em um formato compatível com SAX.

A conversão requer recursos de memória significativos, dependendo do tamanho do modelo:

Modelo Tipo de máquina
7 bi 50-60 GB de memória
13 bi 120 GB de memória
70 bi 500-600 GB de memória (tipo de máquina N2 ou M1)

Para os modelos 7B e 13B, é possível executar a conversão na VM da TPU. Para o modelo 70 bilhões, você precisa criar uma instância do Compute Engine com aproximadamente 1 TB de espaço em disco:

gcloud compute instances create INSTANCE_NAME --project=$PROJECT_ID --zone=$ZONE \
  --machine-type=n2-highmem-128 \
  --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
  --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
  --service-account=$SERVICE_ACCOUNT \
  --scopes=https://www.googleapis.com/auth/cloud-platform \
  --tags=http-server,https-server \
  --create-disk=auto-delete=yes,boot=yes,device-name=bk-workday-dlvm,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
  --no-shielded-secure-boot \
  --shielded-vtpm \
  --shielded-integrity-monitoring \
  --labels=goog-ec-src=vm_add-gcloud \
  --reservation-affinity=any

Se você usa uma TPU ou uma instância do Compute Engine como servidor de conversão, configure o servidor para converter os checkpoints do Llama 2:

  1. Para os modelos 7B e 13B, defina a variável de ambiente do nome do servidor como o nome da sua TPU:

    export CONV_SERVER_NAME=$TPU_NAME
    

    Para o modelo 70B, defina a variável de ambiente do nome do servidor como o nome da sua instância do Compute Engine:

    export CONV_SERVER_NAME=INSTANCE_NAME
    
  2. Conecte-se ao nó de conversão usando SSH.

    Se o nó de conversão for uma TPU, conecte-se a ela:

    gcloud compute tpus tpu-vm ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    

    Se o nó de conversão for uma instância do Compute Engine, conecte-se à VM do Compute Engine:

    gcloud compute ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    
  3. Instale os pacotes necessários no nó de conversão:

    sudo apt update
    sudo apt-get install python3-pip
    sudo apt-get install git-all
    
    pip3 install paxml==1.1.0
    pip3 install torch
    pip3 install jaxlib==0.4.14
    
  4. Faça o download do script de conversão de checkpoint do Llama:

    gcloud storage cp gs://cloud-tpu-inference-public/sax-tokenizers/llama/convert_llama_ckpt.py .
    

Fazer o download dos pesos do Llama 2

Antes de converter o modelo, faça o download dos pesos do Llama 2. Para este tutorial, use os pesos originais do Llama 2 (por exemplo, meta-llama/Llama-2-7b) e não os que foram convertidos para o formato Hugging Face Transformers (por exemplo, meta-llama/Llama-2-7b-hf).

Se você já tiver os pesos do Llama 2, avance para Converter as ponderações.

Para fazer o download dos pesos do hub Hugging Face, é necessário configurar um token de acesso do usuário e solicitar acesso aos modelos do Llama 2. Para solicitar acesso, siga as instruções na página do Hugging Face para o modelo que você quer usar, por exemplo, meta-llama/Llama-2-7b.

  1. Crie um diretório para os pesos:

    sudo mkdir WEIGHTS_DIRECTORY
    
  2. Pegue os pesos Llama2 no hub Hugging Face:

    1. Instale a CLI do hub Hugging Face:

      pip install -U "huggingface_hub[cli]"
      
    2. Mude para o diretório "weights":

      cd WEIGHTS_DIRECTORY
      
    3. Faça o download dos arquivos do Llama 2:

      python3
      from huggingface_hub import login
      login()
      from huggingface_hub import hf_hub_download, snapshot_download
      import os
      PATH=os.getcwd()
      snapshot_download(repo_id="meta-llama/LLAMA2_REPO", local_dir_use_symlinks=False, local_dir=PATH)
      

      Substitua LLAMA2_REPO pelo nome do repositório do Hugging Face de que você quer fazer o download: Llama-2-7b, Llama-2-13b ou Llama-2-70b.

Converter os pesos

Edite o script de conversão e execute-o para converter as ponderações do modelo.

  1. Crie um diretório para armazenar os pesos convertidos:

    sudo mkdir CONVERTED_WEIGHTS
    
  2. Clone o repositório Saxml do GitHub em um diretório em que você tenha permissões de leitura, gravação e execução:

      git clone https://github.com/google/saxml.git -b r1.1.0
    
  3. Altere para o diretório saxml:

    cd saxml
    
  4. Abra o arquivo saxml/tools/convert_llama_ckpt.py.

  5. No arquivo saxml/tools/convert_llama_ckpt.py, altere a linha 169 de:

    'scale': pytorch_vars[0]['layers.%d.attention_norm.weight' % (layer_idx)].type(torch.float16).numpy()
    

    Depois:

    'scale': pytorch_vars[0]['norm.weight'].type(torch.float16).numpy()
    
  6. Execute o script saxml/tools/init_cloud_vm.sh:

    saxml/tools/init_cloud_vm.sh
    
  7. Somente para 70 bilhões: desative o modo de teste:

    1. Abra o arquivo saxml/server/pax/lm/params/lm_cloud.py.
    2. No arquivo saxml/server/pax/lm/params/lm_cloud.py, altere a linha 344 de:

      return True
      

      Depois:

      return False
      
  8. Converta os pesos:

    python3 saxml/tools/convert_llama_ckpt.py --base-model-path WEIGHTS_DIRECTORY \
      --pax-model-path CONVERTED_WEIGHTS \
      --model-size MODEL_SIZE
    

    Substitua:

    • WEIGHTS_DIRECTORY: diretório dos pesos originais.
    • CONVERTED_WEIGHTS: caminho de destino das ponderações convertidas.
    • MODEL_SIZE: 7b, 13b ou 70b.

Preparar o diretório de checkpoint

Depois de converter os checkpoints, o diretório deles precisa ter a seguinte estrutura:

checkpoint_00000000
  metadata/
      metadata
    state/
        mdl_vars.params.lm*/
        ...
        ...
        step/

Crie um arquivo vazio chamado commit_success.txt e coloque uma cópia dele nos diretórios checkpoint_00000000, metadata e state. Isso permite que o SAX saiba que esse checkpoint está totalmente convertido e pronto para ser carregado:

  1. Altere para o diretório do checkpoint:

    cd CONVERTED_WEIGHTS/checkpoint_00000000
    
  2. Crie um arquivo vazio chamado commit_success.txt:

    touch commit_success.txt
    
  3. Altere para o diretório de metadados e crie um arquivo vazio chamado commit_success.txt:

    cd metadata && touch commit_success.txt
    
  4. Mude para o diretório de estado e crie um arquivo vazio chamado commit_success.txt:

    cd .. && cd state && touch commit_success.txt
    

O diretório de checkpoint agora terá a seguinte estrutura:

checkpoint_00000000
    commit_success.txt
metadata/
    commit_success.txt
    metadata
    state/
        commit_success.txt
        mdl_vars.params.lm*/
        ...
        ...
        step/

Produzir um bucket do Cloud Storage

Você precisa armazenar os checkpoints convertidos em um bucket do Cloud Storage para que estejam disponíveis ao publicar o modelo.

  1. Defina uma variável de ambiente para o nome do bucket do Cloud Storage:

    export GSBUCKET=BUCKET_NAME
    
  2. Crie um bucket:

    gcloud storage buckets create gs://${GSBUCKET}
    
  3. Copie os arquivos de checkpoint convertidos para o bucket:

    gcloud storage cp -r CONVERTED_WEIGHTS/checkpoint_00000000  gs://$GSBUCKET/sax_models/llama2/SAX_LLAMA2_DIR/
    

    Substitua SAX_LLAMA2_DIR pelo valor apropriado:

    • 7 bi: saxml_llama27b
    • 13 bi: saxml_llama213b
    • 70 bi: saxml_llama270b

Criar cluster SAX

Para criar um cluster SAX, você precisa:

Em uma implantação típica, você executa o servidor de administrador em uma instância do Compute Engine e o servidor de modelo em uma TPU ou GPU. Para os fins deste tutorial, você implantará o servidor administrador e o servidor de modelo na mesma instância da TPU v5e.

Criar servidor de administrador

Crie o contêiner do Docker do servidor de administrador:

  1. No servidor de conversão, instale o Docker:

    sudo apt-get update
    sudo apt-get install docker.io
    
  2. Inicie o contêiner do Docker do servidor de administrador:

    sudo docker run --name sax-admin-server \
      -it \
      -d \
      --rm \
      --network host \
      --env GSBUCKET=${GSBUCKET} us-docker.pkg.dev/cloud-tpu-images/inference/sax-admin-server:v1.1.0
    

É possível executar o comando docker run sem a opção -d para ver os registros e garantir que o servidor de administrador seja iniciado corretamente.

Criar servidor de modelo

As seções a seguir mostram como criar um servidor de modelo.

Modelo 7b

Inicie o contêiner do Docker do servidor de modelos:

sudo docker run --privileged \
  -it \
  -d \
  --rm \
  --network host \
  --name "sax-model-server" \
  --env SAX_ROOT=gs://${GSBUCKET}/sax-root us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
  --sax_cell="/sax/test" \
  --port=10001 \
  --platform_chip=tpuv5e \
  --platform_topology='4'

Modelo 13b

A configuração de LLaMA13BFP16TPUv5e não está presente em lm_cloud.py. As etapas a seguir mostram como atualizar lm_cloud.py e confirmar uma nova imagem do Docker.

  1. Inicie o servidor de modelo:

    sudo docker run --privileged \
      -it \
      -d \
      --rm \
      --network host \
      --name "sax-model-server" \
      --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
      us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
      --sax_cell="/sax/test" \
      --port=10001 \
      --platform_chip=tpuv5e \
      --platform_topology='8'
    
  2. Conecte-se ao contêiner do Docker usando SSH:

    sudo docker exec -it sax-model-server bash
    
  3. Instale o Vim na imagem do Docker:

    $ apt update
    $ apt install vim
    
  4. Abra o arquivo saxml/server/pax/lm/params/lm_cloud.py. Pesquise LLaMA13B. O seguinte código vai aparecer:

    @servable_model_registry.register
    @quantization.for_transformer(quantize_on_the_fly=False)
    class LLaMA13B(BaseLLaMA):
      """13B model on a A100-40GB.
    
      April 12, 2023
      Latency = 5.06s with 128 decoded tokens. 38ms per output token.
      """
    
      NUM_LAYERS = 40
      VOCAB_SIZE = 32000
      DIMS_PER_HEAD = 128
      NUM_HEADS = 40
      MODEL_DIMS = 5120
      HIDDEN_DIMS = 13824
      ICI_MESH_SHAPE = [1, 1, 1]
    
      @property
      def test_mode(self) -> bool:
        return True
    
  5. Comente ou exclua a linha que começa com @quantization. Após essa alteração, o arquivo ficará assim:

    @servable_model_registry.register
    class LLaMA13B(BaseLLaMA):
      """13B model on a A100-40GB.
    
      April 12, 2023
      Latency = 5.06s with 128 decoded tokens. 38ms per output token.
      """
    
      NUM_LAYERS = 40
      VOCAB_SIZE = 32000
      DIMS_PER_HEAD = 128
      NUM_HEADS = 40
      MODEL_DIMS = 5120
      HIDDEN_DIMS = 13824
      ICI_MESH_SHAPE = [1, 1, 1]
    
      @property
      def test_mode(self) -> bool:
        return True
    
  6. Adicione o código abaixo para oferecer suporte à configuração da TPU.

    @servable_model_registry.register
    class LLaMA13BFP16TPUv5e(LLaMA13B):
    """13B model on TPU v5e-8.
    
    """
    
    BATCH_SIZE = [1]
    BUCKET_KEYS = [128]
    MAX_DECODE_STEPS = [32]
    ENABLE_GENERATE_STREAM = False
    
    ICI_MESH_SHAPE = [1, 1, 8]
    
    @property
    def test_mode(self) -> bool:
      return False
    
  7. Saia da sessão SSH do contêiner do Docker:

    exit
    
  8. Confirme as alterações em uma nova imagem do Docker:

    sudo docker commit sax-model-server sax-model-server:v1.1.0-mod
    
  9. Verifique se a nova imagem Docker foi criada:

    sudo docker images
    

    É possível publicar a imagem Docker no Artifact Registry do seu projeto, mas este tutorial continuará com a imagem local.

  10. Pare o servidor de modelo. O restante do tutorial usará o servidor de modelo atualizado.

    sudo docker stop sax-model-server
    
  11. Inicie o servidor do modelo usando a imagem do Docker atualizada. Especifique o nome atualizado da imagem, sax-model-server:v1.1.0-mod:

    sudo docker run --privileged \
      -it \
      -d \
      --rm \
      --network host \
      --name "sax-model-server" \
      --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
      sax-model-server:v1.1.0-mod \
      --sax_cell="/sax/test" \
      --port=10001 \
      --platform_chip=tpuv5e \
      --platform_topology='8'
    

Modelo 70B

Conecte-se à TPU usando SSH e inicie o servidor de modelo:

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
   sudo docker pull us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0
   # Run model server
   sudo docker run \
     --privileged  \
     -it \
     -d \
     --rm \
     --network host \
     --name "sax-model-server"  \
     --env SAX_ROOT=gs://${GSBUCKET}/sax-root \
     us-docker.pkg.dev/cloud-tpu-images/inference/sax-model-server:v1.1.0 \
     --sax_cell="/sax/test" \
     --port=10001 \
     --platform_chip=tpuv5e \
     --platform_topology='16'
"

Verifique os registros

Verifique os registros do servidor de modelo para confirmar se ele foi iniciado corretamente:

docker logs -f sax-model-server

Se o servidor de modelo não tiver sido iniciado, consulte a seção Resolver problemas para mais informações.

Para o modelo de 70 bilhões, repita essas etapas para cada VM de TPU:

  1. Conecte-se à TPU usando SSH:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --worker=WORKER_NUMBER
    

    WORKER_NUMBER é um índice baseado em 0, indicando a qual VM de TPU você quer se conectar.

  2. Verificar os registros:

    sudo docker logs -f sax-model-server
    

    Três VMs de TPU vão mostrar que estão conectadas às outras instâncias:

    I1117 00:16:07.196594 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.3:10001
    I1117 00:16:07.197484 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.87:10001
    I1117 00:16:07.199437 140613973207936 multi_host_sync.py:152] Received SPMD peer address 10.182.0.13:10001
    

    Uma das VMs da TPU precisa ter registros que mostrem o início do servidor do modelo:

    I1115 04:01:29.479170 139974275995200 model_service_base.py:867] Started joining SAX cell /sax/test
    ERROR: logging before flag.Parse: I1115 04:01:31.479794       1 location.go:141] Calling Join due to address update
    ERROR: logging before flag.Parse: I1115 04:01:31.814721       1 location.go:155] Joined 10.182.0.44:10000
    

Publicar o modelo

O SAX vem com uma ferramenta de linha de comando chamada saxutil, que simplifica a interação com os servidores de modelos do SAX. Neste tutorial, você usa saxutil para publicar o modelo. Para ver a lista completa de comandos saxutil, consulte o arquivo README saxml.

  1. Altere para o diretório em que você clonou o repositório Saxml do GitHub:

    cd  saxml
    
  2. Para o modelo 70B, conecte-se ao seu servidor de conversão:

    gcloud compute ssh ${CONV_SERVER_NAME} \
      --project ${PROJECT_ID} \
      --zone ${ZONE}
    
  3. Instale o Bazel:

    sudo apt-get install bazel
    
  4. Defina um alias para executar saxutil com seu bucket do Cloud Storage:

    alias saxutil='bazel run saxml/bin:saxutil -- --sax_root=gs://${GSBUCKET}/sax-root'
    
  5. Publique o modelo usando saxutil. Isso leva cerca de 10 minutos em um TPU v5litepod-8.

    saxutil --sax_root=gs://${GSBUCKET}/sax-root publish '/sax/test/MODEL' \
        saxml.server.pax.lm.params.lm_cloud.PARAMETERS \
        gs://${GSBUCKET}/sax_models/llama2/SAX_LLAMA2_DIR/checkpoint_00000000/ \
        1
    

    Substitua as seguintes variáveis:

    Tamanho do modelo Valores
    7 bi MODEL: llama27b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama27b
    13 bi MODEL: llama213b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA13BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama213b
    70 bi MODEL: llama270b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA70BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama270b

Testar implantação

Para verificar se a implantação foi bem-sucedida, use o comando saxutil ls:

saxutil ls /sax/test/MODEL

Uma implantação bem-sucedida precisa ter um número de réplicas maior que zero e ser semelhante a esta:

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| MODEL    | MODEL PATH                                            | CHECKPOINT PATH                                                       | # OF REPLICAS | (SELECTED) REPLICAADDRESS |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+
| llama27b | saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e | gs://${MODEL_BUCKET}/sax_models/llama2/7b/pax_7B/checkpoint_00000000/ | 1             | 10.182.0.28:10001         |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+

Os registros do Docker para o servidor de modelo serão semelhantes aos seguintes:

I1114 17:31:03.586631 140003787142720 model_service_base.py:532] Successfully loaded model for key: /sax/test/llama27b

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root is /sax/test/1lama27b

Resolver problemas

Se a implantação falhar, verifique os registros do servidor de modelo:

sudo docker logs -f sax-model-server

Para uma implantação bem-sucedida, você verá a seguinte saída:

Successfully loaded model for key: /sax/test/llama27b

Se os registros não mostrarem que o modelo foi implantado, verifique a configuração e o caminho para o checkpoint do modelo.

Gerar respostas

Use a ferramenta saxutil para gerar respostas a comandos.

Gerar respostas para uma pergunta:

saxutil lm.generate -extra="temperature:0.2"  /sax/test/MODEL "Q: Who is Harry Potter's mother? A:"

A saída será semelhante a esta:

INFO: Running command line: bazel-bin/saxml/bin/saxutil_/saxutil '--sax_rootmgs://sax-admin2/sax-root' lm.generate /sax/test/llama27b 'Q: Who is Harry Potter's mother? A: `
+-------------------------------+------------+
| GENERATE                      | SCORE      |
+-------------------------------+------------+
| 1. Harry Potter's mother is   | -20.214787 |
| Lily Evans. 2. Harry Potter's |            |
| mother is Petunia Evans       |            |
| (Dursley).                    |            |
+-------------------------------+------------+

Interagir com o modelo por um cliente

O repositório SAX inclui clientes que podem ser usados para interagir com uma célula SAX. Os clientes estão disponíveis em C++, Python e Go. O exemplo a seguir mostra como criar um cliente Python.

  1. Crie o cliente Python:

    bazel build saxml/client/python:sax.cc --compile_one_dependency
    
  2. Adicione o cliente a PYTHONPATH. Este exemplo pressupõe que você tenha saxml no seu diretório principal:

    export PYTHONPATH=${PYTHONPATH}:$HOME/saxml/bazel-bin/saxml/client/python/
    
  3. Interaja com o SAX no shell do Python:

    $ python3
    Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] on linux
    Type "help", "copyright", "credits" or "license" for more information.
    >>> import sax
    >>>
    

Interagir com o modelo a partir de um endpoint HTTP

Para interagir com o modelo usando um endpoint HTTP, crie um cliente HTTP:

  1. Crie uma VM do Compute Engine:

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    export HTTP_SERVER_NAME=HTTP_SERVER_NAME
    export SERVICE_ACCOUNT=SERVICE_ACCOUNT
    export MACHINE_TYPE=e2-standard-8
    gcloud compute instances create $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE \
      --machine-type=$MACHINE_TYPE \
      --network-interface=network-tier=PREMIUM,stack-type=IPV4_ONLY,subnet=default \
      --maintenance-policy=MIGRATE --provisioning-model=STANDARD \
      --service-account=$SERVICE_ACCOUNT \
      --scopes=https://www.googleapis.com/auth/cloud-platform \
      --tags=http-server,https-server \
      --create-disk=auto-delete=yes,boot=yes,device-name=$HTTP_SERVER_NAME,image=projects/ml-images/global/images/c0-deeplearning-common-cpu-v20230925-debian-10,mode=rw,size=500,type=projects/$PROJECT_ID/zones/$ZONE/diskTypes/pd-balanced \
      --no-shielded-secure-boot \
      --shielded-vtpm \
      --shielded-integrity-monitoring \
      --labels=goog-ec-src=vm_add-gcloud \
      --reservation-affinity=any
    
  2. Conecte-se à VM do Compute Engine usando SSH:

    gcloud compute ssh $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE
    
  3. Clone a IA no repositório do GKE no GitHub:

    git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git
    
  4. Mude para o diretório do servidor HTTP:

    cd ai-on-gke/tools/saxml-on-gke/httpserver
    
  5. Crie o arquivo Docker:

    docker build -f Dockerfile -t sax-http .
    
  6. Execute o servidor HTTP:

    docker run -e SAX_ROOT=gs://${GSBUCKET}/sax-root -p 8888:8888 -it sax-http
    

Teste o endpoint na máquina local ou em outro servidor com acesso à porta 8888 usando os seguintes comandos:

  1. Exporte as variáveis de ambiente referentes ao endereço IP e à porta do seu servidor:

    export LB_IP=HTTP_SERVER_EXTERNAL_IP
    export PORT=8888
    
  2. Defina o payload do JSON, contendo o modelo e a consulta:

    json_payload=$(cat  << EOF
    {
      "model": "/sax/test/MODEL",
      "query": "Example query"
    }
    EOF
    )
    
  3. Envie a solicitação:

    curl --request POST --header "Content-type: application/json" -s $LB_IP:$PORT/generate --data "$json_payload"
    

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Quando terminar este tutorial, siga as etapas abaixo para limpar seus recursos.

  1. Exclua o Cloud TPU.

    $ gcloud compute tpus tpu-vm delete $TPU_NAME  --zone $ZONE
    
  2. Exclua a instância do Compute Engine, caso tenha criado uma.

    gcloud compute instances delete INSTANCE_NAME
    
  3. Exclua o bucket do Cloud Storage e o conteúdo dele.

    gcloud storage rm --recursive gs://BUCKET_NAME
    

A seguir