הרצת SFT עבור Gemma 4 31B ב-TPU v6e מרובה-מארחים

במדריך הזה מוסבר איך להריץ כוונון מפוקח (SFT) באשכול של Tensor Processing Unit (TPU) v6e עם כמה מארחים, באמצעות MaxText ו-Cluster Toolkit. אתם משתמשים ב-Cluster Toolkit כדי להריץ עומס עבודה של אימון בכמה מארחים, ולייצא את התוצאות בחזרה לפורמט של Hugging Face לצורך הצגה.

סיכום עומס העבודה

כדי להריץ את עומס העבודה של SFT ב-Cloud TPU וב-GKE, צריך להשתמש ברכיבים הבאים:

  • ‫מודל: Gemma 4 31B (google/gemma-4-31b-it), מודל מתקדם עם משקלים פתוחים. במדריך הזה, משקלי המודל יורדו מ-Hugging Face, יומרו לפורמט שתואם ל-MaxText וייפרסו באשכול GKE עם תשתית Cloud TPU v6e (v6e-32).
  • שיטת האימון: SFT באמצעות MaxText. עומס העבודה מריץ צינור עיבוד נתונים של כוונון בקנה מידה גדול ב-Cloud TPU v6e באמצעות JAX. הוא מבצע אופטימיזציה של הפרמטרים של המודל במשך 1,000 שלבים באמצעות מאמן ה-SFT המובנה של MaxText (maxtext.trainers.post_train.sft.train_sft) ומערכי נתונים של כוונון לפי הוראות.
  • מערך נתונים ותבנית: UltraChat 200k‏ (HuggingFaceH4/ultrachat_200k), מערך נתונים גדול ואיכותי של שיחות עם כמה תורות. עומס העבודה משתמש במערך הנתונים הזה כדי לכוונן את המודל ולשפר את היכולות שלו לנהל שיחות ולמלא הוראות.
  • תשתית ומסגרת: GKE, מנוהל באמצעות Cluster Toolkit, ב-Cloud TPU v6e (v6e-32, עם טופולוגיה של 4x8). ה-framework הוא MaxText (tpu_post_training:0.2.4) עם קצה עורפי של JAX.

מטרות

  • מתקינים את Cluster Toolkit ואת יחסי התלות שלו.
  • פריסת אשכול של Cluster Toolkit.
  • המרת מודל של Hugging Face לפורמט MaxText.
  • מריצים עומס עבודה של אימון SFT ב-TPU.
  • ממירים את המודל המכוונן בחזרה לפורמט של Hugging Face כדי להשתמש בו.

עלויות

במסמך הזה משתמשים ברכיבים הבאים של Google Cloud, והשימוש בהם כרוך בתשלום:

בעזרת מחשבון העלויות אתם יכולים להעריך את העלויות בהתאם לשימוש הצפוי.

משתמשים חדשים של Google Cloud ? יכול להיות שאתם זכאים לתקופת ניסיון בחינם.

כשמסיימים את המשימות שמתוארות במסמך הזה אפשר למחוק את המשאבים שיצרתם כדי להימנע מחיובים נוספים. מידע נוסף זמין בקטע הסרת המשאבים.

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

  • כדי להשתמש במדריך הזה, צריך טוקן גישה ל-Hugging Face. אתם יכולים להירשם לחשבון בחינם ב-Hugging Face. אחרי שיוצרים חשבון, יוצרים אסימון גישה:

    1. בדף Welcome to Hugging Face, לוחצים על הדמות של החשבון ובוחרים באפשרות טוקן גישה.
    2. בדף טוקני גישה, לוחצים על יצירת טוקן חדש.
    3. בוחרים את סוג הטוקן Read ומזינים שם לטוקן.
    4. טוקן הגישה שלכם יוצג. שומרים את הטוקן במקום בטוח.

כדי לקבל את ההרשאות שדרושות להשלמת המדריך, בקשו מהאדמין להקצות לכם את התפקידים הבאים ב-IAM בפרויקט:

להסבר על מתן תפקידים, ראו איך מנהלים את הגישה ברמת הפרויקט, התיקייה והארגון.

יכול להיות שאפשר לקבל את ההרשאות הנדרשות גם באמצעות תפקידים בהתאמה אישית או תפקידים מוגדרים מראש.

