‫SAX ב-Cloud TPU v5e

אשכול SAX (תא SAX)

שרת האדמין של SAX ושרת המודל של SAX הם שני רכיבים חיוניים שמופעלים באשכול SAX.

שרת אדמין של SAX

שרת האדמין של SAX מנטר ומתאם את כל שרתי המודלים של SAX באשכול SAX. ב-SAX cluster, אפשר להפעיל כמה שרתי אדמין של SAX, אבל רק אחד מהם פעיל באמצעות בחירת מנהיג. השרתים האחרים הם שרתי המתנה. אם שרת האדמין הפעיל נכשל, שרת אדמין במצב המתנה יהפוך לפעיל. שרת האדמין הפעיל של SAX מקצה רפליקות של מודלים ובקשות הסקה לשרתי מודלים זמינים של SAX.

קטגוריית אחסון של אדמין SAX

כל אשכול SAX צריך קטגוריה של Cloud Storage כדי לאחסן את ההגדרות והמיקומים של שרתי האדמין ושרתי המודל של SAX באשכול SAX.

שרת מודל SAX

שרת המודלים של SAX טוען נקודת ביקורת של מודל ומריץ הסקה באמצעות GSPMD. שרת מודל SAX פועל במכונת עבודה וירטואלית אחת של TPU. פרסום מודל TPU במארח יחיד דורש שרת מודל SAX יחיד במכונה וירטואלית של TPU במארח יחיד. פרסום מודל TPU בכמה מארחים דורש קבוצה של שרתי מודלים של SAX בפרוסת TPU בכמה מארחים. בשלב הזה אין אפשרות להפעיל מודל בכמה מארחים, אבל במסמך הזה יש דוגמה עם מודל בדיקה של 175 מיליארד פרמטרים בגרסת טרום-השקה (Preview).

פרסום מודל SAX

בקטע הבא מפורט תהליך העבודה להצגת מודלים של שפה באמצעות SAX. הוא משתמש במודל GPT-J 6B כדוגמה להצגת מודל של מארח יחיד.

לפני שמתחילים, צריך להתקין את תמונות ה-Docker של Cloud TPU SAX במכונת ה-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}

מגדירים עוד משתנים שישמשו אתכם בהמשך:

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

דוגמה לפרסום מודל GPT-J 6B במארח יחיד

פרסום מודל במארח יחיד רלוונטי לפרוסת TPU במארח יחיד, כלומר v5litepod-1,‏ v5litepod-4 ו-v5litepod-8.

  1. יצירת אשכול SAX

    1. יוצרים קטגוריית אחסון ב-Cloud Storage עבור אשכול SAX:

      SAX_ADMIN_STORAGE_BUCKET=${your_admin_storage_bucket}

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

      יכול להיות שתצטרכו עוד קטגוריית אחסון ב-Cloud Storage כדי לאחסן את נקודת הבדיקה.

      SAX_DATA_STORAGE_BUCKET=${your_data_storage_bucket}

    2. מתחברים ב-SSH למכונת ה-TPU הווירטואלית בטרמינל כדי להפעיל את שרת האדמין של 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}

      כדי לבדוק את יומן הרישום של Docker: 

      docker logs -f ${SAX_ADMIN_SERVER_DOCKER_NAME}

      הפלט ביומן ייראה כך:

      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. מפעילים שרת מודל SAX עם מארח יחיד באשכול SAX:

    בשלב הזה, אשכול ה-SAX מכיל רק את שרת האדמין של ה-SAX. אתם יכולים להתחבר למכונת ה-TPU הווירטואלית באמצעות SSH בטרמינל שני כדי להפעיל שרת מודלים של SAX באשכול 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. המרת נקודת ביקורת של מודל:

    כדי להוריד את נקודת הבדיקה של GPT-J מ-EleutherAI, צריך להתקין את PyTorch ואת Transformers:

    pip3 install accelerate
