SAX auf Cloud TPU v5e

SAX-Cluster (SAX-Zelle)

Der SAX-Administratorserver und der SAX-Modellserver sind zwei wichtige Komponenten, die einen SAX-Cluster ausführen.

SAX-Administratorserver

Der SAX-Administratorserver überwacht und koordiniert alle SAX-Modellserver in einem SAX-Cluster. In einem SAX-Cluster können Sie mehrere SAX-Administratorserver starten. Nur einer der SAX-Administratorserver ist durch die Leaderauswahl aktiv, die anderen sind Stand‑by-Server. Wenn der aktive Administratorserver ausfällt, wird ein Stand‑by-Administratorserver aktiviert. Der aktive SAX-Administratorserver weist Modellreplikate und Inferenzanfragen verfügbaren SAX-Modellservern zu.

SAX-Administrator-Storage-Bucket

Für jeden SAX-Cluster ist ein Cloud Storage-Bucket erforderlich, in dem die Konfigurationen und Standorte von SAX-Administratorservern und SAX-Modellservern im SAX-Cluster gespeichert werden.

SAX-Modellserver

Der SAX-Modellserver lädt einen Modellprüfpunkt und führt die Inferenz mit GSPMD aus. Ein SAX-Modellserver wird auf einem einzelnen TPU-VM-Worker ausgeführt. Für die Bereitstellung von TPU-Modellen mit einem einzelnen Host ist ein einzelner SAX-Modellserver auf einer TPU-VM mit einem einzelnen Host erforderlich. Für die Bereitstellung von TPU-Modellen mit mehreren Hosts ist eine Gruppe von SAX-Modellservern auf einem TPU-Slice mit mehreren Hosts erforderlich. Die Bereitstellung von Modellen auf mehreren Hosts ist derzeit nicht verfügbar. In diesem Dokument finden Sie jedoch ein Beispiel mit einem 175B-Testmodell als Vorschau.

SAX-Modellbereitstellung

Im folgenden Abschnitt wird der Workflow für die Bereitstellung von Sprachmodellen mit SAX beschrieben. Darin wird das GPT‑J 6B-Modell als Beispiel für die Bereitstellung von Modellen mit einem einzelnen Host verwendet.

Installieren Sie zuerst die Cloud TPU-SAX-Docker-Images auf Ihrer TPU-VM:

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}

Legen Sie einige andere Variablen fest, die Sie später verwenden werden:

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

Beispiel für die Bereitstellung eines GPT‑J 6B-Modells mit einem einzelnen Host

Die Bereitstellung von Modellen mit einem einzelnen Host ist für TPU-Slices mit einem einzelnen Host möglich, also v5litepod‑1, v5litepod‑4 und v5litepod‑8.

  1. Erstellen Sie einen SAX-Cluster:

    1. Erstellen Sie einen Cloud Storage-Bucket für den SAX-Cluster:

      SAX_ADMIN_STORAGE_BUCKET=${your_admin_storage_bucket}

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

      Möglicherweise benötigen Sie einen weiteren Cloud Storage-Bucket, um den Prüfpunkt zu speichern.

      SAX_DATA_STORAGE_BUCKET=${your_data_storage_bucket}

    2. Stellen Sie in einem Terminal eine SSH-Verbindung zu Ihrer TPU-VM her, um den SAX-Administratorserver zu starten:

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

      So können Sie das Docker-Log prüfen: 

      docker logs -f ${SAX_ADMIN_SERVER_DOCKER_NAME}

      Die Ausgabe im Log sieht in etwa so aus:

      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. Starten Sie einen SAX-Modellserver mit einem einzelnen Host im SAX-Cluster:

    Der SAX-Cluster enthält zu diesem Zeitpunkt nur den SAX-Administratorserver. Sie können in einem zweiten Terminal eine SSH-Verbindung zu Ihrer TPU-VM herstellen, um einen SAX-Modellserver in Ihrem SAX-Cluster zu starten:

    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. Konvertieren Sie den Modellprüfpunkt:

    Sie müssen PyTorch und Transformers installieren, um den GPT‑J-Prüfpunkt von EleutherAI herunterzuladen:

    pip3 install accelerate
