Implementa y entrega modelos de Llama 2 en TPU v5e con SAX

En este instructivo, se describe cómo realizar las siguientes acciones:

  • Crear una VM de Cloud TPU para implementar la familia Llama 2 de modelos grandes de lenguaje (LLM), disponibles en diferentes tamaños (7B, 13B o 70B)
  • Preparación de puntos de control para los modelos e implementarlos en SAX
  • Interactúa con el modelo a través de un extremo HTTP

La entrega para experimentos de AGI (SAX) es un sistema experimental que entrega modelos Paxml, JAX y PyTorch para inferencia. El código y la documentación para SAX se encuentran en el repositorio de Git en Samxml. La versión estable actual compatible con TPU v5e es v1.1.0.

Acerca de las celdas SAX

Una celda (o clúster) SAX es la unidad principal para entregar tus modelos. Consta de dos partes principales:

  • Servidor de administración: Este servidor realiza un seguimiento de tus servidores de modelos, asigna modelos a esos servidores y ayuda a los clientes a encontrar el servidor de modelos correcto con el que interactuar.
  • Servidores de modelos: Estos servidores ejecutan el modelo. Son responsables de procesar las solicitudes entrantes y generar respuestas.

En el siguiente diagrama, se muestra el diagrama de una celda SAX:

Celda SAX con servidor de administración y servidores de modelos.

Figura 1. Celda SAX con servidor de administrador y servidor de modelos.

Puedes interactuar con una celda SAX mediante clientes escritos en Python, C++ o Go, o directamente a través de un servidor HTTP. En el siguiente diagrama, se muestra cómo un cliente externo puede interactuar con una celda SAX:

Un cliente externo y un extremo HTTP que interactúan con una celda SAX

Figura 2. Arquitectura del entorno de ejecución de un cliente externo que interactúa con una celda SAX.

Objetivos

  • Configura recursos de TPU para la entrega
  • Crea un clúster de SAX
  • Publica el modelo de Llama 2
  • Interactúa con el modelo

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

  • Cloud TPU
  • Compute Engine
  • Cloud Storage

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios. Es posible que los usuarios nuevos de Google Cloud califiquen para obtener una prueba gratuita.

Antes de comenzar

Sigue las instrucciones en Configura el entorno de Cloud TPU para configurar tu proyecto de Google Cloud, activar la API de Cloud TPU y crear una cuenta de servicio.

Crea una TPU

En los siguientes pasos, se muestra cómo crear una VM de TPU que entregará el modelo.

  1. Crea variables de entorno:

    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

    Descripciones de las variables de entorno

    PROJECT_ID
    El ID de tu proyecto de Google Cloud.
    ACCELERATOR_TYPE
    El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Los diferentes tamaños del modelo de Llama 2 tienen diferentes requisitos de tamaño de TPU:
    • 7B: v5litepod-4 o superior
    • 13B: v5litepod-8 o superior
    • 70B: v5litepod-16 o superior
    ZONE
    La zona en la que deseas crear tu Cloud TPU.
    SERVICE_ACCOUNT
    La cuenta de servicio que deseas conectar a tu Cloud TPU.
    TPU_NAME
    Es el nombre de tu Cloud TPU.
    QUEUED_RESOURCE_ID
    Es un identificador para tu solicitud de recurso en cola.

  2. Establece el ID del proyecto y la zona en tu configuración activa de Google Cloud CLI:

    gcloud config set project $PROJECT_ID  && gcloud config set compute/zone $ZONE

  3. Crea la VM de TPU:

    gcloud 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. Comprueba que la TPU esté activa:

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

Configura el nodo de conversión de puntos de control

Para ejecutar los modelos de LLama en un clúster SAX, debes convertir los puntos de control originales de Llama en un formato compatible con SAX.

La conversión requiere una cantidad importante de recursos de memoria, según el tamaño del modelo:

Modelo Tipo de máquina
7,000 millones De 50 a 60 GB de memoria
13,000 millones 120 GB de memoria
70,000 millones De 500 a 600 GB de memoria (tipo de máquina N2 o M1)

Para los modelos 7B y 13B, puedes ejecutar la conversión en la VM de TPU. Para el modelo 70B, debes crear una instancia de Compute Engine con alrededor de 1 TB de espacio en 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

