אימון מודל באמצעות TPU v5e

עם שטח קטן יותר של 256 שבבים לכל Pod,‏ TPU v5e מותאם להיות מוצר בעל ערך גבוה לאימון, לכוונון ולהצגה של מודלים מסוג Transformer, יצירת תמונות לפי טקסט ורשתות נוירונים מלאכותיות (CNN). מידע נוסף על השימוש ב-Cloud TPU v5e להצגה זמין במאמר הסקה באמצעות v5e.

מידע נוסף על חומרת TPU והגדרות אישיות של Cloud TPU v5e זמין במאמר TPU v5e.

קדימה, מתחילים

בקטעים הבאים מוסבר איך להתחיל להשתמש ב-TPU v5e.

בקשת מכסה

כדי להשתמש ב-TPU v5e לאימון, צריך מכסה. יש סוגים שונים של מכסות ל-TPU על פי דרישה, ל-TPU שמור ול-TPU במודל Spot. אם אתם משתמשים ב-TPU v5e להסקה, אתם צריכים מכסות נפרדות. מידע נוסף על מכסות כדי לבקש מכסה ל-TPU v5e, פנו למחלקת המכירות של Cloud.

יצירת חשבון ופרויקט ב- Google Cloud

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

יצירת Cloud TPU

השיטה המומלצת היא להקצות את Cloud TPU v5e כמשאבים בתור באמצעות הפקודה queued-resource create. מידע נוסף זמין במאמר ניהול משאבים בתור.

אפשר גם להשתמש ב-Create Node API‏ (gcloud compute tpus tpu-vm create) כדי להקצות את Cloud TPU v5e. מידע נוסף זמין במאמר ניהול משאבי TPU.

מידע נוסף על ההגדרות הזמינות של v5e לאימון מופיע במאמר סוגי Cloud TPU v5e לאימון.

הגדרת המסגרת

בקטע הזה מתואר תהליך ההגדרה הכללי לאימון מודל בהתאמה אישית באמצעות JAX עם TPU v5e. אפשרויות האימון של PyTorch מפורטות במאמר אימון מודל באמצעות PyTorch ב-Cloud TPU.

הוראות להגדרת הסקה מופיעות במאמר מבוא להסקה ב-v5e.

מגדירים כמה משתני סביבה:

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

הגדרה ל-JAX

אם יש לכם צורות של פרוסות עם יותר מ-8 צ'יפים, יהיו לכם כמה מכונות וירטואליות בפרוסה אחת. במקרה כזה, צריך להשתמש בדגל --worker=all כדי להריץ את ההתקנה בכל ה-TPU VM בשלב אחד, בלי להשתמש ב-SSH כדי להתחבר לכל אחת מהן בנפרד:

gcloud compute tpus tpu-vm ssh $TPU_NAME  \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

תיאורי דגלי הפקודות

  • ‫TPU_NAME: מזהה הטקסט שהמשתמש הקצה ל-TPU, שנוצר כשמוקצית הבקשה למשאב בתור.
  • ‫PROJECT_ID:שם הפרויקט ב- Google Cloud . משתמשים בפרויקט קיים או יוצרים פרויקט חדש בקטע הגדרת פרויקט Google Cloud .
  • ‫ZONE: במסמך אזורים ותחומים של TPU מפורטים התחומים הנתמכים.
  • ‫worker: ה-TPU VM שיש לו גישה ל-TPUs הבסיסיים.

כדי לבדוק את מספר המכשירים, אפשר להריץ את הפקודה הבאה (הפלט שמוצג כאן נוצר באמצעות פרוסת v5litepod-16). הקוד הזה בודק שהכול הותקן כמו שצריך, על ידי בדיקה ש-JAX רואה את Cloud TPU TensorCores ויכול להריץ פעולות בסיסיות:

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

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

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

הערך jax.device_count() מייצג את המספר הכולל של הצ'יפים בפרוסה הנתונה. ‫jax.local_device_count() מציין את מספר הצ'יפים שאליהם יש גישה למכונה וירטואלית אחת בפרוסה הזו.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

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

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

כדי להתחיל לאמן את v5e באמצעות JAX, נסו את המדריכים של JAX במסמך הזה.

הסרת המשאבים

בסיום הסשן, מחקו את ה-TPU ואת המשאב בתור. כדי למחוק משאב בתור, מוחקים את הפרוסה ואז את המשאב בתור בשני שלבים:

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

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=${ZONE} \
   --quiet

אפשר להשתמש בשני השלבים האלה גם כדי להסיר בקשות למשאבים בתור שנמצאות במצב FAILED.

דוגמאות ל-JAX/FLAX

בקטעים הבאים מתוארות דוגמאות לאימון מודלים של JAX ו-FLAX ב-TPU v5e.

אימון ImageNet ב-v5e

במדריך הזה מוסבר איך לאמן את ImageNet ב-v5e באמצעות נתוני קלט מזויפים. אם אתם רוצים להשתמש בנתונים אמיתיים, עיינו בקובץ README ב-GitHub.