pip3 install torch
pip3 install transformers

    Um den Prüfpunkt in einen SAX-Prüfpunkt zu konvertieren, müssen Sie paxml installieren:

    pip3 install paxml==1.1.0

    Mit dem folgenden Script wird der GPT‑J-Prüfpunkt in einen SAX-Prüfpunkt konvertiert:

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

    Nach der Konvertierung: 

    ls checkpoint_00000000/

    Sie müssen eine Datei namens „commit_success“ erstellen und in die Unterverzeichnisse einfügen:

    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. Veröffentlichen Sie das Modell im SAX-Cluster:

    Sie können GPT‑J jetzt mit dem im vorherigen Schritt konvertierten Prüfpunkt veröffentlichen.

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

    Um GPT‑J (und die nachfolgenden Schritte) zu veröffentlichen, stellen Sie in einem dritten Terminal eine SSH-Verbindung zu Ihrer TPU-VM her:

    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}

    Im Docker-Log des Modellservers wird viel Aktivität angezeigt, bis Sie etwas wie das Folgende sehen, was darauf hinweist, dass das Modell erfolgreich geladen wurde:

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

  5. Generieren Sie die Inferenzergebnisse:

    Für GPT‑J müssen Eingabe und Ausgabe als kommagetrennte Token-ID-Strings formatiert werden. Sie müssen die Texteingabe tokenisieren.

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

    Sie können den String mit Token-IDs über den Tokenizer von EleutherAI/gpt‑j-6b abrufen:

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

    Tokenisieren Sie den Eingabetext:

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

    Sie erhalten einen Token-ID-String ähnlich dem folgenden:

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

    So erstellen Sie eine Zusammenfassung für Ihren Artikel:

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

    Die Ausgabe sollte in etwa so aussehen:

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

    So detokenisieren Sie die String-ID der Ausgabetokens:

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

    Der detokenisierte Text sieht dann etwa so aus:

    >>> 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. Bereinigen Sie Ihre Docker-Container und Cloud Storage-Buckets.

Vorschau für die Bereitstellung eines 175B-Modells mit mehreren Hosts

Für einige der Large Language Models ist ein TPU-Slice mit mehreren Hosts erforderlich, also v5litepod-16 oder höher. In diesen Fällen muss auf allen TPU-Hosts mit mehreren Hosts eine Kopie eines SAX-Modellservers vorhanden sein. Alle Modellserver fungieren als SAX-Modellservergruppe, um das große Modell auf einem TPU-Slice mit mehreren Hosts bereitzustellen.

  1. Erstellen Sie einen neuen SAX-Cluster:

    Sie können dem Schritt zum Erstellen eines SAX-Clusters aus der GPT‑J-Schritt-für-Schritt-Anleitung folgen, um einen neuen SAX-Cluster und einen SAX-Administratorserver zu erstellen.

    Wenn Sie bereits einen SAX-Cluster haben, können Sie in Ihrem SAX-Cluster einen Modellserver mit mehreren Hosts starten.

  2. Starten Sie einen SAX-Modellserver mit mehreren Hosts in einem SAX-Cluster:

    Verwenden Sie zum Erstellen eines TPU-Slice mit mehreren Hosts denselben Befehl wie für einen TPU-Slice mit einem einzelnen Host. Geben Sie einfach den entsprechenden Beschleunigertyp für mehrere Hosts an:

    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

    So rufen Sie das SAX-Modellserverimage auf alle TPU-Hosts/Worker ab und starten sie:

    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. Veröffentlichen Sie das Modell im SAX-Cluster:

    In diesem Beispiel wird das Modell LmCloudSpmd175B32Test verwendet:

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

    So veröffentlichen Sie das Testmodell:

    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. Generieren Sie die Inferenzergebnisse:

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

    Da in diesem Beispiel ein Testmodell mit zufälligen Gewichten verwendet wird, ist die Ausgabe möglicherweise nicht aussagekräftig.

  5. Führen Sie die Bereinigung durch.

    Beenden Sie die Docker-Container:

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

    Löschen Sie Ihren Cloud Storage-Administrator-Bucket und alle Datenspeicher-Buckets mit der gcloud CLI, wie unten gezeigt.

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