הרצת כוונון מפוקח (SFT) של מודל Qwen3-14b בכמה מארחים באמצעות MaxText

במדריך הזה נסביר בפירוט איך להריץ כוונון עדין (SFT) מפוקח במודל Qwen3-14b באמצעות MaxText ב-Cloud TPU. תלמדו איך ליצור קובץ אימג' של קונטיינר ייעודי, להקצות אשכול Google Kubernetes Engine ‏ (GKE) באמצעות Pathways בעזרת Accelerated Processing Kit ‏ (XPK), ולהריץ עומס עבודה של אימון מרובה מארחים.

מטרות

  • כאן מוסבר איך ליצור קובץ אימג' של קונטיינר מותאם אישית של MaxText שעברה אופטימיזציה לאימון.
  • הקצאת אשכול GKE באמצעות XPK עם הפעלת Pathways.
  • המרת מודל Qwen3 14b מפורמט Hugging Face לפורמט MaxText.
  • הרצה של עומס עבודה של אימון SFT מרובה-מארחים ב-Cloud TPU.
  • ממירים את המודל שעבר כוונון עדין בחזרה לפורמט של Hugging Face לצורך הצגה.

עלויות

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

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

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

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

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

  • מוודאים שלחשבון המשתמש או לחשבון השירות יש את התפקידים הבאים:
    • ‫roles/compute.admin, כדי ליצור את מכונת ה-VM של ה-Build
    • ‫roles/artifactregistry.admin, כדי לנהל את המאגר ב-Docker
    • ‫roles/storage.admin, כדי לנהל את קטגוריית הנתונים
    • ‫roles/container.admin, כדי ליצור ולנהל את אשכול Google Kubernetes Engine
    • ‫roles/iam.serviceAccountAdmin, כדי ליצור את חשבון השירות של עומס העבודה
    • ‫roles/resourcemanager.projectIamAdmin, כדי להגדיר כללי מדיניות של ניהול זהויות והרשאות גישה (IAM)
    • ‫roles/iam.serviceAccountUser, כדי לפעול בתור חשבון השירות
  • מתקינים ומפעילים את Google Cloud CLI.
  • מוודאים שגרסה Python 3.12 ואילך מותקנת בתחנת העבודה.

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

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

הגדרת הסביבה

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

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

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

  • ‫YOUR_PROJECT_ID: מזהה הפרויקט ב- Google Cloud
  • ‫YOUR_REGION: האזור שבו רוצים להשתמש
  • ‫YOUR_ZONE: האזור שרוצים להשתמש בו
  • ‫YOUR_CLUSTER_NAME: שם לאשכול Google Kubernetes Engine
  • ‫YOUR_GCS_BUCKET: שם ייחודי לקטגוריה של Cloud Storage
  • ‫YOUR_RESERVATION_NAME: הזמנת הקיבולת
  • ‫YOUR_HF_TOKEN: טוקן הגישה שלכם ל-Hugging Face

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

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

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

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

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

    gcloud artifacts repositories create maxtext-images \
    --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.1 --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}"

יצירת אשכול Google Kubernetes Engine

כדי להריץ אימון SFT במודל Qwen3 14b, צריך אשכול Google Kubernetes Engine שמצויד בשבבי TPU. מתקינים את Accelerated Processing Kit (XPK)‎ ויוצרים אשכול GKE עם תמיכה ב-Pathways.

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

הכנת המודל לאימון

להמיר את מודל הבסיס לפורמט MaxText באמצעות עומס עבודה שמבוסס על CPU. אל תריצו את המשימה הזו בכמה מכונות במקביל. הפקודה הבאה כוללת בדיקה כדי לוודא שההמרה פועלת רק בצומת TPU אחד.

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

מעקב אחרי ההתקדמות של המרת המודל

כדי לעקוב אחרי התקדמות ההמרה:

  1. כדי להציג רשימה של ה-Pods שנקבעו באשכול GKE, מריצים את הפקודה kubectl get pod.
  2. מוצאים את הפוד בשם qwen-hf-to-mt-slice-job-0-0-HASH.
  3. כדי לבדוק את הפלט של ה-Pod בזמן אמת, מריצים את הפקודה kubectl logs -f POD_NAME.

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

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

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

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

אפשר לעקוב אחרי הסטטוס של עומס העבודה באמצעות ממשק שורת הפקודה (CLI) של XPK.

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

כדי לראות את היומנים ואת השימוש ב-TPU, משתמשים במסוףGoogle Cloud . אפשר גם להציג את היומנים על ידי הרצת הפקודה הבאה:

kubectl logs -f qwen-training-pathways-head-0-0-HASH

מחליפים את HASH בגיבוב המספרי בשם של ה-Pod. כדי לוודא את הערך של הגיבוב הזה, מריצים את הפקודה kubectl get pod ובודקים את רשימת ה-pods שמוחזרת.

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

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

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

כדי לעקוב אחרי התקדמות ההמרה, מריצים את הפקודה kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH ומחליפים את HASH בגיבוב המספרי בשם של ה-Pod.

אחרי שההמרה תושלם, המודל המותאם שמאוחסן ב-gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ יהיה מוכן לשימוש.

הסרת המשאבים

כדי להימנע מחיובים נוספים, צריך למחוק את המשאבים שנוצרו במהלך המדריך הזה, כולל אשכול Google Kubernetes Engine, קטגוריה של Cloud Storage ומאגר Artifact Registry.

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

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

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

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