הפעלת עומסי עבודה של TPU בקונטיינר ב-Docker

קונטיינרים של Docker מקלים על הגדרת האפליקציות כי הם משלבים את הקוד ואת כל התלויות הנדרשות בחבילה אחת שניתנת להפצה. אתם יכולים להריץ קונטיינרים של Docker במכונות וירטואליות של TPU כדי לפשט את ההגדרה והשיתוף של אפליקציות Cloud TPU. במאמר הזה מוסבר איך להגדיר קונטיינר Docker לכל מסגרת ML שנתמכת על ידי Cloud TPU.

אימון מודל PyTorch בקונטיינר ב-Docker

מכשיר TPU

  1. יצירת מכונת VM של Cloud TPU

    gcloud compute tpus tpu-vm create your-tpu-name \
--zone=europe-west4-a \
--accelerator-type=v2-8 \
--version=tpu-ubuntu2204-base

  2. התחברות למכונת TPU וירטואלית באמצעות SSH

    gcloud compute tpus tpu-vm ssh your-tpu-name \
--zone=europe-west4-a

  3. מוודאים שלמשתמש Google Cloud הוקצה התפקיד Artifact Registry Reader. מידע נוסף זמין במאמר הקצאת תפקידים ב-Artifact Registry.

  4. הפעלת קונטיינר במכונה וירטואלית של TPU באמצעות קובץ אימג' של PyTorch/XLA מגרסת הלילה

    sudo docker run --net=host -ti --rm --name your-container-name --privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
bash

  5. הגדרת זמן ריצה של TPU

    יש שתי אפשרויות לזמן ריצה של PyTorch/XLA: ‏ PJRT ו-XRT. מומלץ להשתמש ב-PJRT אלא אם יש לכם סיבה להשתמש ב-XRT. מידע נוסף על ההגדרות השונות של זמן הריצה זמין במסמכי התיעוד של זמן הריצה של PJRT.

    PJRT

    export PJRT_DEVICE=TPU

    XRT

    export XRT_TPU_CONFIG="localservice;0;localhost:51011"

  6. שכפול המאגר של PyTorch XLA

    git clone --recursive https://github.com/pytorch/xla.git

  7. Train ResNet50

    python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

בסיום סקריפט ההדרכה, חשוב למחוק את המשאבים.

  1. מקלידים exit כדי לצאת ממאגר Docker
  2. מקלידים exit כדי לצאת מ-TPU VM

  3. מחיקת מכונת ה-TPU הווירטואלית

    gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

TPU slice

כשמריצים קוד PyTorch בפרוסת TPU, צריך להריץ את הקוד בכל העובדים של ה-TPU בו-זמנית. אחת הדרכים לעשות זאת היא באמצעות הפקודה gcloud compute tpus tpu-vm ssh עם הדגלים --worker=all ו---command. בתהליך הבא מוסבר איך ליצור תמונת Docker כדי להקל על ההגדרה של כל עובד TPU.

  1. יצירת מכונת TPU וירטואלית

    gcloud compute tpus tpu-vm create your-tpu-name \
--zone=us-central2-b \
--accelerator-type=v4-32 \
--version=tpu-ubuntu2204-base

  2. הוספת המשתמש הנוכחי לקבוצת Docker

    gcloud compute tpus tpu-vm ssh your-tpu-name \
--zone=us-central2-b \
--worker=all \
--command='sudo usermod -a -G docker $USER'

  3. שכפול המאגר של PyTorch XLA

    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=us-central2-b \
--command="git clone --recursive https://github.com/pytorch/xla.git"

  4. הרצת סקריפט האימון בקונטיינר בכל ה-TPU workers

    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=us-central2-b \
--command="docker run --rm --privileged --net=host  -v ~/xla:/xla -e PJRT_DEVICE=TPU us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 python /xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

    התרעות לגבי פקודות Docker:

    • ‫--rm מסיר את הקונטיינר אחרי שהתהליך שלו מסתיים.
    • ‫--privileged חושף את מכשיר ה-TPU למאגר התגים.
    • ‫--net=host מאגד את כל היציאות של הקונטיינר למכונת ה-TPU הווירטואלית כדי לאפשר תקשורת בין המארחים ב-pod.
    • ‫-e מגדיר משתני סביבה.

בסיום סקריפט ההדרכה, חשוב למחוק את המשאבים.

מוחקים את מכונת ה-TPU הווירטואלית באמצעות הפקודה הבאה:

