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

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

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

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

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

אפשר להשתמש בערכת עיבוד מואץ (XPK) כדי ליצור אשכולות GKE לצורך הוכחת היתכנות ובדיקות.

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

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

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

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

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

1. מתקינים XPK. פועלים לפי ההוראות במאגר XPK ב-GitHub.

2. מתקינים את 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
   

3. מגדירים את משתני הסביבה הבאים:

   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: האזור שבו רוצים ליצור את האשכול.
   • ‫YOUR_CLUSTER_NAME: השם של האשכול החדש.
   • ‫YOUR_ACCELERATOR_TYPE: הגרסה והטופולוגיה של ה-TPU. לדוגמה, tpu7x-4x4x8. רשימת הטופולוגיות הנתמכות מופיעה כאן.
   • ‫YOUR_BUCKET_NAME: השם של קטגוריה של Cloud Storage, שתשמש כספריית הפלט לאימון המודל.

4. אם אין לכם קטגוריה קיימת של 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 ולהשתמש בה עבור האשכול, פועלים לפי השלבים הבאים:

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

   export NETWORK_NAME=NETWORK_NAME
   export NETWORK_FW_NAME=FIREWALL_NAME

   מחליפים את מה שכתוב בשדות הבאים:

   • ‫NETWORK_NAME: שם הרשת.
   • ‫FIREWALL_NAME: שם של כלל חומת האש ברשת.

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

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

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

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

4. מגדירים משתנה סביבה לארגומנטים של אשכול XPK כדי להשתמש ברשת שיצרתם:

   export CLUSTER_ARGUMENTS="--network=${NETWORK_NAME} --subnetwork=${NETWORK_NAME}"

5. יוצרים את אשכול ה-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

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

הערה: הגדרת הדגל --cluster-cpu-machine-type לערך n1-standard-8 (או לערך גדול יותר) מבטיחה שמאגר הצמתים שמוגדר כברירת מחדל יכלול מספיק ליבות CPU בשביל פודים של המערכת, למשל webhook של JobSet, וכך מונעת שגיאות. כברירת מחדל, XPK משתמש ב-e2-standard-16. אזורים מסוימים תומכים רק בסוגים ספציפיים של מעבדים, ולכן יכול להיות שתצטרכו לעבור בין הסוגים n1, n2 ו-e2. אחרת, יכול להיות שתיתקלו בשגיאות שקשורות למכסת השימוש.

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

הערה: MaxText תומך רק ב-Python 3.12. יוצרים את הסביבה הווירטואלית עם גרסה 3.12 כדי להתקין את התלות הנכונה.

אפשר ליצור קובץ אימג' של 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) מצורף לשם עומס העבודה כדי להבטיח ייחודיות.

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

• שימוש בפלטפורמת האבחון של Google Cloud ML כדי לבצע אופטימיזציה של עומסי העבודה ולאבחן אותם
• הפעלת עומס עבודה של אימון באמצעות מתכון שעבר אופטימיזציה ל-TPU7x
• הפעלת מיקרו-בנצ'מרק TPU7x