‫Google משתמשת בטכנולוגיית AI כדי לתרגם תוכן לשפה המועדפת עליך. בתרגומים כאלו עשויות להיות שגיאות.

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

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

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

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

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

אפשר להשתמש ב-Cluster Toolkit כדי ליצור אשכולות GKE שמוכנים לייצור ולהריץ עליהם עומסי עבודה של אימון.

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

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

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

מוודאים שיש לכם פרויקט ב- Google Cloud שהחיוב בו מופעל.
מתקינים ומפעילים את Google Cloud CLI.
התקנת Cluster Toolkit

מגדירים את משתני הסביבה הבאים:
```
export PROJECT_ID=YOUR_PROJECT_ID
export ZONE=YOUR_ZONE
export CLUSTER_NAME=YOUR_CLUSTER_NAME
export BASE_OUTPUT_DIR="gs://YOUR_BUCKET_NAME"
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫YOUR_PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
- ‫YOUR_ZONE: האזור שבו רוצים ליצור את האשכול.
- ‫YOUR_CLUSTER_NAME: השם של האשכול החדש.
- ‫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

פריסת אשכול TPU7x ‏ (Ironwood)

פורסים אשכול GKE TPU7x על ידי ביצוע ההוראות במאמר פריסת אשכול GKE TPU 7x. מוודאים שהטופולוגיה של מאגר הצמתים באשכול תואמת לדרישות של עומס העבודה של האימון (למשל 4x4x8 עבור עומס העבודה לדוגמה של MaxText שמוצג בהמשך).

פיתוח והעלאה של קובץ האימג' של Docker ל-MaxText

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

# 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

export CLOUD_IMAGE_NAME="${USER}-maxtext-runner"
bash docker_upload_runner.sh CLOUD_IMAGE_NAME=${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_training \
        steps=30"

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

שולחים את משימת האימון באמצעות gcluster job submit. פרטים על אפשרויות שליחת משימות זמינים במדריך המשימות של Cluster Toolkit.

gcluster job submit \
    --name="maxtext-1b-$(date +%H%M)" \
    --cluster=${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${ZONE} \
    --compute-type=tpu7x-standard-4t \
    --topology=4x4x8 \
    --image="gcr.io/${PROJECT_ID}/${CLOUD_IMAGE_NAME}" \
    --command="${MAXTEXT_COMMAND}"

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