gcloud compute tpus tpu-vm delete your-tpu-name \
--zone=us-central2-b

אימון מודל JAX בקונטיינר Docker

מכשיר TPU

  1. יצירת TPU VM

    gcloud compute tpus tpu-vm create your-tpu-name \
--zone=europe-west4-a \
--accelerator-type=v2-8 \
--version=tpu-ubuntu2204-base

  2. התחברות למכונת TPU וירטואלית באמצעות SSH

    gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a

  3. הפעלת Docker daemon ב-TPU VM

    sudo systemctl start docker

  4. הפעלת קונטיינר Docker

    sudo docker run --net=host -ti --rm --name your-container-name \
--privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
bash

  5. התקנה של JAX

    pip install jax[tpu]

  6. התקנה של FLAX

    pip install --upgrade clu
git clone https://github.com/google/flax.git
pip install --user -e flax

  7. התקנת חבילות tensorflow ו-tensorflow-dataset

    pip install tensorflow
pip install tensorflow-datasets

  8. הרצת סקריפט האימון של FLAX MNIST

    cd flax/examples/mnist
python3 main.py --workdir=/tmp/mnist \
--config=configs/default.py \
--config.learning_rate=0.05 \
--config.num_epochs=5

בסיום סקריפט ההדרכה, חשוב למחוק את המשאבים.

  1. מקלידים exit כדי לצאת ממאגר Docker
  2. מקלידים exit כדי לצאת מ-TPU VM

  3. מחיקת מכונת ה-TPU הווירטואלית

    gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

TPU slice

כשמריצים קוד JAX על פרוסת TPU, צריך להריץ את קוד JAX על כל העובדים של TPU בו-זמנית. אחת הדרכים לעשות זאת היא באמצעות הפקודה gcloud compute tpus tpu-vm ssh עם הדגלים --worker=all ו---command. בתהליך הבא מוסבר איך ליצור קובץ אימג' של Docker כדי להקל על ההגדרה של כל עובד TPU.

  1. יוצרים קובץ בשם Dockerfile בספרייה הנוכחית ומדביקים בו את הטקסט הבא:

    FROM python:3.10
RUN pip install jax[tpu]
RUN pip install --upgrade clu
RUN git clone https://github.com/google/flax.git
RUN pip install --user -e flax
RUN pip install tensorflow
RUN pip install tensorflow-datasets
WORKDIR ./flax/examples/mnist

  2. הכנת Artifact Registry

    gcloud artifacts repositories create your-repo \
--repository-format=docker \
--location=europe-west4 --description="Docker repository" \
--project=your-project

gcloud artifacts repositories list \
--project=your-project

gcloud auth configure-docker europe-west4-docker.pkg.dev

  3. יצירת קובץ האימג' של Docker

    docker build -t your-image-name .

  4. מוסיפים תג לקובץ אימג' של Docker לפני שמעבירים אותה בדחיפה ל-Artifact Registry. מידע נוסף על עבודה עם Artifact Registry זמין במאמר עבודה עם תמונות של קונטיינרים.

    docker tag your-image-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

  5. העברת קובץ אימג' של Docker ל-Artifact Registry

    docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

  6. יצירת מכונת TPU וירטואלית

    gcloud compute tpus tpu-vm create your-tpu-name \
--zone=europe-west4-a \
--accelerator-type=v2-8 \
--version=tpu-ubuntu2204-base

  7. שליפת קובץ האימג' של Docker מ-Artifact Registry בכל העובדים של TPU

    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=europe-west4-a \
--command='sudo usermod -a -G docker ${USER}'
     
    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=europe-west4-a \
--command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
     
    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=europe-west4-a \
--command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"

  8. הרצת המאגר בכל עובדי ה-TPU

    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=europe-west4-a \
--command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag bash"

  9. הפעלת סקריפט האימון בכל ה-TPU workers

    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=europe-west4-a \
--command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
--config=configs/default.py \
--config.learning_rate=0.05 \
--config.num_epochs=5"

בסיום סקריפט ההדרכה, חשוב למחוק את המשאבים.

  1. השבתת הקונטיינר בכל העובדים

    gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
--zone=europe-west4-a \
--command="docker kill your-container-name"

  2. מחיקת מכונת ה-TPU הווירטואלית

    gcloud compute tpus tpu-vm delete your-tpu-name \
--zone=europe-west4-a

המאמרים הבאים