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

במאמר הזה מוסבר איך להקצות משאבי TPU7x באמצעות Cluster Toolkit, ואיך לפרוס עומס עבודה של אימון באמצעות MaxText‏ (JAX) או TorchTitan ב-TPU‏ (PyTorch).

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

דוגמאות נוספות שמותאמות ל-TPU7x זמינות במאגרים הבאים ב-GitHub:

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

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

לאחר מכן, מבצעים את שלבי ההגדרה הבאים:

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

    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 לפלט האימון.

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

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

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

פורסים אשכול TPU7x ב-Google Kubernetes Engine (GKE) לפי ההוראות במאמר פריסת אשכול GKE TPU 7x. חשוב לוודא שהטופולוגיה של מאגר הצמתים באשכול תואמת לדרישות של עומס העבודה לאימון, כמו 4x4x8 בדוגמה של MaxText או 2x2x4 בדוגמה של TorchTitan.

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

בקטעים הבאים מוסבר איך ליצור קובץ אימג' של קונטיינר ולשלוח משימת אימון של MaxText באמצעות Cluster Toolkit.

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

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

# 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 מאמנת מודל שפה עם כמיליארד פרמטרים. המודל הזה משמש לבדיקת צ'יפים קטנים. הביצועים לא אופטימליים.

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"

שליחת משימת האימון של MaxText

שולחים את משימת האימון באמצעות 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).

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

‫TorchTitan משתמש ב-TorchTPU כדי להריץ אימון PyTorch מבוזר ב-Cloud TPU. מידע נוסף זמין בקובץ README של TorchTitan ב-TPU.

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