הפעלת אימון למידה מרובה מארחים לחיזוק עבור Gemma 4 26B על TPU v6e

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

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

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

  • מודל: Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it), מודל מתקדם עם משקלים פתוחים שעבר כוונון לפי הוראות. במדריך הזה, משקלי המודל יורדו מ-Hugging Face, יומרו לפורמט שתואם ל-MaxText וייפרסו באשכול GKE עם תשתית Cloud TPU v6e.
  • אלגוריתם למידת חיזוק (RL): אופטימיזציה של מדיניות יחסית לקבוצה (GRPO) באמצעות MaxText ו-Pathways. עומס העבודה מפעיל צינור עיבוד נתונים שניתן להתאמה אחרי האימון, באמצעות מאמן ה-RL המובנה של MaxText (maxtext.trainers.post_train.rl.train_rl) בשילוב עם קצה עורפי של דוגם vLLM. ‫GRPO מבצע אופטימיזציה של המדיניות על ידי יצירת קבוצת פלטים לכל הנחיה והערכת תגמולים יחסיים בלי צורך במודל ביקורת נפרד.
  • קבוצת נתונים ותבנית: OpenMathInstruct-2 (openmathinstruct2_rl.json), קבוצת נתונים גדולה לכוונון הוראות במתמטיקה. עומס העבודה משתמש בתבניות צ'אט מובנות במהלך שלב ההשקה של RL כדי לשפר את יכולות החשיבה המתמטית בכמה שלבים.
  • תשתית ומסגרת: GKE, מנוהל באמצעות Cluster Toolkit, ב-Cloud TPU v6e (פרוסה של ct6e-standard-4t, v6e-32 עם טופולוגיה של 4x8). המסגרת היא MaxText‏ (tpu_post_training:0.2.4) עם Pathways וקצה עורפי של JAX.

מטרות

  • מתקינים את Cluster Toolkit ואת יחסי התלות שלו.
  • פריסת אשכול של Cluster Toolkit.
  • המרת מודל של Hugging Face לפורמט MaxText.
  • מריצים עומס עבודה של אימון RL באשכול TPU v6e.
  • ממירים את המודל המכוונן בחזרה לפורמט של 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 COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-32"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

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

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

התקנת יחסי התלות של 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

אימות הסביבה

לפני שמריצים את הסקריפטים לפריסה, צריך לאמת את Google Cloud CLI ולהגדיר את Application Default Credentials ‏ (ADC). כך Cluster Toolkit יוכל להקצות משאבים ולשלוח קובצי אימג' ל-Artifact Registry בשמכם.

gcloud auth login
gcloud auth application-default login --no-launch-browser
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. משתמשים בפקודה gcluster deploy כדי לפרוס את האשכול של Cluster Toolkit באמצעות תוכנית ה-blueprint‏ gke-tpu-v6e-advanced.yaml, ומעבירים את המשתנים הנדרשים באמצעות הדגל --vars:

    TMP_DIR=$(mktemp -d)
BLUEPRINT_FILE="${TMP_DIR}/gke-tpu-v6e-advanced.yaml"

cp ./examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml "${BLUEPRINT_FILE}"

# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" "${BLUEPRINT_FILE}"
# 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)]' "${BLUEPRINT_FILE}"
# Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" "${BLUEPRINT_FILE}"