הגדרת משתני הסביבה

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

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

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

  • ‫YOUR_PROJECT_ID: מזהה הפרויקט ב- Google Cloud.
  • ‫YOUR_REGION: האזור שבו נמצאת ההזמנה של ה-TPU ושבו פורסים את האשכול.
  • ‫YOUR_ZONE: האזור שבו נמצאת ההזמנה של ה-TPU ושבו פורסים את האשכול.
  • ‫YOUR_CLUSTER_NAME: השם של אשכול GKE. השם צריך להתחיל באות קטנה, לכלול רק אותיות קטנות, מספרים ומקפים, ולהיות באורך של עד 20 תווים, כדי שמזהי חשבונות השירות שמבוססים עליו לא יחרגו ממגבלת 30 התווים.
  • ‫YOUR_BUCKET_NAME: שם ייחודי גלובלי לקטגוריה של Cloud Storage.
  • ‫YOUR_RESERVATION_NAME: השם של ההזמנה של TPU v6e.
  • ‫YOUR_HF_TOKEN: טוקן הגישה שלכם ל-Hugging Face.

כברירת מחדל, הסקריפט מגדיר את MODEL_NAME ל-gemma4-31b (שם המודל של MaxText ל-google/gemma-4-31b-it).

התקנת יחסי התלות של Cluster Toolkit

כדי להשלים את המדריך הזה מתחנת עבודה או מלקוח של Linux או macOS, פועלים לפי השלבים הרלוונטיים בקטע התקנת יחסי תלות במסמכי התיעוד של Cluster Toolkit.

אם אתם משתמשים ב-Cloud Shell, אתם יכולים לדלג על החלק הזה.

התקנת Cluster Toolkit

מורידים ומחלצים את החבילה שנוצרה מראש לגרסה הספציפית של Cluster Toolkit שמוגדרת ב-CLUSTER_TOOLKIT_VERSION (v1.103.0) לספריית העבודה הנוכחית:

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

אימות הסביבה

כדי לוודא ש-Cluster Toolkit יכול להקצות משאבים ולשלוח קובצי אימג' ל-Artifact Registry בשמכם, צריך לאמת את Google Cloud CLI ולהגדיר את Application Default Credentials:

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

יצירת אשכול ב-Cluster Toolkit

כדי ליצור ולפרוס אשכול של Cluster Toolkit עם 32 שבבי TPU מסוג v6e:

  1. יוצרים קטגוריה של Cloud Storage:

    gcloud storage buckets create "gs://${GCS_BUCKET:?}" --project="${PROJECT:?}" --location="${REGION:?}" || true

  2. מכינים את הגדרת התוכנית ומעתיקים את תוכנית הפריסה לדוגמה לספריית עבודה מבודדת:

    mkdir -p tmp
cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

  3. מעדכנים את tmp/gke-tpu-v6e-advanced.yaml כדי להקצות לחשבון השירות של מאגר הצמתים של האשכול (gke-np-sa) את התפקיד "אדמין של Storage"‏ (roles/storage.admin) לכתיבת נקודות ביקורת בקטגוריה של Cloud Storage, מחליפים את סוג המכונה של מאגר הצמתים של המערכת ל-e2-standard-8 ומצמידים את מאגר הצמתים של המערכת לאזור היעד כדי למנוע מחסור בסוגי מכונות באזור:

    # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
# Pin system node pool to a single zone to avoid cross-zone stockouts
sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' tmp/gke-tpu-v6e-advanced.yaml

  4. משתמשים בפקודה gcluster deploy כדי לפרוס את האשכול של Cluster Toolkit באמצעות תוכנית ה-blueprint‏ tmp/gke-tpu-v6e-advanced.yaml, ומעבירים את המשתנים הנדרשים באמצעות הדגל --vars:

    ./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -l IGNORE \
    --auto-approve -w

המרת המודל לפורמט MaxText

כדי לאמן את המודל בפורמט MaxText, צריך להמיר אותו מפורמט Hugging Face לפורמט MaxText.

  1. כדי לפשט את הפקודות הבאות, משתמשים בפקודה gcluster job config כדי להגדיר את פרויקט ברירת המחדל, האשכול והמיקום:

    # Configure gcluster Defaults
