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

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

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

במאמר PyTorch ב-Cloud TPU מתוארות ספריות PyTorch ל-TPU, ויש בו קישורים למסמכי תיעוד לקבלת חבילות Python וקובצי אימג' של Docker. אתם יכולים להשתמש ישירות ב-TorchTPU (ראו תחילת העבודה ב-GKE לקובץ האימג' של הקונטיינר בגרסה שזמינה לכולם) או להשתמש בספריית העזר TorchTitan כדי ליצור ולהריץ עומסי עבודה של אימון בקונטיינרים.

אימון מודל 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 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-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

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

  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 workers

    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 workers

    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. סגירת הקונטיינר בכל ה-workers

    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

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