אימון מודל באמצעות 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.
הגדרה
יוצרים משתני סביבה:
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: מזהה הטקסט שהמשתמש הקצה לבקשת המשאבים בתור.
-
-
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 מתקינים את הגרסה העדכנית של 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'משכפלים את המודל של 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"כדי ליצור נתונים מזויפים, המודל צריך מידע על המאפיינים של מערך הנתונים. אפשר לאסוף את המידע הזה מהמטא-נתונים של קבוצת הנתונים 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.
הגדרה
יוצרים משתני סביבה:
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: מזהה הטקסט שהמשתמש הקצה לבקשת המשאבים בתור.
-
-
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 מתקינים את 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'מורידים את המאגר של 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'מורידים את מערך הנתונים 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 הוא מודל סמוי של יצירת תמונות לפי טקסט, שיוצר תמונות פוטו-ריאליסטיות מכל קלט טקסט. למידע נוסף, תוכלו לקרוא את המאמרים הבאים:
הגדרה
מגדירים משתנה סביבה לשם של קטגוריית האחסון:
export GCS_BUCKET_NAME=your_bucket_nameמגדירים קטגוריית אחסון לפלט מהמודל:
Google Cloud CLI storage buckets create gs://$GCS_BUCKET_NAME \ --project=$PROJECT_ID \ --location=$ZONEיוצרים משתני סביבה:
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: מזהה הטקסט שהמשתמש הקצה לבקשת המשאבים בתור.
-
-
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 מתקינים את 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'מורידים את המאגר של 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.
מוחקים את ה-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מוחקים את הקטגוריה של 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.