Ya sea que uses una instancia de TPU o Compute Engine como servidor de conversión, configura el servidor para convertir los puntos de control de Llama 2:

  1. Para los modelos 7B y 13B, configura la variable de entorno de nombre del servidor con el nombre de tu TPU:

    export CONV_SERVER_NAME=$TPU_NAME

    Para el modelo 70B, configura la variable de entorno de nombre del servidor como el nombre de tu instancia de Compute Engine:

    export CONV_SERVER_NAME=INSTANCE_NAME

  2. Conéctate al nodo de conversión mediante SSH.

    Si tu nodo de conversión es una TPU, conéctate a la TPU:

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

    Si tu nodo de conversión es una instancia de Compute Engine, conéctate a la VM de Compute Engine:

    gcloud compute ssh $CONV_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

  3. Instala los paquetes obligatorios en el nodo de conversión:

    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. Descarga el guion de conversión del punto de control de Llama:

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

Descargar pesos de Llama 2

Antes de convertir el modelo, debes descargar los pesos de Llama 2. Para este instructivo, debes usar los pesos originales de Llama 2 (por ejemplo, meta-llama/Llama-2-7b) y no los que se convirtieron para el formato de Transformers de Hugging Face (por ejemplo, meta-llama/Llama-2-7b-hf).

Si ya tienes los pesos de Llama 2, avanza a Convertir los pesos.

Para descargar los pesos del centro de Hugging Face, debes configurar un token de acceso de usuario y solicitar acceso a los modelos de Llama 2. Para solicitar acceso, sigue las instrucciones de la página Hugging Face del modelo que quieres usar, por ejemplo, meta-llama/Llama-2-7b.

  1. Crea un directorio para los pesos:

    sudo mkdir WEIGHTS_DIRECTORY

  2. Obtén los pesos de Llama2 del centro de Hugging Face:

    1. Instala la CLI de Hugging Face:

      pip install -U "huggingface_hub[cli]"

    2. Cambia al directorio de pesos:

      cd WEIGHTS_DIRECTORY

    3. Descarga los archivos de 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)

      Reemplaza LLAMA2_REPO por el nombre del repositorio de Hugging Face del que deseas descargar: Llama-2-7b, Llama-2-13b o Llama-2-70b.

Convierte las ponderaciones

Edita la secuencia de comandos de conversiones y, luego, ejecútala para convertir los pesos del modelo.

  1. Crea un directorio para contener las ponderaciones convertidas:

    sudo mkdir CONVERTED_WEIGHTS

  2. Clona el repositorio de Saxml GitHub en un directorio donde tengas permisos de lectura, escritura y ejecución:

      git clone https://github.com/google/saxml.git -b r1.1.0

  3. Cambia al directorio saxml:

    cd saxml

  4. Abre el archivo saxml/tools/convert_llama_ckpt.py.

  5. En el archivo saxml/tools/convert_llama_ckpt.py, cambia la línea 169 de la siguiente manera:

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

    A:

    'scale': pytorch_vars[0]['norm.weight'].type(torch.float16).numpy()

  6. Ejecuta la secuencia de comandos saxml/tools/init_cloud_vm.sh:

    saxml/tools/init_cloud_vm.sh

  7. Solo para 70B: Para desactivar el modo de prueba, haz lo siguiente:

    1. Abre el archivo saxml/server/pax/lm/params/lm_cloud.py.

    2. En el archivo saxml/server/pax/lm/params/lm_cloud.py, cambia la línea 344 de la siguiente manera:

      return True

      A:

      return False

  8. Convierte las ponderaciones:

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

    Reemplaza lo siguiente:

    • WEIGHTS_DIRECTORY: Directorio para las ponderaciones originales
    • CONVERTED_WEIGHTS: Ruta de destino para las ponderaciones convertidas.
    • MODEL_SIZE: 7b, 13b o 70b.

Prepara el directorio del punto de control

Después de convertir los puntos de control, el directorio de puntos de control debe tener la siguiente estructura:

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

Crea un archivo vacío llamado commit_success.txt y pon una copia de este en los directorios checkpoint_00000000, metadata y state. Esto le permite a SAX saber que este punto de control está completamente convertido y listo para cargarse:

  1. Cambia al directorio del punto de control:

    cd CONVERTED_WEIGHTS/checkpoint_00000000

  2. Crea un nuevo archivo llamado commit_success.txt:

    touch commit_success.txt

  3. Cambia al directorio de metadatos y crea un archivo vacío llamado commit_success.txt:

    cd metadata && touch commit_success.txt

  4. Cambia al directorio de estado y crea un archivo vacío llamado commit_success.txt:

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

