הרצת כוונון מפוקח (SFT) ב-Gemma 4 ב-TPU v6e

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

מטרות

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

עלויות

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

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

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

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

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

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

  1. בדף Welcome to Hugging Face, לוחצים על הדמות שמייצגת את החשבון ובוחרים באפשרות Access tokens.
  2. בדף Access tokens (אסימוני גישה), לוחצים על Create new token (יצירת אסימון חדש).
  3. בוחרים את סוג הטוקן Read (קריאה) ומזינים שם לטוקן.
  4. טוקן הגישה יוצג. שומרים את האסימון במקום בטוח.
  • באתר Hugging Face, מאשרים את הסכם הרישיון של המודל שמתכננים לאמן. במדריך הזה נעשה שימוש במודל gemma4-31b.

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

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

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

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

מגדירים את משתני הסביבה על ידי הרצת הסקריפט הבא:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

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

  • ‫YOUR_PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
  • ‫YOUR_REGION: האזור שבו רוצים לפרוס את האשכול.
  • ‫YOUR_ZONE: האזור שבו רוצים לפרוס את האשכול.
  • ‫YOUR_REPOSITORY_NAME: השם של מאגר Artifact Registry לתמונות MaxText.
  • ‫YOUR_RESERVATION_NAME: השם של ההזמנה.
  • ‫YOUR_HF_TOKEN: טוקן הגישה שלכם ל-Hugging Face.
  • ‫YOUR_BUCKET_NAME: שם ייחודי גלובלית לקטגוריה של Cloud Storage.

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

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

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

התקנת Cluster Toolkit

כדי להתקין את חבילת ה-Cluster Toolkit שנבנתה מראש, פועלים לפי ההוראות במאמר בנושא התקנת Cluster Toolkit.

הכנת קובץ אימג' של קונטיינר MaxText

כדי להכין את קובץ האימג' של קונטיינר MaxText, כולל התקנת התלויות הנדרשות, מבצעים את השלבים הבאים:

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

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

  2. יוצרים מאגר Artifact Registry:

    gcloud artifacts repositories create ${REPOSITORY_NAME} \
    --repository-format=docker \
    --location=${REGION} \
    --project=${PROJECT} \
    --description="Docker repository for MaxText images in ${REGION}" || true

  3. יוצרים קובץ בספריית השורש של המאגר עם שם הקובץ cloudbuild.yaml והתוכן הבא:

    steps:
  - name: 'gcr.io/cloud-builders/docker'
    entrypoint: 'bash'
    args:
      - '-c'
      - |
        set -euo pipefail

        # 0. Install prerequisites (if needed)
        apt-get update && apt-get install -y curl || apk add curl || true

        # 1. Install uv
        curl -LsSf https://astral.sh/uv/install.sh | sh
        source $$HOME/.local/bin/env

        # 2. Setup Python environment and install MaxText runner
        uv venv --python 3.12 --seed maxtext_venv
        source maxtext_venv/bin/activate
        uv pip install maxtext[runner]==0.2.3 --resolution=lowest

        # 3. Build the Docker image (Cloud Build has Docker pre-configured)
        build_maxtext_docker_image WORKFLOW=post-training

        # 4. Tag the image properly
        docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}

# Cloud Build automatically pushes images listed here
images:
  - '${_CLOUD_IMAGE_NAME}'

options:
  # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
  machineType: 'E2_HIGHCPU_32'

  4. משתמשים ב-Cloud Build כדי ליצור את קובץ האימג' של MaxText Docker:

    gcloud builds submit . \
    --project=${PROJECT} \
    --region=${REGION} \
    --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

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

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

  1. יוצרים תפקיד בהתאמה אישית לניהול זהויות והרשאות גישה (IAM) בשם gke.gcsfuse.profileUser:

    # The GKE TPU v6e blueprint uses GCS Fuse CSI Storage Profiles which requires a custom IAM role.
# If this role is not already created in your project, you must create it before deploying.
gcloud iam roles create gke.gcsfuse.profileUser \
  --project=${PROJECT} \
  --title="GKE GCSFuse Profile User" \
  --description="Allows scanning GCS buckets for objects, retrieving bucket metadata, and creating Anywhere Caches." \
  --permissions="storage.objects.list,storage.buckets.get,storage.anywhereCaches.create,storage.anywhereCaches.get,storage.anywhereCaches.list,storage.anywhereCaches.update"

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

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

  3. כברירת מחדל, לחשבון השירות של מאגר הצמתים של האשכול אין את ההרשאות הנדרשות לכתיבה לקטגוריה של Cloud Storage. כדי לאפשר לחשבון השירות של מאגר הצמתים לכתוב לקטגוריה של Cloud Storage, צריך להקצות לו את התפקיד Storage Admin. כדי להעניק את התפקיד הזה, צריך לערוך את הקובץ gke-tpu-v6e-advanced.yaml ולעדכן את המודול node_pool_service_account:

    - id: node_pool_service_account
  source: modules/project/service-account
  settings:
    name: gke-np-sa
    project_roles:
    - logging.logWriter
    - monitoring.metricWriter
    - monitoring.viewer
    - stackdriver.resourceMetadata.writer
    - storage.admin            # Change from storage.objectViewer
    - artifactregistry.reader

  4. פורסים את אשכול Cluster Toolkit באמצעות תוכנית הבסיס gke-tpu-v6e-advanced.yaml ומעבירים את המשתנים הנדרשים באמצעות הדגל --vars:

    ./gcluster deploy examples/gke-tpu-v6e/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 \
    -w

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

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

  1. אחרי שמסיימים ליצור את אשכול Cluster Toolkit, צריך להגדיר את Docker:

    # Configure docker for pulling images
gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

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

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

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

    ./gcluster job submit --name hf-to-mt \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

כדי לבדוק את הסטטוס של עבודת ההמרה, מריצים את הפקודה הבאה:

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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

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

./gcluster job submit --name sft \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

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

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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

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

./gcluster job submit --name mt-to-hf \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

כדי לבדוק את הסטטוס של עבודת ההמרה, מריצים את הפקודה הבאה:

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ - though again, it's ~2x the size of the original...

הסרת המשאבים

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

gcluster destroy ${CLUSTER_NAME} --robust
gcloud storage rm -r gs://${GCS_BUCKET}
gcloud artifacts repositories delete ${REPOSITORY_NAME} --location=${REGION} --project=${PROJECT} --quiet

# To delete the local deployment folder
rm -rf .ghpc ${CLUSTER_NAME}

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

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