הגדרה

  1. יוצרים משתני סביבה:

    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
    export SERVICE_ACCOUNT=your-service-account
    export QUEUED_RESOURCE_ID=your-queued-resource-id

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

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

      לדוגמה: tpu-service-account@PROJECT_ID.iam.

    • ‫QUEUED_RESOURCE_ID: מזהה הטקסט שהמשתמש הקצה לבקשת המשאבים בתור.

  2. יצירת משאב TPU:

    gcloud 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
    

    תוכלו להתחבר ב-SSH ל-TPU VM ברגע שהמשאב בתור יהיה במצב ACTIVE:

    gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
       --project=$PROJECT_ID \
       --zone=$ZONE
    

    כשה-QueuedResource במצב ACTIVE, הפלט אמור להיראות כך:

     state: ACTIVE
    
  3. מתקינים את הגרסה העדכנית של JAX ו-jaxlib:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --worker=all \
       --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. משכפלים את המודל של ImageNet ומתקינים את הדרישות המתאימות:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --worker=all \
       --command="git clone https://github.com/coolkp/flax.git && cd flax && git checkout pmap-orbax-conversion && git pull"
    
  5. כדי ליצור נתונים מזויפים, המודל צריך מידע על המאפיינים של מערך הנתונים. אפשר לאסוף את המידע הזה מהמטא-נתונים של קבוצת הנתונים ImageNet:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --worker=all \
       --command="cd flax/examples/imagenet && pip install -r requirements-cloud-tpu.txt"
    

אימון המודל

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

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command="cd flax/examples/imagenet && bash ../../tests/download_dataset_metadata.sh && JAX_PLATFORMS=tpu python imagenet_fake_data_benchmark.py"

מחיקת ה-TPU והמשאב בתור

בסיום הסשן, מחקו את ה-TPU ואת המשאב בתור.

gcloud compute tpus tpu-vm delete $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet
gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

מודלים של Hugging Face FLAX

מודלים של Hugging Face שמוטמעים ב-FLAX פועלים באופן מיידי ב-Cloud TPU v5e. בקטע הזה מוסבר איך להריץ מודלים פופולריים.

אימון ViT ב-Imagenette

במדריך הזה נסביר איך לאמן את המודל Vision Transformer (ViT) מ-HuggingFace באמצעות מערך הנתונים Imagenette של Fast AI ב-Cloud TPU v5e.

מודל ViT היה הראשון שאימן בהצלחה מקודד Transformer ב-ImageNet עם תוצאות מצוינות בהשוואה לרשתות קונבולוציה. מידע נוסף זמין בסקירה הכללית על ViT.

הגדרה

  1. יוצרים משתני סביבה:

    export PROJECT_ID=your-project-id
    export TPU_NAME=your-tpu-name
    export ZONE=us-west4-a
    export ACCELERATOR_TYPE=v5litepod-16
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=your-service-account
    export QUEUED_RESOURCE_ID=your-queued-resource-id

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

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

      לדוגמה: tpu-service-account@PROJECT_ID.iam.

    • ‫QUEUED_RESOURCE_ID: מזהה הטקסט שהמשתמש הקצה לבקשת המשאבים בתור.

  2. יצירת משאב TPU:

    gcloud 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
    

    תוכלו להתחבר ב-SSH ל-TPU VM ברגע שהמשאב בתור יהיה במצב ACTIVE:

    gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
       --project=$PROJECT_ID \
       --zone=$ZONE
    

    כשהמשאב בתור נמצא במצב ACTIVE, הפלט אמור להיראות כך:

     state: ACTIVE
    
  3. מתקינים את JAX ואת הספרייה שלו:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --worker=all \
       --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. מורידים את המאגר של Hugging Face ומתקינים את הדרישות:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --worker=all \
       --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.19 && sed -i "s/torchvision==0.12.0+cpu/torchvision==0.22.1/" examples/flax/vision/requirements.txt && pip install -r examples/flax/vision/requirements.txt && pip install tf-keras'
    
  5. מורידים את מערך הנתונים Imagenette:

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --worker=all \
       --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
    

אימון המודל

מאמנים את המודל עם מאגר נתונים ממופה מראש בנפח 4GB.

gcloud compute tpus tpu-vm ssh $TPU_NAME \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

מחיקת ה-TPU והמשאב בתור

בסיום הסשן, מוחקים את ה-TPU ואת המשאב בתור.

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

gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
   --project=$PROJECT_ID \
   --zone=$ZONE \
   --quiet

תוצאות של הערכת ביצועים של ViT

סקריפט האימון הופעל ב-v5litepod-4, ב-v5litepod-16 וב-v5litepod-64. בטבלה הבאה מוצג קצב העיבוד עם סוגים שונים של מאיצים.

סוג המאיץ v5litepod-4 v5litepod-16 v5litepod-64
תקופת זמן לקביעת נושאים 3 3 3
גודל אצווה גלובלי 32 128 512
תפוקה (דוגמאות לשנייה) 263.40 429.34 470.71