El directorio del punto de control ahora debería tener la siguiente estructura:

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

Crear un bucket de Cloud Storage

Debes almacenar los puntos de control convertidos en un bucket de Cloud Storage para que estén disponibles cuando publiques el modelo.

  1. Establece una variable de entorno para el nombre de tu bucket de Cloud Storage:

    export GSBUCKET=BUCKET_NAME

  2. Crea un bucket:

    gcloud storage buckets create gs://${GSBUCKET}

  3. Copia los archivos de punto de control convertidos en tu bucket:

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

    Reemplaza SAX_LLAMA2_DIR por el valor adecuado:

    • 7B: saxml_llama27b
    • 13,000 millones: saxml_llama213b
    • 70,000 millones: saxml_llama270b

Crea un clúster de SAX

Para crear un clúster de SAX, debes hacer lo siguiente:

En una implementación típica, ejecutarías el servidor de administrador en una instancia de Compute Engine y el servidor de modelos en una TPU o GPU. A los fines de este instructivo, implementarás el servidor de administrador y el servidor de modelos en la misma instancia de TPU v5e.

Crear servidor de administrador

Crea el contenedor de Docker del servidor de administrador:

  1. En el servidor de conversión, instala Docker:

    sudo apt-get update
sudo apt-get install docker.io

  2. Inicia el contenedor de Docker del 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

Puedes ejecutar el comando docker run sin la opción -d para ver los registros y asegurarte de que el servidor de administrador se inicie de forma correcta.

Crear servidor de modelos

En las siguientes secciones, se muestra cómo crear un servidor de modelos.

Modelo 7b

Inicia el contenedor de Docker del 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

Falta la configuración de LLaMA13BFP16TPUv5e en lm_cloud.py. En los siguientes pasos, se muestra cómo actualizar lm_cloud.py y confirmar una imagen de Docker nueva.

  1. Inicia el 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='8'

  2. Conéctate al contenedor de Docker con SSH:

    sudo docker exec -it sax-model-server bash

  3. Instala Vim en la imagen de Docker:

    $ apt update
$ apt install vim

  4. Abre el archivo saxml/server/pax/lm/params/lm_cloud.py. Busca LLaMA13B. Deberías ver el siguiente código:

    @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. Comenta o borra la línea que comienza con @quantization. Después de este cambio, el archivo debería verse de la siguiente manera:

    @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. Agrega el siguiente código para admitir la configuración de la 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. Sal de la sesión de SSH del contenedor de Docker:

    exit

  8. Confirma los cambios en una nueva imagen de Docker:

    sudo docker commit sax-model-server sax-model-server:v1.1.0-mod

  9. Verifica que se haya creado la nueva imagen de Docker:

    sudo docker images

    Puedes publicar la imagen de Docker en Artifact Registry de tu proyecto, pero este instructivo continuará con la imagen local.

  10. Detén el servidor del modelo. En el resto del instructivo, se usará el servidor de modelos actualizado.

    sudo docker stop sax-model-server

  11. Inicia el servidor de modelos con la imagen de Docker actualizada. Asegúrate de especificar el nombre de imagen actualizado, 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

Conéctate a la TPU con SSH y, luego, inicia el servidor del 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'
"

Verifica los registros

Verifica los registros del servidor de modelos para asegurarte de que este se haya iniciado de forma correcta:

docker logs -f sax-model-server

Si el servidor de modelo no se inició, consulta la sección Cómo solucionar problemas para obtener más información.

Para el modelo 70B, repite estos pasos para cada VM de TPU:

  1. Conéctate a la TPU con SSH:

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

    WORKER_NUMBER es un índice basado en 0 que indica a qué VM de TPU deseas conectarte.

  2. Verifica los registros:

    sudo docker logs -f sax-model-server

    Tres VM de TPU deberían mostrar que se conectaron a las otras instancias:

    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

    Una de las VMs de TPU debería tener registros que muestren cómo se inicia el servidor de modelos:

    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

Publica el modelo

