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, wobei nur einer der SAX-Administratorserver bei der Auswahl des Leaders aktiv ist und die anderen Standby-Server sind. Wenn der aktive Admin-Server ausfällt, wird ein Standby-Admin-Server aktiv. Der aktive SAX-Administratorserver weist den verfügbaren SAX-Modellservern Modellreplikate und Inferenzanfragen zu.

SAX-Administratorspeicher-Bucket

Jeder SAX-Cluster benötigt einen Cloud Storage-Bucket, um die Konfigurationen und Standorte von SAX-Administratorservern und SAX-Modellservern im SAX-Cluster zu speichern.

SAX-Modellserver

Der SAX-Modellserver lädt einen Modellprüfpunkt und führt eine Inferenz mit GSPMD aus. Ein SAX-Modellserver wird auf einem einzelnen TPU-VM-Worker ausgeführt. Für die TPU-Modellbereitstellung 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 mit mehreren Hosts ist derzeit nicht verfügbar. Dieses Dokument enthält jedoch ein Beispiel mit einem 175B-Testmodell für die Vorschau.

SAX-Modellbereitstellung

Im folgenden Abschnitt wird der Workflow zum Bereitstellen von Sprachmodellen mit SAX beschrieben. Als Beispiel für die Bereitstellung eines Einzelhost-Modells wird das GPT-J 6B-Modell 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:

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 des GPT-J 6B-Modells mit einem einzelnen Host

Die Modellbereitstellung mit einem einzelnen Host gilt für TPU-Slice mit einem einzelnen Host, d. h. v5litepod-1, v5litepod-4 und v5litepod-8.

  1. SAX-Cluster erstellen

    1. Erstellen Sie einen Cloud Storage-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}

      Zum Speichern des Prüfpunkts benötigen Sie möglicherweise einen anderen Cloud Storage-Storage-Bucket.

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

    An dieser Stelle enthält der SAX-Cluster nur den SAX-Administratorserver. Sie können in einem zweiten Terminal über SSH eine 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. Modellprüfpunkt konvertieren:

    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

    Wenn Sie den Prüfpunkt in einen SAX-Prüfpunkt konvertieren möchten, müssen Sie paxml installieren:

    pip3 install paxml==1.1.0

    Das folgende Skript wandelt den GPT-J-Prüfpunkt in einen SAX-Prüfpunkt um:

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

    Nach der Umwandlung: 

    ls checkpoint_00000000/

    Sie müssen eine commit_success-Datei erstellen und in den untergeordneten Verzeichnissen speichern:

    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/

  4. Modell im SAX-Cluster veröffentlichen

    Sie können jetzt GPT-J veröffentlichen, wobei der Prüfpunkt im vorherigen Schritt konvertiert wurde.

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

    Verwenden Sie SSH, um in einem dritten Terminal eine Verbindung zu Ihrer TPU-VM herzustellen, um GPT-J und die nachfolgenden Schritte zu veröffentlichen:

    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 so lange eine Menge Aktivitäten angezeigt, bis Sie einen Wert wie den folgenden sehen, der angibt, dass das Modell erfolgreich geladen wurde:

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

  5. Inferenzergebnisse generieren

    Für GPT-J müssen Eingabe und Ausgabe als kommagetrennter Token-ID-String formatiert werden. Die Texteingabe muss tokenisiert werden.

    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 Token-ID-String über den EleutherAI/gpt-j-6b-Tokenizer abrufen:

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

    Eingabetext tokenisieren:

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

    Der Token-ID-String sieht in etwa so aus:

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

    Sie sollten in etwa Folgendes erwarten:

    +----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+--------------+
|                                                                                                                                                    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 heben Sie die Tokenisierung des Ausgabetoken-IDs-Strings auf:

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

    Sie können den detokenisierten Text so erwarten:

    >>> 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. Docker-Container und Cloud Storage-Storage-Buckets bereinigen

Vorschau des 175-Milliarden-Modells mit mehreren Hosts

Einige der Large Language Models erfordern ein TPU-Slice mit mehreren Hosts, d. h. v5litepod-16 und höher. In diesen Fällen müssen alle TPU-Hosts mit mehreren Hosts eine Kopie eines SAX-Modellservers haben. Alle Modellserver fungieren als SAX-Modellservergruppe, um das große Modell auf einem TPU-Slice mit mehreren Hosts bereitzustellen.

  1. Neuen SAX-Cluster erstellen

    Sie können dieselben Schritte wie unter „SAX-Cluster erstellen“ in der GPT-J-Anleitung ausführen, um einen neuen SAX-Cluster und einen SAX-Administratorserver zu erstellen.

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

  2. SAX-Modellserver mit mehreren Hosts in einem SAX-Cluster starten

    Verwenden Sie denselben Befehl, um ein TPU-Segment mit mehreren Hosts zu erstellen wie für ein TPU-Segment mit einem einzelnen Host. Geben Sie einfach den entsprechenden Beschleunigertyp mit mehreren 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 Server-Image des SAX-Modells auf alle TPU-Hosts/-Worker ab und starten diese:

    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. Modell im SAX-Cluster veröffentlichen

    In diesem Beispiel wird ein LmCloudSpmd175B32Test-Modell 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. Inferenzergebnisse generieren

    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 Gewichtungen verwendet wird, ist die Ausgabe möglicherweise nicht aussagekräftig.

  5. Bereinigen

    Beenden Sie die Docker-Container:

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

    Löschen Sie den Cloud Storage-Administratorspeicher-Bucket und alle Datenspeicher-Buckets mithilfe von gsutil, wie unten gezeigt.

    gsutil rm -rf gs://${SAX_ADMIN_STORAGE_BUCKET}
gsutil rm -rf gs://${SAX_DATA_STORAGE_BUCKET}