pip3 install torch
pip3 install transformers

    כדי להמיר את נקודת הבדיקה לנקודת בדיקה של SAX, צריך להתקין את paxml:

    pip3 install paxml==1.1.0

    הסקריפט הבא ממיר את נקודת הבדיקה של GPT-J לנקודת בדיקה של SAX:

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

    אחרי שההמרה מסתיימת: 

    ls checkpoint_00000000/

    צריך ליצור קובץ commit_success ולמקם אותו בספריות המשנה:

    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. פרסום המודל באשכול SAX

    עכשיו אפשר לפרסם את GPT-J עם נקודת הבדיקה שהומרה בשלב הקודם.

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

    כדי לפרסם את GPT-J (ואת השלבים הבאים), משתמשים ב-SSH כדי להתחבר למכונת ה-TPU הווירטואלית בטרמינל השלישי:

    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}

    תראו הרבה פעילות ביומן Docker של שרת המודל, עד שתראו משהו כמו הדוגמה הבאה שמציין שהמודל נטען בהצלחה:

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

  5. יצירת תוצאות של הסקה

    ב-GPT-J, הקלט והפלט צריכים להיות בפורמט של מחרוזת מזהי טוקנים מופרדת בפסיקים. תצטרכו ליצור טוקניזציה של קלט הטקסט.

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

    אפשר לקבל את מחרוזת מזהי האסימונים באמצעות tokenizer של EleutherAI/gpt-j-6b:

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

    פיצול טקסט הקלט לטוקנים:

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

    מחרוזת של מזהה אסימון שנראית בערך כך:

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

    כדי ליצור סיכום של המאמר:

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

    התוצאה תהיה דומה לזו:

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

    כדי לבטל את הטוקניזציה של מחרוזת מזהי הטוקנים של הפלט:

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

    הטקסט שיוחזר אחרי ההסרה של הטוקנים יהיה:

    >>> 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 ובקטגוריות של Cloud Storage.

תצוגה מקדימה של פרסום מודל עם 175 מיליארד פרמטרים בכמה מארחים

חלק מהמודלים הגדולים של שפה ידרשו פרוסת TPU מרובת מארחים, כלומר, v5litepod-16 ומעלה. במקרים כאלה, כל מארחי ה-TPU מרובי המארחים יצטרכו להכיל עותק של שרת מודל SAX, וכל שרתי המודל יפעלו כקבוצת שרתי מודל SAX כדי להציג את המודל הגדול בפרוסת TPU מרובת מארחים.

  1. יצירת אשכול SAX חדש

    אפשר לפעול לפי השלב 'יצירת אשכול SAX' בהמדריך ל-GPT-J כדי ליצור אשכול SAX חדש ושרת אדמין SAX.

    לחלופין, אם כבר יש לכם אשכול SAX, אתם יכולים להפעיל שרת מודלים מרובה מארחים באשכול SAX.

  2. הפעלת שרת מודלים של SAX מרובה מארחים באשכול SAX

    משתמשים באותה פקודה שבה משתמשים כדי ליצור פלח TPU במארח יחיד, רק מציינים את סוג המאיץ המתאים למארחים מרובים:

    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

    כדי למשוך את תמונת השרת של מודל SAX לכל המארחים או העובדים של TPU ולהפעיל אותם:

    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. פרסום המודל באשכול SAX

    בדוגמה הזו נעשה שימוש במודל LmCloudSpmd175B32Test:

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

    כדי לפרסם את מודל הבדיקה:

    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. יצירת תוצאות של הסקה

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

    שימו לב: בדוגמה הזו נעשה שימוש במודל בדיקה עם משקלים אקראיים, ולכן יכול להיות שהפלט לא יהיה משמעותי.

  5. הסרת המשאבים

    עוצרים את קונטיינרים של Docker:

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

    כדי למחוק את קטגוריית האחסון של האדמין ב-Cloud Storage וכל קטגוריית אחסון נתונים, משתמשים ב-CLI של gcloud כמו בדוגמה שלמטה.

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