SAX incluye una herramienta de línea de comandos llamada saxutil, que simplifica la interacción con los servidores de modelos SAX. En este instructivo, usarás saxutil para publicar el modelo. Para obtener la lista completa de los comandos de saxutil, consulta el archivo README de Saxml.

  1. Cambia al directorio en el que clonaste el repositorio de Saxml de GitHub:

    cd  saxml

  2. Para el modelo 70B, conéctate a tu servidor de conversiones:

    gcloud compute ssh ${CONV_SERVER_NAME} \
  --project ${PROJECT_ID} \
  --zone ${ZONE}

  3. Instala Bazel:

    sudo apt-get install bazel

  4. Establece un alias para ejecutar saxutil con tu bucket de Cloud Storage:

    alias saxutil='bazel run saxml/bin:saxutil -- --sax_root=gs://${GSBUCKET}/sax-root'

  5. Publica el modelo con saxutil. Esto tarda alrededor de 10 minutos en una 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

    Reemplaza las siguientes variables:

    Tamaño del modelo Valores
    7,000 millones MODEL: llama27b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA7BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama27b
    13,000 millones MODEL: llama213b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA13BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama213b
    70,000 millones MODEL: llama270b
    PARAMETERS: saxml.server.pax.lm.params.lm_cloud.LLaMA70BFP16TPUv5e
    SAX_LLAMA2_DIR: saxml_llama270b

Implementación de prueba

Para verificar si la implementación se realizó correctamente, usa el comando saxutil ls:

saxutil ls /sax/test/MODEL

Una implementación exitosa debe tener una cantidad de réplicas mayor que cero y ser similar a lo siguiente:

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         |
+----------+-------------------------------------------------------+-----------------------------------------------------------------------+---------------+---------------------------+

Los registros de Docker para el servidor de modelos serán similares a los siguientes:

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

Solucionar problemas

Si la implementación falla, verifica los registros del servidor de modelos:

sudo docker logs -f sax-model-server

Para que la implementación se realice de forma correcta, deberías ver el siguiente resultado:

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

Si los registros no muestran que el modelo se implementó, verifica su configuración y la ruta de acceso a tu punto de control del modelo.

Genera respuestas

Puedes usar la herramienta saxutil para generar respuestas a las instrucciones.

Genera respuestas a una pregunta:

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

El resultado debería ser similar al siguiente ejemplo:

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).                    |            |
+-------------------------------+------------+

Interactúa con el modelo desde un cliente

El repositorio SAX incluye clientes que puedes usar para interactuar con una celda SAX. Los clientes están disponibles en C++, Python y Go. En el siguiente ejemplo, se muestra cómo compilar un cliente Python.

  1. Compila el cliente de Python:

    bazel build saxml/client/python:sax.cc --compile_one_dependency

  2. Agrega el cliente a PYTHONPATH. En este ejemplo, se supone que tienes saxml en tu directorio principal:

    export PYTHONPATH=${PYTHONPATH}:$HOME/saxml/bazel-bin/saxml/client/python/

  3. Interactúa con SAX desde la shell de 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
>>>

Interactúa con el modelo desde un extremo HTTP

Para interactuar con el modelo desde un extremo HTTP, crea un cliente HTTP:

  1. Crea una VM de 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. Conéctate a la VM de Compute Engine mediante SSH:

    gcloud compute ssh $HTTP_SERVER_NAME --project=$PROJECT_ID --zone=$ZONE

  3. Clona la IA en el repositorio de GitHub de GKE:

    git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

  4. Cambia al directorio del servidor HTTP:

    cd ai-on-gke/tools/saxml-on-gke/httpserver

  5. Compila el archivo Docker:

    docker build -f Dockerfile -t sax-http .

  6. Ejecuta el servidor HTTP:

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

Prueba el extremo desde tu máquina local o desde otro servidor con acceso al puerto 8888 mediante los siguientes comandos:

  1. Exporta las variables de entorno para la dirección IP y el puerto de tu servidor:

    export LB_IP=HTTP_SERVER_EXTERNAL_IP
export PORT=8888

  2. Configura la carga útil de JSON, que contiene el modelo y la consulta:

    json_payload=$(cat  << EOF
{
  "model": "/sax/test/MODEL",
  "query": "Example query"
}
EOF
)

  3. Envía la solicitud:

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

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Cuando termines con este instructivo, sigue estos pasos para limpiar tus recursos.

  1. Borra tu Cloud TPU.

    $ gcloud compute tpus tpu-vm delete $TPU_NAME  --zone $ZONE

  2. Borrar tu instancia de Compute Engine, si creaste una

    gcloud compute instances delete INSTANCE_NAME

  3. Borrar el bucket de Cloud Storage y su contenido

    gcloud storage rm --recursive gs://BUCKET_NAME

¿Qué sigue?