אימון מודל דיפוזיה על פוקימונים

במדריך הזה מוסבר איך לאמן את המודל Stable Diffusion מ-HuggingFace באמצעות מערך הנתונים Pokémon ב-Cloud TPU v5e.

המודל Stable Diffusion הוא מודל סמוי של יצירת תמונות לפי טקסט, שיוצר תמונות פוטו-ריאליסטיות מכל קלט טקסט. למידע נוסף, תוכלו לקרוא את המאמרים הבאים:

הגדרה

  1. מגדירים משתנה סביבה לשם של קטגוריית האחסון:

    export GCS_BUCKET_NAME=your_bucket_name
    
  2. מגדירים קטגוריית אחסון לפלט מהמודל:

    Google Cloud CLI storage buckets create gs://$GCS_BUCKET_NAME \
         --project=$PROJECT_ID \
         --location=$ZONE
    
  3. יוצרים משתני סביבה:

    export PROJECT_ID=your-project-id
    export TPU_NAME=your-tpu-name
    export ZONE=us-west1-c
    export ACCELERATOR_TYPE=v5litepod-16
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=your-service-account
    export QUEUED_RESOURCE_ID=your-queued-resource-id

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

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

      לדוגמה: tpu-service-account@PROJECT_ID.iam.

    • ‫QUEUED_RESOURCE_ID: מזהה הטקסט שהמשתמש הקצה לבקשת המשאבים בתור.

  4. יצירת משאב TPU:

    gcloud 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
    

    תוכלו להתחבר ב-SSH ל-TPU VM ברגע שהמשאב בתור יהיה במצב ACTIVE:

    gcloud compute tpus queued-resources describe $QUEUED_RESOURCE_ID  \
       --project=$PROJECT_ID \
       --zone=$ZONE
    

    כשהמשאב בתור נמצא במצב ACTIVE, הפלט אמור להיראות כך:

     state: ACTIVE
    
  5. מתקינים את JAX ואת הספרייה שלו.

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --worker=all \
       --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  6. מורידים את המאגר של HuggingFace ומתקינים את הדרישות.

    gcloud compute tpus tpu-vm ssh $TPU_NAME \
         --project=$PROJECT_ID \
         --zone=$ZONE \
         --worker=all \
         --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install -U -r examples/text_to_image/requirements_flax.txt && pip install tensorflow==2.17.1 clu && pip install tensorboard==2.17.1'
    

אימון המודל

מאמנים את המודל עם מאגר נתונים ממופה מראש בנפח 4GB.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --project=$PROJECT_ID --worker=all --command="
    git clone https://github.com/google/maxdiffusion
    cd maxdiffusion
    pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    pip3 install -r requirements.txt
    pip3 install .
    pip3 install gcsfs
    export LIBTPU_INIT_ARGS=''
    python -m src.maxdiffusion.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run \
    jax_cache_dir=gs://${GCS_BUCKET_NAME} activations_dtype=bfloat16 weights_dtype=bfloat16 \
    per_device_batch_size=1 precision=DEFAULT dataset_save_location=gs://${GCS_BUCKET_NAME} \
    output_dir=gs://${GCS_BUCKET_NAME}/ attention=flash"

הסרת המשאבים

בסיום הסשן, מוחקים את ה-TPU, את המשאב בתור ואת הקטגוריה של Cloud Storage.

  1. מוחקים את ה-TPU:

    gcloud compute tpus tpu-vm delete $TPU_NAME \
        --project=$PROJECT_ID \
        --zone=$ZONE \
        --quiet
    
  2. מוחקים את המשאב שבתור:

    gcloud compute tpus queued-resources delete $QUEUED_RESOURCE_ID \
        --project=$PROJECT_ID \
        --zone=$ZONE \
        --quiet
    
  3. מוחקים את הקטגוריה של Cloud Storage:

    gcloud storage rm -r gs://$GCS_BUCKET_NAME
    

תוצאות של הערכת ביצועים לדיפוזיה

סקריפט האימון פעל ב-v5litepod-4, ב-v5litepod-16 וב-v5litepod-64. בטבלה הבאה מוצגים קצבי העברת הנתונים.

סוג המאיץ v5litepod-4 v5litepod-16 v5litepod-64
שלב האימון 1500 1500 1500
גודל אצווה גלובלי 32 64 128
תפוקה (דוגמאות לשנייה) 36.53 43.71 49.36

אימון מודל באמצעות PyTorch ב-Cloud TPU

אתם יכולים להשתמש ב-TorchTPU ישירות כדי לאמן את המודל או להשתמש שוב בספריית העזר TorchTitan. כדי להתחיל לעבוד עם TorchTPU ב-TPU VM, קראו את המאמר תחילת העבודה עם TorchTPU ב-TPU VM. מידע נוסף על סביבת PyTorch ב-Cloud TPU זמין במאמר PyTorch ב-Cloud TPU.