אימון מודל באמצעות TPU7x ‏ (Ironwood)

במאמר הזה מוסבר איך להקצות משאבי TPU7x ומוצגת דוגמה לפריסת עומס עבודה של אימון באמצעות MaxText ו-XPK.

‫TPU7x היא הגרסה הראשונה במשפחת Ironwood, Google Cloud's TPU מהדור השביעי. דור Ironwood מיועד לאימון AI ולהסקת מסקנות בקנה מידה גדול. מידע נוסף זמין במאמר בנושא TPU7x.

דוגמאות נוספות שעברו אופטימיזציה ל-TPU7x זמינות ב-Training Recipes for Ironwood TPU ב-GitHub.

הקצאת מעבדי TPU

אפשר להקצות ולנהל TPU7x באמצעות השיטות הבאות:

‫GKE: אתם יכולים להשתמש ב-GKE כדי להקצות ולנהל יחידות TPU כמאגר של מאיצים לעומסי עבודה של למידת מכונה בקונטיינרים. אתם יכולים להשתמש ב-Google Cloud CLI כדי ליצור ידנית את מופע אשכול GKE, כדי לבצע התאמה אישית מדויקת או להרחיב סביבות GKE קיימות בסביבת הייצור. מידע נוסף זמין במאמר בנושא TPU ב-GKE.
‫GKE ו-XPK: XPK הוא כלי שורת פקודה שמפשט את יצירת האשכולות והרצת עומסי העבודה ב-GKE. הוא מיועד למומחים בתחום למידת המכונה (ML) כדי להקצות TPU ולהריץ משימות אימון ללא צורך במומחיות מעמיקה ב-Kubernetes. אפשר להשתמש ב-XPK כדי ליצור במהירות אשכולות GKE ולהפעיל עומסי עבודה (workloads) לצורך הוכחת היתכנות ובדיקות. למידע נוסף, אפשר לעיין במאגר XPK ב-GitHub.
‫GKE ו-TPU Cluster Director: TPU Cluster Director זמין דרך הזמנה במצב 'כל הקיבולת', שנותנת לכם גישה מלאה לכל הקיבולת שהזמנתם (בלי הגבלות) ושקיפות מלאה לגבי טופולוגיית החומרה של TPU, סטטוס הניצול וסטטוס התקינות. מידע נוסף זמין במאמר סקירה כללית של מצב 'כל הקיבולת'.

פריסת עומס עבודה של אימון באמצעות MaxText ו-XPK

אפשר להשתמש ב-Accelerated Processing Kit ‏ (XPK) כדי ליצור אשכולות GKE לצורך הוכחת היתכנות ובדיקות. XPK הוא כלי שורת פקודה שנועד לפשט את ההקצאה, הניהול וההפעלה של עומסי עבודה של למידת מכונה.

בקטעים הבאים מוסבר איך לפרוס עומס עבודה של אימון באמצעות MaxText ו-XPK.

לפני שמתחילים

לפני שמתחילים, צריך לבצע את השלבים הבאים:

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

התקנה של XPK ויחסי תלות

מתקינים XPK. פועלים לפי ההוראות במאגר XPK ב-GitHub.
מתקינים את Docker לפי ההוראות שקיבלתם מהאדמין או לפי הוראות ההתקנה הרשמיות. אחרי ההתקנה, מריצים את הפקודות הבאות כדי להגדיר את Docker ולבדוק את ההתקנה:
```
gcloud auth configure-docker
sudo usermod -aG docker $USER # relaunch the terminal and activate venv after running this command
docker run hello-world # Test Docker
```
מגדירים את משתני הסביבה הבאים:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export ACCELERATOR_TYPE=YOUR_ACCELERATOR_TYPE
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫YOUR_PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
- ‫YOUR_ZONE: האזור שבו רוצים ליצור את האשכול. בתצוגה המקדימה, יש תמיכה רק ב-us-central1-c.
- ‫YOUR_CLUSTER_NAME: השם של האשכול החדש.
- ‫YOUR_ACCELERATOR_TYPE: הגרסה והטופולוגיה של ה-TPU. לדוגמה, tpu7x-4x4x8. רשימת הטופולוגיות הנתמכות מופיעה כאן.
- ‫YOUR_BUCKET_NAME: השם של קטגוריה של Cloud Storage, שתהיה ספריית הפלט לאימון המודל.

אם אין לכם קטגוריה קיימת ב-Cloud Storage, אתם יכולים ליצור אחת באמצעות הפקודה הבאה:

gcloud storage buckets create ${BASE_OUTPUT_DIR} \
    --project=${PROJECT_ID} \
    --location=US \
    --default-storage-class=STANDARD \
    --uniform-bucket-level-access

יצירת אשכול עם כרטיס רשת יחיד ופרוסת שרת יחידה

בוחרים אחת מהאפשרויות הבאות כדי ליצור את האשכול. כדי להשיג ביצועים אופטימליים, מומלץ להשתמש ברשת מותאמת אישית עם MTU של 8,896.

רשת בהתאמה אישית

כדי ליצור רשת בהתאמה אישית עם MTU של 8,896 ולהשתמש בה באשכול, פועלים לפי השלבים הבאים:

