הרצת עומס עבודה של אצווה באמצעות Pathways

לצורך המאמר הזה, עומסי עבודה של אצווה מוגדרים כעומסי עבודה של JAX שמופעלים עד להשלמה ונפרסים באותו אשכול GKE כמו אשכול Pathways, במיוחד לצד רכיבי בקר Pathways (שרת proxy של IFRT ומנהל משאבים של Pathways). השלמת עומס העבודה של JAX מפסיקה את רכיבי האשכול של Pathways. במדריך הזה נשתמש בעומס עבודה של אימון JAX כדי להדגים את זה.

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

ודאו שיש לכם:

יצירת תמונת אימון באמצעות MaxText

‫MaxText הוא פרויקט קוד פתוח של מודל שפה גדול (LLM) שפותח על ידי Google. הוא כתוב ב-JAX ומתוכנן להיות בעל ביצועים גבוהים וניתן להרחבה, והוא פועל ביעילות במעבדי TPU ובמעבדי GPU של Google Cloud.

כדי ליצור קובץ אימג' של MaxText Docker באמצעות הגרסה היציבה העדכנית ביותר של JAX ממאגר ה-OSS GitHub, מריצים את הפקודה הבאה:

git clone https://github.com/AI-Hypercomputer/maxtext
cd maxtext/dependencies/scripts
gcloud config set project PROJECT_ID
bash ./docker_build_dependency_image.sh MODE=stable
gcloud auth configure-docker
bash ./docker_upload_runner.sh CLOUD_IMAGE_NAME=USER_runner # This script needs bash version >= 4.2 to execute.

הפקודה הזו מעבירה את תמונת Kubernetes של MaxText אל gcr.io/$PROJECT_ID/${USER}_runner. אפשר להשתמש בקובץ האימג' של Docker הזה כדי להריץ אימון ב-TPU באמצעות קצה העורף של Pathways.

הרצת עומס עבודה של אצווה באמצעות Cluster Toolkit

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

gcluster job submit \
    --pathways \
    --pathways-gcs-location="gs://BUCKET_NAME/pathways-artifacts" \
    --name=WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE \
    --num-slices=WORKLOAD_NODEPOOL_COUNT \
    --compute-type=COMPUTE_TYPE \
    --topology=TOPOLOGY \
    --image="gcr.io/PROJECT_ID/USER_runner" \
    --command="python3 -m MaxText.train /deps/src/MaxText/configs/base.yml base_output_directory=gs://BUCKET_NAME per_device_batch_size=1 enable_checkpointing=false remat_policy=full global_parameter_scale=1 steps=20 max_target_length=2048 use_iota_embed=true reuse_example_batch=1 dataset_type=synthetic attention=flash gcs_metrics=True enable_single_controller=True run_name=RUN_NAME-pathways-job"

מידע נוסף על אפשרויות לשליחת משימות זמין במדריך לשליחת משימות ב-Cluster Toolkit.

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

  • ‫WORKLOAD: שם ייחודי לזיהוי עומס העבודה. בגלל מגבלות התווים של תווית DNS, השם צריך לכלול 22 תווים או פחות.
  • ‫CLUSTER: השם של אשכול GKE
  • ‫WORKLOAD_NODEPOOL_COUNT: מספר מאגרי הצמתים של פרוסת ה-TPU
  • ‫COMPUTE_TYPE: סוג מכונת ה-TPU (לדוגמה, ct6e-standard-4t). מידע נוסף על סוגי ה-TPU הנתמכים לכל גרסת TPU זמין במאמר בנושא גרסאות TPU.
  • ‫TOPOLOGY: טופולוגיית המיקום של ה-TPU (לדוגמה, 2x4)
  • ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud
  • ‫ZONE: האזור שבו אתם מתכננים להריץ את עומס העבודה
  • ‫USER: מזהה המשתמש ב- Google Cloud
  • ‫BUCKET_NAME: הקטגוריה של Cloud Storage לפלט
  • ‫RUN_NAME: שם שהמשתמש מקצה כדי לזהות את ההרצה של תהליך העבודה

עוקבים אחרי ההתקדמות של עומס העבודה באמצעות הפקודה gcluster job logs:

gcluster job logs WORKLOAD \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=ZONE \
    --main-only=false

completed step: 1, seconds: 0.484, TFLOP/s/device: 87.349, Tokens/s/device: 2117.382, total_weights: 2945, loss: 10.888
completed step: 2, seconds: 0.407, TFLOP/s/device: 103.699, Tokens/s/device: 2513.735, total_weights: 3253, loss: 9.697
completed step: 3, seconds: 0.248, TFLOP/s/device: 170.300, Tokens/s/device: 4128.167, total_weights: 3154, loss: 9.641
completed step: 4, seconds: 0.216, TFLOP/s/device: 195.122, Tokens/s/device: 4729.880, total_weights: 3119, loss: 9.547
completed step: 5, seconds: 0.272, TFLOP/s/device: 155.298, Tokens/s/device: 3764.512, total_weights: 2837, loss: 10.179
completed step: 6, seconds: 0.472, TFLOP/s/device: 89.489, Tokens/s/device: 2169.266, total_weights: 3069, loss: 9.776

כדי לבטל את עומס העבודה לפני שהוא מסתיים, משתמשים בפקודה gcluster job cancel:

gcluster job cancel WORKLOAD --cluster=CLUSTER --project=PROJECT_ID --location=ZONE

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