הפעלת אימון RL עם כמה מארחים עבור Qwen3-30b-a3b ב-TPU v6e

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

מטרות

  • מתקינים את 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, לוחצים על האווטאר של החשבון ובוחרים באפשרות Access tokens.
    2. בדף Access tokens (אסימוני גישה), לוחצים על Create new token (יצירת אסימון חדש).
    3. בוחרים את סוג הטוקן Read (קריאה) ומזינים שם לטוקן.
    4. טוקן הגישה יוצג. שומרים את האסימון במקום בטוח.

  • באתר Hugging Face, מאשרים את הסכם הרישיון של המודל שרוצים לאמן. במדריך הזה נעשה שימוש במודל qwen3-30b-a3b.

כדי לקבל את ההרשאות שדרושות להשלמת המדריך הזה, צריך לבקש מהאדמין להקצות לכם בפרויקט את תפקידי ה-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 TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-30b-a3b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

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

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

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

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

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

התקנת Cluster Toolkit

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

לדוגמה, אפשר להוריד ולחלץ את החבילה לספריית העבודה הנוכחית באופן הבא:

wget -qO- "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:-v1.103.0}/gcluster_bundle_linux_amd64.tgz" | tar -xz

חילוץ החבילה לספריית העבודה הנוכחית מספק את קובץ ה-gcluster binary ואת תוכניות ה-examples/blueprint שבהן משתמשים בשלבים הבאים.

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

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

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

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

  2. מעתיקים את תרשים האשכול לתיקיית העבודה הנוכחית:

    cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml .

  3. כברירת מחדל, לחשבון השירות של מאגר הצמתים של האשכול אין את ההרשאות הנדרשות לכתיבה לקטגוריה של Cloud Storage. כדי לאפשר לחשבון השירות של מאגר הצמתים לכתוב לקטגוריה של Cloud Storage, צריך להקצות לו את התפקיד Storage Admin. כדי להעניק את התפקיד הזה, צריך לערוך את הקובץ gke-tpu-v6e-advanced.yaml ולעדכן את מודול service-account שנקרא 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
    - artifactregistry.reader

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

    ./gcluster deploy gke-tpu-v6e-advanced.yaml \
    --vars project_id="${PROJECT}" \
    --vars deployment_name="${CLUSTER_NAME}" \
    --vars region="${REGION}" \
    --vars zone="${ZONE}" \
    --vars num_slices="${CLUSTER_NODEPOOL_COUNT}" \
    --vars tpu_topology="${TOPOLOGY}" \
    --vars authorized_cidr="0.0.0.0/0" \
    --vars reservation="${RESERVATION:-}" \
    -l IGNORE --auto-approve -w

  5. מגדירים אימות של Container Registry ומקצים את התפקיד 'אדמין אחסון' (roles/storage.admin) לחשבונות השירות של GKE:

    gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION}-docker.pkg.dev" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet

המרת המודל לפורמט 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. כדי להמיר את המודל מפורמט Hugging Face לפורמט MaxText ולאחסן אותו בקטגוריה של Cloud Storage, משתמשים בפקודה gcluster job submit:

    ./gcluster job submit \
  --name qwen-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='Qwen/Qwen3-30B-A3B-Instruct-2507' \
  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"

  3. משתמשים בפקודה 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 qwen-hf-to-mt --main-only -f

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

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

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

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

./gcluster job submit \
  --name="qwen-rl" \
  --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="(echo 190G > /sys/fs/cgroup/memory.max || echo 190G > /sys/fs/cgroup/memory/memory.limit_in_bytes) 2>/dev/null || true && \
    export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, tunix.generate.vllm_sampler as vs; p = pathlib.Path(vs.__file__); p.write_text(p.read_text().replace('reshard_chunk_size: Optional[int] = None', 'reshard_chunk_size: Optional[int] = 4').replace('reshard_chunk_size=self.config.reshard_chunk_size', 'reshard_chunk_size=4'))\" && \
    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} \
    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 \
    num_test_batches=0 \
    batch_size=8 \
    train_micro_batch_size=8 \
    max_target_length=512 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=1 \
    rollout_expert_parallelism=4 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='Qwen/Qwen3-30B-A3B-Instruct-2507' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.55 \
    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'}}\""

  • כדי לבדוק את הסטטוס של משימת האימון:

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

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

# Check progress of the job
kubectl logs -f \
    -l jobset.sigs.k8s.io/replicatedjob-name=pathways-head \
    -c workload-container

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

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

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

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

./gcluster job submit \
  --name="qwen-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=True \
  use_multimodal=False \
  weight_dtype=bfloat16 \
  --override_model_architecture"

  • כדי לבדוק את הסטטוס של עבודת ההמרה:

    # 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 qwen-mt-to-hf --main-only -f
# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

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

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

הסרת המשאבים

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

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

# To delete the local deployment folder and copied blueprint
rm -rf .ghpc "${CLUSTER_NAME}" gke-tpu-v6e-advanced.yaml

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