./gcluster deploy "${BLUEPRINT_FILE}" \
    --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=${CLUSTER_NODEPOOL_COUNT:?},tpu_topology=${TOPOLOGY:?},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    -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="gemma4-hf-to-mt" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --await-job-completion \
  --command="[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
    python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name=${MODEL_NAME:?} \
    hf_access_token=${HF_TOKEN:?} \
    --hf_model_path='google/gemma-4-26b-a4b-it' \
    base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
    scan_layers=False \
    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 gemma4-hf-to-mt --main-only -f

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

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

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

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

    ./gcluster job submit \
  --name="gemma4-training" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET:?}/pathways/" \
  --gke-ttl-after-finished="24h" \
  --restarts=0 \
  --env="GRPC_DNS_RESOLVER=native" \
  --env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, re; p = pathlib.Path('/deps/src/maxtext/trainers/post_train/rl/utils_rl.py'); p.write_text(re.sub('optax[.]adamw[(][^)]+[)]', 'optax.adafactor(learning_rate=learning_rate)', p.read_text()))\" && \
    JAX_PLATFORMS=proxy,cpu ENABLE_PATHWAYS_PERSISTENCE=1 \
    python3 -m maxtext.trainers.post_train.rl.train_rl \
    run_name=rl \
    base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME:?}/trained/ \
    model_name=${MODEL_NAME} \
    scan_layers=False \
    load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
    hf_access_token=${HF_TOKEN:?} \
    data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \
    num_batches=50 \
    batch_size=4 \
    train_micro_batch_size=2 \
    max_target_length=640 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=4 \
    rollout_expert_parallelism=1 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='google/gemma-4-26b-a4b-it' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.7 \
    remat_policy=full \
    async_scheduling=False \
    allow_split_physical_axes=true \
    ragged_gather_reduce_fallback=True \
    enable_dp_attention=False \
    decode_sampling_temperature=0.8 \
    decode_sampling_top_k=50 \
    decode_sampling_nucleus_p=0.95 \
    learning_rate=2e-5 \
    learning_rate_schedule_steps=100 \
    rl.num_generations=4 \
    rl.reshard_chunk_size=4 \
    debug=True \
    vllm_hf_overrides='{\"architectures\": [\"MaxTextForCausalLM\"]}' \
    vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'model_call_mode': 'inference', 'enable_dp_attention': false, 'allow_split_physical_axes': true, 'use_ragged_sort': false, 'ragged_gather_reduce_fallback': true, 'prefuse_moe_weights': true, 'weight_dtype': 'bfloat16'}}\""

  2. בודקים את הסטטוס, מחכים שה-Pod של המתאם יהיה מוכן ומעבירים בסטרימינג את היומנים של משימת האימון:

    # Use the list command to check status (initially shows "Suspended" briefly while the CPU node pool scales up)
./gcluster job list

# Ensure kubectl is configured
gcloud container clusters get-credentials "${CLUSTER_NAME:?}" \
  --location="${REGION:?}" \
  --project="${PROJECT:?}"

# Wait for the Pathways head pod to be ready before streaming logs
kubectl wait --for=condition=Ready pod \
    -l job-name=gemma4-training-pathways-head-0 \
    --timeout=600s

# Stream training logs (runs for 50 steps; wait for "Post RL Training: ... mean_reward=..." and job status Succeeded before converting to Hugging Face)
kubectl logs -f \
    -l job-name=gemma4-training-pathways-head-0 \
    -c workload-container

    מיד אחרי השליחה, הפלט מהפקודה ./gcluster job list עשוי להציג לזמן קצר את Suspended בזמן ש-GKE מגדיל את מאגר הצמתים של המעבד (CPU) של Pathways coordinator‏ (cpu-np). כשהמשימה מתחילה, היא מפיקה הערכה ראשונית לפני האימון בשלב 0 (Pre RL Training: ... mean_reward=...), משהה את פלט היומן למשך כמה דקות במהלך הידור ה-XLA הראשוני, ואז מאמנת במשך 50 שלבים, ומדפיסה את Reshard finished ואת דורות ההשקה בכל שלב בלי מספרי שלבים.

    לפני שממשיכים לשלב הבא, מחכים עד שפלט היומנים Post RL Training: ... mean_reward=... והפלט מהפקודה ./gcluster job list יציגו את הסטטוס של gemma4-training בתור Succeeded.

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

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/rl/checkpoints/actor/"

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

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

    ./gcluster job submit \
  --name="gemma4-mt-to-hf" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --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/rl/checkpoints/actor/50/model_params/ \
      base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
      skip_jax_distributed_system=true \
      hardware=cpu \
      scan_layers=False \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

  2. אופציונלי: כדי לעקוב אחרי יומני ההמרה בזמן אמת בזמן שפקודת --await-job-completion הקודמת ממתינה בטרמינל הראשון, פותחים טרמינל נפרד בספריית העבודה שבה חילצתם את 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 gemma4-mt-to-hf --main-only -f

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

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

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

הסרת המשאבים

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

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

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

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