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

שימו לב: הדף הזה מתייחס ל-Cloud TPU API, שכבר לא נמצא בפיתוח פעיל. ממשק Cloud TPU API יקבל רק תיקוני באגים ועדכוני אבטחה. כדי להשתמש בתכונות העדכניות ביותר ולקבל תמיכה בגרסאות ה-TPU העדכניות ביותר, מומלץ להשתמש ב-Compute Engine או ב-Google Kubernetes Engine כדי לנהל משאבי TPU. מידע נוסף זמין במאמרים משאבי TPU ב-Compute Engine ומידע על TPUs ב-GKE. ‫Cloud TPU API לא נתמך ב-TPU7x ובגרסאות TPU מאוחרות יותר.

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

חשוב: אפשר להריץ עד 256 פודים של Docker במכונת TPU וירטואלית.

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

לפני שמריצים את הפקודות הבאות, יוצרים את משתני הסביבה הבאים:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

תיאורים של משתני סביבה

• PROJECT_ID: מזהה הפרויקט ב- Google Cloud . משתמשים בפרויקט קיים או יוצרים פרויקט חדש.
• ‫TPU_NAME: השם של ה-TPU.
• ‫ZONE: האזור שבו יוצרים את ה-TPU VM. מידע נוסף על אזורים נתמכים זמין במאמר אזורים ותחומים של TPU.
• ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של Cloud TPU שרוצים ליצור. מידע נוסף על סוגי המאיצים הנתמכים בכל גרסת TPU זמין במאמר בנושא גרסאות TPU.
• ‫RUNTIME_VERSION: גרסת התוכנה של Cloud TPU.

מכשיר TPU

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

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. התחברות ל-TPU VM באמצעות SSH

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

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

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

   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 VM

   gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE
   

TPU slice

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

1. יצירת TPU VM

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --zone=$ZONE \
   --worker=all \
   --command='sudo usermod -a -G docker $USER'
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="git clone --recursive https://github.com/pytorch/xla.git"
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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 VM באמצעות הפקודה הבאה:

gcloud compute tpus tpu-vm delete $TPU_NAME \
--zone=$ZONE

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

לפני שמריצים את הפקודות הבאות, יוצרים את משתני הסביבה הבאים:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

תיאורים של משתני סביבה

• PROJECT_ID: מזהה הפרויקט ב- Google Cloud . משתמשים בפרויקט קיים או יוצרים פרויקט חדש.
• ‫TPU_NAME: השם של ה-TPU.
• ‫ZONE: האזור שבו יוצרים את ה-TPU VM. מידע נוסף על אזורים נתמכים זמין במאמר אזורים ותחומים של TPU.
• ‫ACCELERATOR_TYPE: סוג המאיץ מציין את הגרסה והגודל של Cloud TPU שרוצים ליצור. מידע נוסף על סוגי המאיצים הנתמכים בכל גרסת TPU זמין במאמר בנושא גרסאות TPU.
• ‫RUNTIME_VERSION: גרסת התוכנה של Cloud TPU.

מכשיר TPU

1. יצירת TPU VM

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. התחברות ל-TPU VM באמצעות SSH

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

3. הפעלת שד Docker במכונת TPU וירטואלית

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

   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 VM

   gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE
   

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=$PROJECT_ID
   
   gcloud artifacts repositories list \
   --project=$PROJECT_ID
   
   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/$PROJECT_ID/your-repo/your-image-name:your-tag
   

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

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

6. יצירת TPU VM

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command='sudo usermod -a -G docker ${USER}'
   

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="docker pull europe-west4-docker.pkg.dev/$PROJECT_ID/your-repo/your-image-name:your-tag"
   

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

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

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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 $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="docker kill your-container-name"
   

2. מחיקת ה-TPU VM

   gcloud compute tpus tpu-vm delete $TPU_NAME \
   --zone=$ZONE
   

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

• מדריכים ל-Cloud TPU
• ניהול של יחידות TPU
• ארכיטקטורת Cloud TPU
• הרצת קוד JAX בפרוסות TPU