SAX en Cloud TPU v5e

Clúster de SAX (celda de SAX)

El servidor de administrador de SAX y el servidor de modelos de SAX son dos componentes esenciales que ejecutan un clúster de SAX.

Servidor de administrador de SAX

El servidor de administrador de SAX supervisa y coordina todos los servidores de modelos de SAX en un clúster de SAX. Allí, puedes iniciar varios servidores de administrador de SAX, en los que solo uno de ellos está activo a través de la elección de líder. Los demás son servidores en espera. Cuando falle el servidor de administrador activo, se activará un servidor de administrador en espera. El servidor de administrador de SAX activo asigna réplicas de modelos y solicitudes de inferencia a los servidores de modelos de SAX disponibles.

Bucket de almacenamiento del administrador de SAX

Cada clúster de SAX requiere un bucket de Cloud Storage para almacenar los parámetros de configuración y las ubicaciones de los servidores de administrador y los servidores de modelos de SAX en el clúster de SAX.

Servidor de modelos de SAX

El servidor de modelos de SAX carga un punto de control del modelo y ejecuta la inferencia con GSPMD. Un servidor de modelos de SAX se ejecuta en un solo trabajador de VM de TPU. La entrega de modelos de TPU de host único requiere un solo servidor de modelos de SAX en una VM de TPU de host único. La entrega de modelos de TPU de varios hosts requiere un grupo de servidores de modelos de SAX en una porción de TPU de varios hosts. Actualmente, no está disponible la entrega de modelos en varios hosts, pero este documento proporciona un ejemplo con un modelo de prueba 175B para obtener una vista previa.

Entrega de modelos de SAX

En la siguiente sección, se explica el flujo de trabajo para entregar modelos de lenguaje con SAX. Utiliza el modelo GPT-J 6B como ejemplo de entrega de modelos de host único.

Antes de comenzar, instala las imágenes de Docker de SAX de Cloud TPU en tu VM de 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}

Configura otras variables que usarás más adelante:

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

Ejemplo de entrega del modelo GPT-J 6B de un solo host

La entrega de modelos de host único se aplica a la porción de TPU de host único, es decir, v5litepod-1, v5litepod-4 y v5litepod-8.

  1. Crea un clúster de SAX.

    1. Crea un bucket de almacenamiento de Cloud Storage para el clúster de SAX:

      SAX_ADMIN_STORAGE_BUCKET=${your_admin_storage_bucket}

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

      Es posible que necesites otro bucket de almacenamiento de Cloud Storage para almacenar el punto de control.

      SAX_DATA_STORAGE_BUCKET=${your_data_storage_bucket}

    2. Establece una conexión SSH a tu VM de TPU en una terminal para iniciar el servidor de administración de 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}

      Puedes verificar el registro de Docker de la siguiente manera: 

      docker logs -f ${SAX_ADMIN_SERVER_DOCKER_NAME}

      El resultado en el registro será similar al siguiente:

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

  2. Inicia un servidor de modelos de SAX de host único en el clúster de SAX:

    En este punto, el clúster de SAX solo contiene el servidor de administrador de SAX. Puedes conectarte a tu VM de TPU a través de SSH en una segunda terminal para iniciar un servidor de modelos de SAX en tu clúster de 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

  3. Convierte el punto de control del modelo:

    Debes instalar PyTorch y Transformers para descargar el punto de control de GPT-J de EleutherAI:

    pip3 install accelerate
pip3 install torch
pip3 install transformers

    Para convertir el punto de control en un punto de control de SAX, debes instalar paxml:

    pip3 install paxml==1.1.0

    La siguiente secuencia de comandos convierte el punto de control de GPT-J en un punto de control de SAX:

    python3 -m convert_gptj_ckpt --base EleutherAI/gpt-j-6b --pax pax_6b

    Después de la conversión, haz lo siguiente: 

    ls checkpoint_00000000/

    Debes crear un archivo commit_success y colocarlo en los subdirectorios:

    gcloud storage cp checkpoint_00000000 ${CHECKPOINT_PATH} --recursive