מגדירים משתני סביבה לשמות של הרשת ושל חומת האש:
```
export NETWORK_NAME=NETWORK_NAME
export NETWORK_FW_NAME=FIREWALL_NAME
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫NETWORK_NAME: שם לרשת.
- ‫FIREWALL_NAME: שם של כלל חומת אש בין רשתות.

כדי ליצור רשת בהתאמה אישית עם MTU של 8,896:

gcloud compute networks create ${NETWORK_NAME} \
    --mtu=8896 \
    --project=${PROJECT_ID} \
    --subnet-mode=auto \
    --bgp-routing-mode=regional

יוצרים כלל של חומת אש שמאפשר תעבורת נתונים מסוג TCP,‏ ICMP ו-UDP ברשת:

gcloud compute firewall-rules create ${NETWORK_FW_NAME} \
    --network=${NETWORK_NAME} \
    --allow tcp,icmp,udp \
    --project=${PROJECT_ID}

מגדירים משתנה סביבה לארגומנטים של אשכול XPK כדי להשתמש ברשת שיצרתם:
```
export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"
```

יוצרים את אשכול ה-XPK. הפקודה הבאה מקצה קיבולת לפי דרישה:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand \
    --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"

כדי להשתמש בקיבולת שמורה, מחליפים את --on-demand ב---reservation=RESERVATION_NAME. כדי להשתמש במכונות וירטואליות מסוג TPU Spot, מחליפים את --on-demand ב---spot.

רשת ברירת המחדל

אם אתם לא צריכים רשת עם MTU גבוה, אתם יכולים ליצור אשכול שמשתמש ברשת VPC שמוגדרת כברירת מחדל. הפקודה הבאה מקצה קיבולת לפי דרישה:

xpk cluster create --cluster=${CLUSTER_NAME} \
    --cluster-cpu-machine-type=n1-standard-8 \
    --num-slices=${NUM_SLICES} \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone=${ZONE} \
    --project=${PROJECT_ID} \
    --on-demand

איך יוצרים או מעלים את קובץ אימג' של Docker של MaxText

אפשר ליצור קובץ אימג' של Docker באופן מקומי באמצעות סקריפטים שסופקו על ידי MaxText, או להשתמש בקובץ אימג' מוכן מראש.

פיתוח באופן מקומי

הפקודות הבאות מעתיקות את הספרייה המקומית אל המאגר:

# Make sure you're running on a virtual environment with python3.12. If nothing is printed, you have the correct version.
[[ "$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")' 2>/dev/null)" == "3.12" ]] || { >&2 echo "Error: Python version must be 3.12."; false; }

# Clone MaxText
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext
git checkout maxtext-tutorial-v1.0.0

# Build the Docker image
bash docker_build_dependency_image.sh MODE=stable JAX_VERSION=0.8.2

אחרי הרצה מוצלחת של הפקודות, אמורה להופיע תמונה בשם maxtext_base_image שנוצרה באופן מקומי. אפשר להשתמש בתמונה המקומית ישירות בפקודת עומס העבודה xpk.

העלאת תמונה (אופציונלי)

אחרי שיוצרים את קובץ האימג' של Docker באופן מקומי באמצעות ההוראות שבקטע הקודם, אפשר להעלות את קובץ האימג' של MaxText Docker למאגר באמצעות הפקודה הבאה:

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${CLOUD_IMAGE_NAME}

אחרי שהפקודה הזו תופעל בהצלחה, התמונה MaxText תופיע ב-gcr.io עם השם gcr.io/PROJECT_ID/CLOUD_IMAGE_NAME.

הגדרת פקודת האימון של MaxText

מכינים את הפקודה להרצת סקריפט ההדרכה בקונטיינר Docker.

מודל MaxText 1B הוא הגדרה במסגרת MaxText שנועדה לאימון מודל שפה עם כמיליארד פרמטרים. אפשר להשתמש במודל הזה כדי להתנסות בשבבים קטנים. הביצועים לא אופטימליים.

export MAXTEXT_COMMAND="JAX_PLATFORMS=tpu,cpu \
    ENABLE_PJRT_COMPATIBILITY=true \
    python3 src/MaxText/train.py src/MaxText/configs/base.yml \
        base_output_directory=${BASE_OUTPUT_DIR} \
        dataset_type=synthetic \
        per_device_batch_size=2 \
        enable_checkpointing=false \
        gcs_metrics=true \
        run_name=maxtext_xpk \
        steps=30"

פריסת עומס העבודה של האימון

מריצים את הפקודה xpk workload create כדי לפרוס את משימת האימון. צריך לציין את הדגל --base-docker-image כדי להשתמש בתמונת הבסיס של MaxText, או לציין את הדגל --docker-image ואת התמונה שרוצים להשתמש בה. אפשר לכלול את הדגל --enable-debug-logs כדי להפעיל רישום של ניפוי באגים.

xpk workload create \
    --cluster ${CLUSTER_NAME} \
    --base-docker-image maxtext_base_image \
    --workload maxtext-1b-$(date +%H%M) \
    --tpu-type=${ACCELERATOR_TYPE} \
    --zone ${ZONE} \
    --project ${PROJECT_ID} \
    --command "${MAXTEXT_COMMAND}"
    # [--enable-debug-logs]

שמות של עומסי עבודה חייבים להיות ייחודיים בתוך האשכול. בדוגמה הזו, הערך $(date +%H%M) מצורף לשם עומס העבודה כדי להבטיח שהוא יהיה ייחודי.