./gcluster job config set project "${PROJECT:?}"
./gcluster job config set cluster "${CLUSTER_NAME:?}"
./gcluster job config set location "${REGION:?}"

  2. משתמשים בפקודה gcluster job submit כדי להמיר את המודל מפורמט Hugging Face לפורמט MaxText, ולאחסן אותו בקטגוריה של Cloud Storage:

    ./gcluster job submit --name hf-to-mt \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_maxtext \
        model_name=${MODEL_NAME:?} \
        hf_access_token=${HF_TOKEN:?} \
        base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
        scan_layers=True \
        use_multimodal=False \
        skip_jax_distributed_system=true \
        checkpoint_storage_use_zarr3=0 \
        checkpoint_storage_use_ocdbt=0 \
        hardware=cpu \
        --lazy_load_tensors=True"

    הפקודה כוללת את הדגל --await-job-completion, כלומר היא ממתינה בטרמינל הנוכחי עד לסיום ההמרה.

  3. אופציונלי: כדי להעביר בסטרימינג את יומני ההמרה בזמן שהמשימה פועלת, פותחים טרמינל נפרד בספריית העבודה שבה חילצתם את gcluster, ומריצים את הפקודה gcluster job logs:

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs hf-to-mt --main-only -f

  4. בטרמינל המקורי, מוודאים שקובצי המודל שהומרו זמינים בקטגוריה של Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

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

  1. אחרי שתהליך ההמרה מסתיים, מפעילים את עומס העבודה של SFT:

    ./gcluster job submit --name sft \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
      python3 -m maxtext.trainers.post_train.sft.train_sft \
      run_name=sft \
      base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/ \
      model_name=${MODEL_NAME} \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
      hf_access_token=${HF_TOKEN:?} \
      dataset_type=hf \
      hf_path=HuggingFaceH4/ultrachat_200k \
      per_device_batch_size=1 steps=1000 \
      profiler=xplane \
      checkpoint_storage_use_zarr3=0 \
      checkpoint_storage_use_ocdbt=0 \
      skip_jax_distributed_system=False"

    הפקודה כוללת את הדגל --await-job-completion, כלומר היא ממתינה בטרמינל הנוכחי עד לסיום משימת האימון.

  2. אופציונלי: כדי להעביר בסטרימינג את יומני האימון בזמן אמת בזמן שהמשימה פועלת, פותחים טרמינל נפרד בספריית העבודה שבה חילצתם את gcluster, ומריצים את הפקודה gcluster job logs:

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs sft --main-only -f

    אחרי שההידור הראשוני של XLA מסתיים, הרכז מתעד את התקדמות הפלט של כל שלב בהרצה של 1,000 השלבים בפורמט completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ..., כשההפסד באימון יורד מכ-1.57 בשלב 5 לכ-0.82 בשלב 1,000. לפני שממשיכים לחלק הבא, מחכים עד ששלב 1000 יסתיים והפלט מהפקודה ./gcluster job list יציג את הסטטוס של sft בתור Succeeded.

  3. בטרמינל המקורי, מוודאים שנקודות הבדיקה של אימון ה-SFT נוצרו בקטגוריה של Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

המרת המודל המאומן חזרה לפורמט של Hugging Face

  1. אחרי שעומס העבודה של האימון מסתיים, ממירים את המודל חזרה לפורמט של Hugging Face:

    ./gcluster job submit --name mt-to-hf \
    --compute-type "${TPU_TYPE:?}" \
    --num-slices 1 \
    --image "${CLOUD_IMAGE_NAME:?}" \
    --await-job-completion \
    --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
      python3 -m maxtext.checkpoint_conversion.to_huggingface \
        model_name=${MODEL_NAME:?} \
        hf_access_token=${HF_TOKEN:?} \
        load_parameters_path=gs://${GCS_BUCKET:?}/${MODEL_NAME}/trained/sft/checkpoints/1000/model_params/ \
        base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
        skip_jax_distributed_system=true \
        hardware=cpu \
        scan_layers=True \
        use_multimodal=False \
        weight_dtype=bfloat16 \
        --override_model_architecture"

    הפקודה כוללת את הדגל --await-job-completion, כלומר היא ממתינה בטרמינל הנוכחי עד לסיום ההמרה.

  2. אופציונלי: כדי להציג את יומני ההמרות בשידור חי בזמן שהמשימה פועלת, פותחים טרמינל נפרד בספריית העבודה שבה חילצתם את gcluster, ומריצים את הפקודה gcluster job logs:

    # Use the list command to check status
./gcluster job list

# Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
./gcluster job logs mt-to-hf --main-only -f

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  3. בטרמינל המקורי, מוודאים שהמשקולות הסופיות של Hugging Face Safetensors נוצרו בקטגוריה של Cloud Storage:

    gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

הסרת המשאבים

כדי להימנע מחיובים נוספים, השתמשו בפקודה gcluster destroy כדי למחוק את המשאבים שנוצרו במהלך המדריך הזה:

./gcluster destroy "${CLUSTER_NAME:?}" --robust --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

# To delete the local deployment folder and staged blueprint
rm -rf .ghpc "${CLUSTER_NAME:?}" tmp

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

  • למידע נוסף על Cloud TPU, קראו את המאמר מבוא ל-Cloud TPU.
  • פרטים על הארכיטקטורה וההגדרות של TPU v6e-32 מופיעים במאמר TPU v6e.