touch commit_success.txt
gcloud storage cp commit_success.txt ${CHECKPOINT_PATH}/
gcloud storage cp commit_success.txt ${CHECKPOINT_PATH}/metadata/
gcloud storage cp commit_success.txt ${CHECKPOINT_PATH}/state/

  4. Publica el modelo en el clúster de SAX.

    Ahora puedes publicar GPT-J con el punto de control convertido en el paso anterior.

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

    Para publicar GPT-J (y los pasos posteriores), usa SSH y conectate a tu VM de TPU en una tercera 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}

    Verás mucha actividad en el registro de Docker del servidor de modelos hasta que veas algo como lo siguiente, que indicará que el modelo se cargó correctamente:

    I0829 01:33:49.287459 139865140229696 servable_model.py:697] loading completed.

  5. Genera resultados de inferencia.

    En el caso de GPT-J, la entrada y salida deben tener el formato de una cadena de ID de tokens separados por comas. Deberás tokenizar la 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\:

    Puedes obtener la cadena de IDs de token a través del tokenizador EleutherAI/gpt-j-6b:

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

    Asigna tokens al 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])

    Puedes esperar una cadena de ID de token similar a la siguiente:

    >>> 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 generar un resumen de tu artículo, sigue estos pasos:

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

    Verás algo similar a lo siguiente:

    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
|                                                                                                                                                    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 destokenizar la cadena de IDs de tokens de salida, haz lo siguiente:

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

    Puedes esperar el siguiente texto sin tokenizar:

    >>> OUTPUT_TEXT
'This page includes the show Transcript.\nUse the Transcript to help
students with reading comprehension and vocabulary.\nAt the bottom of
the page, comment for a chance to be mentioned on CNN Student News.
You must be a teacher or a student age 13 or older to request a mention on the CNN Student News Roll Call.'

  6. Limpia tus contenedores de Docker y buckets de almacenamiento de Cloud Storage.

Vista previa de la entrega de modelos multihost 175B

Algunos de los modelos de lenguaje grandes requerirán una porción de TPU de varios hosts, es decir, v5litepod-16 y versiones posteriores. En esos casos, todos los hosts de TPU de varios hosts deberán tener una copia de un servidor de modelos de SAX, y todos los servidores de modelos funcionarán como un grupo de servidores de modelos de SAX para entregar el modelo grande en una porción de TPU de varios hosts.

  1. Crea un clúster de SAX nuevo.

    Puedes seguir el mismo paso sobre crear un clúster de SAX en la explicación de GPT-J para crear un nuevo clúster de SAX y un servidor de administrador de SAX.

    O bien, si ya tienes un clúster de SAX existente, puedes iniciar un servidor de modelos de varios hosts en tu clúster de SAX.

  2. Inicia un servidor de modelos de SAX de varios hosts en un clúster de SAX.

    Usa el mismo comando para crear una porción de TPU, ya sea de varios hosts o de un solo host. Solo debes especificar el tipo de acelerador de varios hosts adecuado:

    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 extraer la imagen del servidor de modelos de SAX a todos los hosts o trabajadores TPU y, luego, inicialos, haz lo siguiente:

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

  3. Publica el modelo en el clúster de SAX.

    En este ejemplo, se usa un modelo LmCloudSpmd175B32Test:

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

    Para publicar el modelo de prueba, haz lo siguiente:

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

  4. Genera resultados de inferencia.

    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\: "

    Ten en cuenta que, dado que en este ejemplo se muestra el uso de un modelo de prueba con pesos aleatorios, es posible que el resultado no sea significativo.

  5. Realiza una limpieza.

    Detén los contenedores de Docker:

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

    Borra tu bucket de almacenamiento de administrador de Cloud Storage y cualquier bucket de almacenamiento de datos con gcloud CLI, como se muestra a continuación.

    gcloud storage rm gs://${SAX_ADMIN_STORAGE_BUCKET} --recursive --continue-on-error
gcloud storage rm gs://${SAX_DATA_STORAGE_BUCKET} --recursive --continue-on-error