‫Google משתמשת בטכנולוגיית AI כדי לתרגם תוכן לשפה המועדפת עליך. בתרגומים כאלו עשויות להיות שגיאות.

כוונון עדין של מודלים במארחים מרובים של TPU באמצעות Ray ו-TorchTPU

Autopilot Standard

במדריך הזה נסביר איך להריץ עומסי עבודה מבוזרים של PyTorch ולכוונן מודל Gemma 4 בפרוסות Cloud TPU בכמה מארחים ב-Google Kubernetes Engine ‏ (GKE) באמצעות Ray Core,‏ Ray Train,‏ PyTorch Fully Sharded Data Parallel ‏ (FSDP2) ו-TorchTPU.

המדריך הזה מיועד למהנדסי למידת מכונה (ML), לאדמינים ולמפעילי פלטפורמות ולמומחי נתונים ו-AI שרוצים לתזמר עומסי עבודה מבוזרים של אימון PyTorch בפרוסות TPU עם כמה מארחים ב-Kubernetes. למידע נוסף על תפקידים נפוצים ומשימות לדוגמה שמוזכרות בתוכן של Google Cloud , קראו את המאמר תפקידים ומשימות נפוצים של משתמשי GKE.

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

רקע

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

יחידות TPU

יחידות Tensor Processing Unit‏ (TPU) הן מאיצים ש-Google פיתחה בהתאמה אישית לעומסי עבודה של למידת מכונה. המדריך הזה תומך גם ב-TPU Trillium (v6e) וגם ב-Ironwood (TPU7x). בפרוסת TPU עם כמה מארחים, כמה מארחי TPU מתקשרים ישירות דרך רשת Inter-Chip Interconnect‏ (ICI) במהירות גבוהה, ומספקים את הזיכרון המצטבר ברוחב פס גבוה (HBM) ואת התקשורת הקולקטיבית עם זמן אחזור קצר שנדרשים לאימון מודלים שחורגים מקיבולת הזיכרון של שבב יחיד או של מארח TPU יחיד.

TorchTPU

‫TorchTPU הוא קצה עורפי של PyTorch בקוד פתוח ל-Cloud TPU, שפותח בשיתוף פעולה בין Google לבין Meta. בעזרת TorchTPU, אתם יכולים להריץ מודלים קיימים של PyTorch על חומרת TPU עם שינויים מינימליים בקוד, ולהרחיב את עומסי העבודה של האימון וההסקה משבב יחיד למארזי TPU Pod עם כמה פרוסות באמצעות מהדר Accelerated Linear Algebra ‏ (XLA). ‫TorchTPU משתמש במודל הפעלה של תהליך אחד לכל מכשיר ברשת TPU ICI, ותומך בממשקי API סטנדרטיים של PyTorch, כולל ביצוע מיידי, torch.compile, torch.distributed, DeviceMesh, DTensor ו-Fully Sharded Data Parallel (FSDP2)‎.

‫KubeRay ו-Ray ב-TPU

‫KubeRay הוא אופרטור של Kubernetes לפריסה ולניהול של אפליקציות Ray. כשמפעילים את התוסף Ray on GKE, נפרסים מפעיל KubeRay המנוהל וה-webhook של KubeRay TPU. כש-KubeRay יוצר Pods של Worker שמבקשים משאבי google.com/tpu, ה-webhook מגדיר באופן אוטומטי את שמות המארחים של ה-Pod, את השירות ללא ממשק משתמש ואת משתני הסביבה של ה-TPU שנדרשים לתיאום בין כמה מארחים.

‫Ray מזהה את הטופולוגיה של ה-TPU בכל צומת בזמן ההפעלה, ומתזמן עומסי עבודה מבוזרים בפרוסות TPU. אתם יכולים לתזמר עומסי עבודה מבוזרים של TorchTPU בשתי רמות:

‫Ray Core מספק קבוצות מיקום (placement groups) עם מודעות לטופולוגיה ותזמון שחקנים ללולאות ביצוע מבוזרות בהתאמה אישית, כמו חישוב מקבילי של טנזורים או צינורות עיבוד נתונים של למידת חיזוק (RL).
‫Ray Train מספק אימון מנוהל של מודלים ונתונים במקביל, כולל PyTorch Distributed Data Parallel ‏ (DDP), ‏PyTorch FSDP2 ו-JAX single program, multiple data ‏ (SPMD), עם נקודות ביקורת מובנות ושחזור אוטומטי במקרה של כשל.

מטרות

במדריך הזה מוסבר איך:

יוצרים אשכול GKE במצב Autopilot או Standard ב-TPU Trillium (v6e) או Ironwood (TPU7x).
יוצרים קובץ אימג' של קונטיינר שמוסיף שכבות של ספריות אימון של Ray ו-Hugging Face מעל קובץ האימג' הבסיסי של TorchTPU.
פורסים משאב בהתאמה אישית של RayCluster כדי להקצות את צומת ה-head ואת צומתי העובדים של ה-TPU.
הפעלת חישוב מבוזר של טנזור מקביל באמצעות PyTorch DeviceMesh ו-DTensor בפרוסת TPU באמצעות Ray Core.
כוונון עדין של מודל Gemma 4 באמצעות PyTorch FSDP2 בפרוסת TPU עם כמה מארחים באמצעות Ray Train.

עלויות

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

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

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

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

לפני שמתחילים, חשוב לוודא שביצעתם את המשימות הבאות:

מפעילים את Artifact Registry API, את Cloud Storage API ואת Google Kubernetes Engine API.

הפעלת ממשקי ה-API

כדי להשתמש ב-Google Cloud CLI למשימה הזו, מתקינים ואז מאתחלים את ה-CLI של gcloud. אם כבר התקנתם את ה-CLI של gcloud, הריצו את הפקודה gcloud components update כדי לקבל את הגרסה העדכנית. יכול להיות שבגרסאות קודמות של ה-CLI של gcloud אי אפשר להריץ את הפקודות שבמסמך הזה.
הערה: אם ה-CLI של gcloud כבר מותקן, צריך להגדיר את המאפיין compute/region. אם אתם משתמשים בעיקר באשכולות אזוריים, הגדירו במקום זאת את compute/zone. אם מגדירים מיקום ברירת מחדל, אפשר להימנע משגיאות ב-CLI של gcloud כמו: One of [--zone, --region] must be supplied: Please specify location. יכול להיות שתצטרכו לציין את המיקום בפקודות מסוימות אם המיקום של האשכול שונה מברירת המחדל שהגדרתם.

מוודאים שיש לפרויקט מכסה מספיקה לקיבולת של TPU Trillium (v6e) או Ironwood (TPU7x) באזור שנבחר. מידע נוסף זמין במאמר מכסות של Cloud TPU.
מוודאים שלחשבון שלכם יש את התפקידים הבאים ב-IAM:
- אדמין ב-Kubernetes Engine ‏ (roles/container.admin)
- אדמין ב-Artifact Registry ‏ (roles/artifactregistry.admin)
- אדמין באחסון (roles/storage.admin)

הכנת הסביבה

במדריך הזה תשתמשו ב-Cloud Shell כדי לנהל משאבים שמתארחים ב- Google Cloud. ב-Cloud Shell מותקנת מראש התוכנה שדרושה למדריך הזה, כולל kubectl,‏ docker,‏ python3 ו-gcloud CLI.

כדי להגדיר את הסביבה באמצעות Cloud Shell:

במסוף Google Cloud , לוחצים על Activate Cloud Shell במסוףGoogle Cloud כדי להפעיל סשן של Cloud Shell. הפעולה הזו מפעילה סשן בחלונית התחתונה של מסוף Google Cloud .

משכפלים את המאגר kubernetes-engine-samples ועוברים לספריית הדוגמאות של המדריך הזה:

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples/ai-ml/gke-ray/raytrain/torchtpu

יוצרים ומפעילים סביבה וירטואלית מקומית של Python עם Ray CLI מותקן:
```
python3 -m venv ~/ray-env
source ~/ray-env/bin/activate
pip install "ray[default]==2.58.0"
```
כדי להוריד את משקלי המודל של Gemma 4, צריך לאשר את הסכם הרישיון של Google ב-Hugging Face:
1. עוברים אל דף המודל של Gemma 4 ב-Hugging Face.
2. נכנסים לחשבון ומאשרים את תנאי הרישיון בלחיצה על Agree and access repository.
3. בהגדרות החשבון ב-Hugging Face, יוצרים טוקן גישה למשתמש עם הרשאות Read.
מגדירים את משתני הסביבה של הפרויקט, האשכול, הקטגוריה של Cloud Storage והאסימון של Hugging Face ב- Google Cloud :
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=CLUSTER_NAME
export REGION=REGION
export ZONE=ZONE
export ARTIFACT_REGISTRY=ARTIFACT_REGISTRY
export BUCKET_NAME=BUCKET_NAME
export HF_TOKEN=HF_TOKEN
export IMAGE_URI=${REGION}-docker.pkg.dev/${PROJECT_ID}/${ARTIFACT_REGISTRY}/ray-torchtpu:latest
gcloud config set project ${PROJECT_ID}
```
מחליפים את מה שכתוב בשדות הבאים:
- ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
- ‫CLUSTER_NAME: השם של אשכול GKE, למשל ray-torchtpu-cluster.
- ‫REGION: האזור שבו נמצאים המאגר של Artifact Registry, הקטגוריה של Cloud Storage והקיבולת של ה-TPU.
- ‫ZONE: האזור שבו קיבולת ה-TPU זמינה. במאמר אזורים ותחומים של TPU מפורטים האזורים והתחומים שזמינים לכל גרסת TPU.
- ARTIFACT_REGISTRY: השם של מאגר Docker ב-Artifact Registry, למשל ray-torchtpu.
- ‫BUCKET_NAME: השם של הקטגוריה ב-Cloud Storage שבה מאוחסנות נקודות הביקורת של Ray Train, למשל PROJECT_ID-ray-torchtpu.
- ‫HF_TOKEN: טוקן הגישה ל-Hugging Face.

יצירת אשכול GKE

אתם יכולים להריץ עומסי עבודה מבוזרים של אימון TorchTPU באשכול GKE Autopilot או Standard, כשהתוסף המנוהל Ray operator מופעל.

יוצרים את אשכול GKE ומגדירים את פרטי הכניסה לכלי שורת הפקודה kubectl:
טייס אוטומטי
1. יוצרים אשכול GKE Autopilot אזורי עם תוסף Ray operator מופעל:
  
  gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION} \ --cluster-version=1.36 \ --enable-ray-operator
  
  כשפורסים את המשאב בהתאמה אישית של RayCluster בשלב מאוחר יותר, GKE Autopilot מקצה באופן אוטומטי את צומת ה-CPU הראשי ואת פרוסת ה-TPU עם כמה מארחים שתואמים לבקשות המשאבים של Pod ולבוררי הצמתים.
2. מגדירים את כלי שורת הפקודה kubectl כדי לתקשר עם האשכול:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --region=${REGION}
רגילה
1. יוצרים אשכול GKE Standard אזורי עם מאגר צומתי CPU לאירוח ה-Pod הראשי של Ray, עם הזרמת תמונות, GKE Dataplane V2 ואיחוד זהויות של עומסי עבודה ל-GKE:
  
  gcloud container clusters create ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --cluster-version=1.36 \ --machine-type=n2-standard-8 \ --num-nodes=1 \ --addons=RayOperator \ --enable-image-streaming \ --enable-dataplane-v2 \ --workload-pool=${PROJECT_ID}.svc.id.goog
  הערה: אם אתם מנהלים בעצמכם את KubeRay בקוד פתוח במקום להשתמש בתוסף RayOperator של GKE, התקינו את KubeRay ואת KubeRay TPU webhook בגרסה 1.4.0 ואילך כדי להחדיר את משתנה הסביבה TorchTPU.
2. מגדירים את כלי שורת הפקודה kubectl כדי לתקשר עם האשכול:
  
  gcloud container clusters get-credentials ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
3. יוצרים מאגר של פרוסות TPU עם מספר מארחים ו-16 מכשירים, עם טופולוגיית 4x4 בארבעה מארחים ב-TPU Trillium (v6e) או טופולוגיית 2x2x2 בשני מארחים ב-Ironwood (TPU7x):
  
  Trillium (v6e)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=ct6e-standard-4t \ --tpu-topology=4x4 \ --num-nodes=4
  
  ‫Ironwood‏ (TPU7x)
  gcloud container node-pools create tpu-multi-host-pool \ --project=${PROJECT_ID} \ --cluster=${CLUSTER_NAME} \ --zone=${ZONE} \ --machine-type=tpu7x-standard-4t \ --tpu-topology=2x2x2 \ --num-nodes=2
  
  טיפ: כדי להשתמש בהזמנה ספציפית של TPU, מוסיפים את הדגלים --reservation-affinity=specific ו---reservation=RESERVATION_NAME לפקודה node-pools create, כאשר RESERVATION_NAME הוא השם של הזמנת ה-TPU. כדי להשתמש במכונות וירטואליות במודל Spot, מוסיפים את הדגל --spot.

יוצרים קטגוריה של Cloud Storage כדי לאחסן נקודות בקרה משותפות של Ray Train בצומתי TPU עם כמה מארחים, ונותנים למשאב default Kubernetes ServiceAccount גישה לקטגוריה:

gcloud storage buckets create gs://${BUCKET_NAME} --location=${REGION}

PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/default" \
    --role="roles/storage.objectAdmin"

יוצרים Kubernetes Secret עם טוקן הגישה של Hugging Face כדי ש-Ray Pods יוכלו להוריד את המשקלים של Gemma 4:
```
kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN}
```

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

יוצרים קובץ אימג' של קונטיינר שמתקין את Ray ואת ספריות האימון של Hugging Face על גבי קובץ האימג' הבסיסי של TorchTPU, שכולל את TorchTPU ואת הספרייה libtpu.

יוצרים מאגר Docker ב-Artifact Registry ומגדירים את האימות של Docker:

gcloud artifacts repositories create ${ARTIFACT_REGISTRY} \
    --repository-format=docker \
    --location=${REGION}

gcloud auth configure-docker ${REGION}-docker.pkg.dev,us-docker.pkg.dev

בודקים את Dockerfile בספרייה לדוגמה:

FROM us-docker.pkg.dev/ml-oss-artifacts-published/torch-tpu-public-release-container/torch-tpu:latest

# Install wget for KubeRay health probes, Ray (Core, Dashboard, and Train), and Hugging Face libraries for fine-tuning
RUN apt-get update && apt-get install -y --no-install-recommends wget \
    && rm -rf /var/lib/apt/lists/* \
    && pip install --no-cache-dir \
    "ray[default,train]==2.58.0" \
    "transformers>=5.5.0" \
    datasets

כדי לאמן פרוסה אחת עם המחלקה TorchTrainer של Ray ולתזמן לכל צ'יפלט ב-Ironwood‏ (TPU7x), צריך את Ray בגרסה 2.58.0 ואילך. כדי לבצע כוונון עדין של Gemma 4, צריך את הספרייה transformers בגרסה 5.5.0 ואילך.

יוצרים את קובץ האימג' ומעבירים אותו בדחיפה ל-Artifact Registry:
```
docker build -t ${IMAGE_URI} .
docker push ${IMAGE_URI}
```

פריסת המשאב בהתאמה אישית של RayCluster

פורסים משאב בהתאמה אישית של RayCluster שמקצה Pod ראשי של מעבד (CPU), נפח /dev/shm בזיכרון למאגר האובייקטים של הזיכרון המשותף של Ray וקבוצת וורקרים של TPU עם כמה מארחים.

בודקים את המניפסט של RayCluster לגרסת ה-TPU במאגר לדוגמה:

Trillium (v6e)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-torchtpu-cluster
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: v6e-16
    replicas: 1
    numOfHosts: 4
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          resources:
            limits:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "200Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 100Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x4

‫Ironwood‏ (TPU7x)

apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: ray-torchtpu-cluster
spec:
  headGroupSpec:
    rayStartParams:
      num-cpus: "0"
    template:
      spec:
        containers:
        - name: ray-head
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          ports:
          - containerPort: 6379
            name: gcs
          - containerPort: 8265
            name: dashboard
          - containerPort: 10001
            name: client
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
            requests:
              cpu: "4"
              memory: "16Gi"
  workerGroupSpecs:
  - groupName: tpu7x-16
    replicas: 1
    numOfHosts: 2
    resources:
      # Advertises 8 logical TPU cores (2 chiplets per physical chip) per host.
      TPU: "8"
    template:
      spec:
        containers:
        - name: ray-worker
          image: ${IMAGE_URI}
          env:
          - name: HF_TOKEN
            valueFrom:
              secretKeyRef:
                name: hf-secret
                key: hf_api_token
          - name: RAY_TPU_RESOURCE_PER_CHIP
            value: "2"
          resources:
            limits:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
            requests:
              cpu: "64"
              memory: "250Gi"
              google.com/tpu: "4"
          volumeMounts:
          - mountPath: /dev/shm
            name: dshm
        volumes:
        - name: dshm
          emptyDir:
            medium: Memory
            sizeLimit: 128Gi
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu7x
          cloud.google.com/gke-tpu-topology: 2x2x2

ב-Ironwood (TPU7x), כל שבב פיזי מכיל שני צ'יפלטים, שמספקים שמונה ליבות TPU לוגיות לכל מארח עם ארבעה שבבים. המניפסט מגדיר את הערך של השדה TPU ל-"8" בקטע resources, ומגדיר את משתנה הסביבה RAY_TPU_RESOURCE_PER_CHIP ל-"2" כדי לפרסם ולתזמן שמונה משאבי TPU לוגיים לכל מארח.

מחילים את המניפסט כדי ליצור את המשאב בהתאמה אישית של RayCluster:

Trillium (v6e)

envsubst < ray-cluster.tpu-v6e-multihost.yaml | kubectl apply -f -

‫Ironwood‏ (TPU7x)

envsubst < ray-cluster.tpu-v7x-multihost.yaml | kubectl apply -f -

מחכים עד שה-Pod הראשי של RayCluster וכל ה-Pods של TPU worker יגיעו למצב Running:

kubectl get pods -l ray.io/cluster=ray-torchtpu-cluster

הפלט אמור להיראות כך:

Trillium (v6e)

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-pdpk6              1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-brwpc     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-dx4w7     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-h4d4w     1/1     Running   0          2m30s
ray-torchtpu-cluster-v6e-16-worker-t2g24     1/1     Running   0          2m30s

‫Ironwood‏ (TPU7x)

NAME                                         READY   STATUS    RESTARTS   AGE
ray-torchtpu-cluster-head-k9x2m              1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-7cwpq   1/1     Running   0          2m30s
ray-torchtpu-cluster-tpu7x-16-worker-m98r9   1/1     Running   0          2m30s

בודקים את משתני הסביבה באחד מה-Pods של ה-workers כדי לוודא שה-webhook של KubeRay TPU החדיר את משתני התיאום של ה-TPU עם כמה מארחים ואת משתני הטופולוגיה של TorchTPU:

kubectl get pods -l ray.io/node-type=worker \
    -o=jsonpath='{range .items[0].spec.containers[0].env[*]}{.name}={.value}{"\n"}{end}' \
    | grep -E "TPU_WORKER_ID|TPU_WORKER_HOSTNAMES|TPU_NAME|TORCH_TPU_TOPOLOGY|TORCH_TPU_SLICEBUILDER_ADDRESSES"

הפלט אמור להיראות כך:

Trillium (v6e)

TPU_WORKER_HOSTNAMES=v6e-16-0-0.ray-torchtpu-cluster-headless,v6e-16-0-1.ray-torchtpu-cluster-headless,v6e-16-0-2.ray-torchtpu-cluster-headless,v6e-16-0-3.ray-torchtpu-cluster-headless
TORCH_TPU_TOPOLOGY=4,4,1
TORCH_TPU_SLICEBUILDER_ADDRESSES=v6e-16-0-0.ray-torchtpu-cluster-headless:8471,v6e-16-0-0.ray-torchtpu-cluster-headless:8472,v6e-16-0-0.ray-torchtpu-cluster-headless:8473,v6e-16-0-0.ray-torchtpu-cluster-headless:8474,v6e-16-0-1.ray-torchtpu-cluster-headless:8471,v6e-16-0-1.ray-torchtpu-cluster-headless:8472,v6e-16-0-1.ray-torchtpu-cluster-headless:8473,v6e-16-0-1.ray-torchtpu-cluster-headless:8474,v6e-16-0-2.ray-torchtpu-cluster-headless:8471,v6e-16-0-2.ray-torchtpu-cluster-headless:8472,v6e-16-0-2.ray-torchtpu-cluster-headless:8473,v6e-16-0-2.ray-torchtpu-cluster-headless:8474,v6e-16-0-3.ray-torchtpu-cluster-headless:8471,v6e-16-0-3.ray-torchtpu-cluster-headless:8472,v6e-16-0-3.ray-torchtpu-cluster-headless:8473,v6e-16-0-3.ray-torchtpu-cluster-headless:8474
TPU_WORKER_ID=3
TPU_NAME=v6e-16-0

‫Ironwood‏ (TPU7x)

TPU_WORKER_HOSTNAMES=tpu7x-16-0-0.ray-torchtpu-cluster-headless,tpu7x-16-0-1.ray-torchtpu-cluster-headless
TORCH_TPU_TOPOLOGY=2,2,2,2
TORCH_TPU_SLICEBUILDER_ADDRESSES=tpu7x-16-0-0.ray-torchtpu-cluster-headless:8471,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8472,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8473,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8474,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8475,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8476,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8477,tpu7x-16-0-0.ray-torchtpu-cluster-headless:8478,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8471,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8472,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8473,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8474,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8475,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8476,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8477,tpu7x-16-0-1.ray-torchtpu-cluster-headless:8478
TPU_WORKER_ID=1
TPU_NAME=tpu7x-16-0

מגדירים סשן של העברה ליציאה אחרת לשירות הראשי של Ray:
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
אפשר לבדוק את צומתי העובדים של ה-TPU ואת התוויות של ray.io/tpu-* שזוהו באופן אוטומטי בכרטיסייה Cluster ב-Ray Dashboard.

הפעלת מקביליות מבוזרת של טנזורים באמצעות Ray Core

‫Ray Core מספק את הפונקציה slice_placement_group במודול ray.util.tpu כדי להקצות פרוסת TPU שלמה עם כמה מארחים, ואת המחלקה PlacementGroupSchedulingStrategy כדי לתזמן שחקן Ray אחד לכל מכשיר TPU לוגי בפרוסה.

בודקים את הסקריפט ray_core_collectives.py במאגר לדוגמה:

import argparse
import os

import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from ray.util.tpu import (
    get_tpu_version_from_type,
    slice_placement_group,
)
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    ColwiseParallel,
    RowwiseParallel,
    parallelize_module,
)
import torch.nn as nn
import torch_tpu  # noqa: F401

RANDOM_SEED = 42


class FeedForward(nn.Module):
    """Two-layer feed-forward block for tensor-parallel sharding across a TPU slice."""

    def __init__(self, hidden_dim: int):
        super().__init__()
        self.w1 = nn.Linear(hidden_dim, 4 * hidden_dim, bias=False)
        self.act = nn.SiLU()
        self.w2 = nn.Linear(4 * hidden_dim, hidden_dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.w2(self.act(self.w1(x)))


@ray.remote
class TorchTPUSliceWorker:
    """Ray actor bound to a single logical TPU device within a slice placement group."""

    def __init__(self, devices_per_host: int):
        self.node_ip = ray.util.get_node_ip_address()
        hosts = os.environ["TPU_WORKER_HOSTNAMES"].split(",")
        self.worker_id = int(os.environ.get("TPU_WORKER_ID", "0"))
        self.local_rank = int(
            ray.get_runtime_context().get_accelerator_ids()["TPU"][0]
        )
        self.rank = self.worker_id * devices_per_host + self.local_rank
        # Set per-process PyTorch env:// variables from the webhook's host-level
        # TPU metadata and Ray's assigned local device ID (TorchTrainer sets these
        # automatically in the Ray Train section below).
        os.environ.update(
            {
                "MASTER_ADDR": hosts[0],
                "MASTER_PORT": "29500",
                "WORLD_SIZE": str(len(hosts) * devices_per_host),
                "RANK": str(self.rank),
                "LOCAL_RANK": str(self.local_rank),
                "LOCAL_WORLD_SIZE": str(devices_per_host),
            }
        )

    def run_tensor_parallel_step(
        self,
        batch_size: int = 64,
        hidden_dim: int = 128,
    ) -> dict:
        dist.init_process_group(backend="tpu_dist")
        world_size = dist.get_world_size()
        device = torch.device("tpu")

        # 1. Compute single-device reference output on CPU with identical seed.
        torch.manual_seed(RANDOM_SEED)
        ref_model = FeedForward(hidden_dim).eval()
        inputs_cpu = torch.randn(batch_size, hidden_dim)
        with torch.no_grad():
            ref_output = ref_model(inputs_cpu)

        # 2. Shard the model across all TPU devices in the slice with DTensor TP.
        torch.manual_seed(RANDOM_SEED)
        tp_model = FeedForward(hidden_dim).to(device).eval()
        mesh = init_device_mesh("tpu", (world_size,))
        tp_model = parallelize_module(
            tp_model,
            mesh,
            {"w1": ColwiseParallel(), "w2": RowwiseParallel()},
        )

        with torch.no_grad():
            tp_output = tp_model(inputs_cpu.to(device)).cpu()

        torch.testing.assert_close(tp_output, ref_output, rtol=1e-3, atol=1e-2)
        max_abs_diff = float((tp_output - ref_output).abs().max().item())

        # 3. Verify direct ICI all_reduce across all ranks in the slice.
        rank_tensor = torch.full(
            (4,), float(self.rank + 1), dtype=torch.float32, device=device
        )
        dist.all_reduce(rank_tensor, op=dist.ReduceOp.SUM)
        all_reduce_sum = rank_tensor.cpu().tolist()

        dist.destroy_process_group()
        return {
            "rank": self.rank,
            "local_rank": self.local_rank,
            "worker_id": self.worker_id,
            "node_ip": self.node_ip,
            "all_reduce_sum": all_reduce_sum,
            "max_abs_diff": max_abs_diff,
        }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run distributed TorchTPU tensor parallelism across a Ray TPU slice."
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU physical topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()

    slice_handle = slice_placement_group(
        topology=args.topology,
        accelerator_version=get_tpu_version_from_type(args.accelerator_type),
        num_slices=1,
        resources_per_bundle={"TPU": 1},
    )
    ray.get(slice_handle.placement_group.ready())

    devices_per_host = slice_handle.num_bundles // slice_handle.num_hosts
    workers = [
        TorchTPUSliceWorker.options(
            resources={"TPU": 1},
            scheduling_strategy=PlacementGroupSchedulingStrategy(
                placement_group=slice_handle.placement_group,
                placement_group_bundle_index=bundle_idx,
            ),
        ).remote(devices_per_host)
        for bundle_idx in range(slice_handle.num_bundles)
    ]
    results = sorted(
        ray.get([w.run_tensor_parallel_step.remote() for w in workers]),
        key=lambda r: r["rank"],
    )

    expected_sum = sum(range(1, slice_handle.num_bundles + 1))
    for r in results:
        print(
            f"Rank {r['rank']:2d} (worker_id={r['worker_id']}, "
            f"local_rank={r['local_rank']}, ip={r['node_ip']}): "
            f"all_reduce_sum={r['all_reduce_sum']}"
        )
        assert r["all_reduce_sum"] == [float(expected_sum)] * 4

    print(
        f"All {len(results)} TorchTPU ranks verified DTensor TP parity "
        f"(max_abs_diff={results[0]['max_abs_diff']:.6f}) and "
        f"all_reduce sum = {expected_sum:.1f}"
    )
    for w in workers:
        ray.kill(w)
    slice_handle.shutdown()


if __name__ == "__main__":
    main()

הסקריפט מריץ מקביליות טנזור מבוזרת בפרוסת ה-TPU בשלושה שלבים:

הפונקציה slice_placement_group שומרת פרוסת TPU של 16 מכשירים ומתזמנת שחקן TorchTPUSliceWorker אחד בכל מכשיר TPU.
כל שחקן מחשב את הדירוג הגלובלי שלו לפי אינדקס המארח ומזהה המכשיר המקומי, ואז מאתחל את קבוצת התהליכים של tpu_dist.
הפונקציות init_device_mesh ו-parallelize_module מפצלות את שכבות ה-feed-forward בין כל 16 מכשירי ה-TPU, מאמתות את הפלט מול קובץ עזר של מעבד במכשיר אחד ומריצות all_reduce קולקטיבי ב-ICI.

שולחים את המשימה של Ray Core לגרסת ה-TPU:

Trillium (v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=4x4 --accelerator-type=TPU-V6E

‫Ironwood‏ (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_core_collectives.py --topology=2x2x2 --accelerator-type=TPU-V7X

הפלט אמור להיראות כך:

Trillium (v6e)

Rank  0 (worker_id=0, local_rank=0, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=10.72.3.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=1, local_rank=0, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=1, local_rank=1, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=1, local_rank=2, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=1, local_rank=3, ip=10.72.14.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=2, local_rank=0, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=2, local_rank=1, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=2, local_rank=2, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=2, local_rank=3, ip=10.72.7.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=3, local_rank=0, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=3, local_rank=1, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=3, local_rank=2, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=3, local_rank=3, ip=10.72.6.20): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_HFTkKrjSnfza3ZFD' succeeded
------------------------------------------

‫Ironwood‏ (TPU7x)

Rank  0 (worker_id=0, local_rank=0, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  1 (worker_id=0, local_rank=1, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  2 (worker_id=0, local_rank=2, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  3 (worker_id=0, local_rank=3, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  4 (worker_id=0, local_rank=4, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  5 (worker_id=0, local_rank=5, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  6 (worker_id=0, local_rank=6, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  7 (worker_id=0, local_rank=7, ip=172.16.14.6): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  8 (worker_id=1, local_rank=0, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank  9 (worker_id=1, local_rank=1, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 10 (worker_id=1, local_rank=2, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 11 (worker_id=1, local_rank=3, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 12 (worker_id=1, local_rank=4, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 13 (worker_id=1, local_rank=5, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 14 (worker_id=1, local_rank=6, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
Rank 15 (worker_id=1, local_rank=7, ip=172.16.15.7): all_reduce_sum=[136.0, 136.0, 136.0, 136.0]
All 16 TorchTPU ranks verified DTensor TP parity (max_abs_diff=0.002470) and all_reduce sum = 136.0
------------------------------------------
Job 'raysubmit_MrJbAe3hvARBcBN2' succeeded
------------------------------------------

כוונון עדין של מודל Gemma 4 באמצעות Ray Train ו-PyTorch FSDP2

הכיתה TorchTrainer של Ray Train הופכת את ההזמנה של פרוסת TPU, את מיקום ה-worker ואת האתחול של torch.distributed לאוטומטיים בפרוסת TPU עם כמה מארחים.

בודקים את הסקריפט ray_train_llm_finetune.py במאגר לדוגמה, שמבצע כוונון עדין של המודל google/gemma-4-E2B-it במערך הנתונים של ההוראות databricks/databricks-dolly-15k:

import argparse
import os
import tempfile
import time

from datasets import load_dataset
import ray
import ray.train
from ray.train import (
    Checkpoint,
    CheckpointConfig,
    RunConfig,
    ScalingConfig,
)
from ray.train.torch import TorchConfig, TorchTrainer
import torch
from torch.distributed import ReduceOp, all_reduce
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_model_state_dict,
)
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
from torch.utils.data import DataLoader
import torch_tpu  # noqa: F401
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
)


def create_dataloader(
    tokenizer: AutoTokenizer,
    batch_size: int,
    dataset_split: str = "train[:1024]",
    max_seq_len: int = 512,
) -> DataLoader:
    """Loads and tokenizes the Dolly 15k dataset for causal LM fine-tuning."""
    dataset = load_dataset(
        "databricks/databricks-dolly-15k", split=dataset_split
    )
    tokenized = dataset.map(
        lambda x: tokenizer(
            f"Instruction: {x['instruction']}\nResponse: {x['response']}",
            truncation=True,
            max_length=max_seq_len,
            padding="max_length",
        ),
        remove_columns=dataset.column_names,
    )
    dataloader = DataLoader(
        tokenized,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True,
        collate_fn=DataCollatorForLanguageModeling(
            tokenizer=tokenizer, mlm=False
        ),
    )
    return ray.train.torch.prepare_data_loader(dataloader)


def train_func(config: dict) -> None:
    """Per-worker distributed FSDP2 training loop executed on each logical TPU core."""
    ctx = ray.train.get_context()
    world_rank = ctx.get_world_rank()
    world_size = ctx.get_world_size()
    device = ray.train.torch.get_device()

    model_id = config["model_id"]
    dataset_split = config["dataset_split"]
    batch_size = config["batch_size_per_device"]
    num_epochs = config["num_epochs"]
    max_seq_len = 512

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    dataloader = create_dataloader(
        tokenizer=tokenizer,
        batch_size=batch_size,
        dataset_split=dataset_split,
        max_seq_len=max_seq_len,
    )

    torch.manual_seed(42)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        dtype=torch.bfloat16,
        attn_implementation="eager",
        vision_config=None,
        audio_config=None,
    ).to(device)
    model.config.get_text_config().use_cache = False

    mesh = init_device_mesh("tpu", (world_size,))
    mp_policy = MixedPrecisionPolicy(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    )
    text_model = model.model.language_model
    text_model.embed_tokens_per_layer.requires_grad_(False)
    for layer in text_model.layers:
        fully_shard(layer.mlp, mesh=mesh, mp_policy=mp_policy)
    fully_shard(
        model,
        mesh=mesh,
        mp_policy=mp_policy,
        ignored_params=set(text_model.embed_tokens_per_layer.parameters()),
    )

    optimizer = torch.optim.AdamW(
        model.parameters(), lr=2e-5, weight_decay=0.01
    )
    global_batch_size = batch_size * world_size

    for epoch in range(num_epochs):
        dataloader.sampler.set_epoch(epoch)
        model.train()
        epoch_start = time.perf_counter()
        running_loss = torch.zeros(1, device=device, dtype=torch.float32)

        for batch in dataloader:
            optimizer.zero_grad()
            outputs = model(
                input_ids=batch["input_ids"],
                labels=batch["labels"],
            )
            outputs.loss.backward()
            optimizer.step()
            running_loss += outputs.loss.detach().float()

        all_reduce(running_loss, op=ReduceOp.AVG)
        avg_loss = running_loss.item() / len(dataloader)
        epoch_time = time.perf_counter() - epoch_start
        tokens_per_sec = (
            len(dataloader) * global_batch_size * max_seq_len
        ) / epoch_time

        metrics = {
            "epoch": epoch + 1,
            "loss": round(avg_loss, 4),
            "epoch_time_s": round(epoch_time, 2),
            "tokens_per_s": round(tokens_per_sec, 1),
        }
        if world_rank == 0:
            print(
                f"Epoch {epoch + 1}/{num_epochs} | loss={metrics['loss']:.4f} | "
                f"epoch_time={metrics['epoch_time_s']:.2f}s | "
                f"tokens/s={metrics['tokens_per_s']:.1f}"
            )

        full_state_dict = get_model_state_dict(
            model,
            options=StateDictOptions(
                full_state_dict=True,
                cpu_offload=True,
                ignore_frozen_params=True,
            ),
        )
        if world_rank == 0:
            with tempfile.TemporaryDirectory() as tmpdir:
                torch.save(full_state_dict, os.path.join(tmpdir, "model.pt"))
                ray.train.report(
                    metrics, checkpoint=Checkpoint.from_directory(tmpdir)
                )
        else:
            ray.train.report(metrics)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Fine-tune a Hugging Face CausalLM on TPUs with Ray Train and TorchTPU FSDP2."
    )
    parser.add_argument(
        "--model-id",
        type=str,
        default="google/gemma-4-E2B-it",
        help="Hugging Face model ID to fine-tune.",
    )
    parser.add_argument(
        "--topology",
        type=str,
        default="4x4",
        help="TPU slice topology.",
    )
    parser.add_argument(
        "--accelerator-type",
        type=str,
        default="TPU-V6E",
        help="Ray TPU accelerator type (TPU-V6E, TPU-V7X, etc.).",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=16,
        help="Total TPU worker processes across the slice.",
    )
    parser.add_argument(
        "--dataset-split",
        type=str,
        default="train[:1024]",
        help="Hugging Face dataset split to load.",
    )
    parser.add_argument(
        "--num-epochs",
        type=int,
        default=3,
        help="Number of fine-tuning epochs.",
    )
    parser.add_argument(
        "--batch-size-per-device",
        type=int,
        default=4,
        help="Per-device training batch size.",
    )
    parser.add_argument(
        "--storage-path",
        type=str,
        default=os.environ.get("STORAGE_PATH"),
        required="STORAGE_PATH" not in os.environ,
        help="Cloud Storage or shared file system path for checkpoints.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    ray.init()
    run_name = f"torchtpu_gemma4_fsdp2_sft_{int(time.time())}"

    print(
        f"Starting Ray Train FSDP2 SFT job: model={args.model_id}, "
        f"topology={args.topology}, accelerator_type={args.accelerator_type}, "
        f"num_workers={args.num_workers}"
    )

    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        train_loop_config={
            "model_id": args.model_id,
            "dataset_split": args.dataset_split,
            "num_epochs": args.num_epochs,
            "batch_size_per_device": args.batch_size_per_device,
        },
        torch_config=TorchConfig(backend="tpu_dist"),
        scaling_config=ScalingConfig(
            use_tpu=True,
            num_workers=args.num_workers,
            topology=args.topology,
            accelerator_type=args.accelerator_type,
            resources_per_worker={"TPU": 1},
        ),
        run_config=RunConfig(
            name=run_name,
            storage_path=args.storage_path,
            checkpoint_config=CheckpointConfig(
                num_to_keep=1,
                checkpoint_score_attribute="loss",
                checkpoint_score_order="min",
            ),
        ),
    )

    result = trainer.fit()
    print(f"Training completed! Final metrics: {result.metrics}")
    print(f"Saved checkpoint: {result.checkpoint}")


if __name__ == "__main__":
    main()

הסקריפט מגדיר את המחלקה TorchTrainer עם הרכיבים הבאים:

הכיתה ScalingConfig מגדירה את טופולוגיית ה-worker המבוזרת למשימת האימון. הגדרת הפרמטרים use_tpu=True,‏ num_workers=16,‏ topology ו-accelerator_type שומרת פרוסת TPU של 16 מכשירים ומתזמנת וורקר אחד לכל ליבת TPU.
הכיתה TorchConfig מגדירה את קבוצת התהליכים torch.distributed בכל ה-workers. הגדרת הפרמטר backend="tpu_dist" מחברת את כל 16 ה-workers דרך רשת ה-TPU ICI לפני שהפונקציה train_func פועלת.
הפונקציה create_dataloader מורידה את מערך הנתונים של ההוראות ומבצעת לו טוקניזציה, יוצרת אצוות של רצפי הטוקנים באמצעות אובייקט DataLoader של PyTorch, וקוראת לפונקציה ray.train.torch.prepare_data_loader כדי שכל worker יקבל חלק נפרד מכל אצווה גלובלית במכשיר ה-TPU שהוקצה לו.
הפונקציה train_func טוענת את המודל ברמת הדיוק bfloat16 וקוראת לפונקציה fully_shard של PyTorch FSDP2 בכל בלוק של feed-forward ובמודל הבסיס כדי לחלק את הפרמטרים, הגרדיאנטים ומצבי האופטימיזציה בין כל 16 ליבות ה-TPU.
הכיתה Checkpoint והפונקציה ray.train.report אוספות את מילון המצבים המלא של המודל בדירוג 0 בסוף כל תקופה, ומעלות את נקודת הבדיקה לקטגוריה של Cloud Storage.

שולחים את משימת הכוונון העדין של TorchTrainer FSDP2 לגרסת ה-TPU:

Trillium (v6e)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=4x4 \
        --accelerator-type=TPU-V6E \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

‫Ironwood‏ (TPU7x)

ray job submit \
    --address http://localhost:8265 \
    --working-dir . \
    --runtime-env-json '{"excludes": ["ray-env", ".git", "*.whl"]}' \
    -- python3 ray_train_llm_finetune.py \
        --topology=2x2x2 \
        --accelerator-type=TPU-V7X \
        --storage-path=gs://${BUCKET_NAME}/ray_train_torchtpu

הפלט אמור להיראות כך:

Trillium (v6e)

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=4x4, accelerator_type=TPU-V6E, num_workers=16
(TrainController pid=5407) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V6E': 0.001}] * 16
(RayTrainWorker pid=19733, ip=10.72.3.20) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 1/3 | loss=3.1402 | epoch_time=28.17s | tokens/s=18613.5
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-08-10.762175)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 2/3 | loss=2.3318 | epoch_time=26.07s | tokens/s=20112.8
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-10-02.548088)
(RayTrainWorker pid=19733, ip=10.72.3.20) Epoch 3/3 | loss=2.2050 | epoch_time=25.52s | tokens/s=20545.2
(RayTrainWorker pid=19733, ip=10.72.3.20) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.205, 'epoch_time_s': 25.52, 'tokens_per_s': 20545.2}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791864427/checkpoint_2026-10-13_04-11-48.701853)
------------------------------------------
Job 'raysubmit_zbm4pnFZitA1YECA' succeeded
------------------------------------------

‫Ironwood‏ (TPU7x)

Starting Ray Train FSDP2 SFT job: model=google/gemma-4-E2B-it, topology=2x2x2, accelerator_type=TPU-V7X, num_workers=16
(TrainController pid=1162) Attempting to start training worker group of size 16 with the following resources: [{'TPU': 1, 'accelerator_type:TPU-V7X': 0.001}] * 16
(RayTrainWorker pid=257, ip=172.16.26.6) Setting up process group for: env:// [rank=0, world_size=16]
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 1/3 | loss=3.1407 | epoch_time=26.28s | tokens/s=19953.3
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-52-00.026930)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 2/3 | loss=2.3235 | epoch_time=16.68s | tokens/s=31429.9
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-53-33.533341)
(RayTrainWorker pid=257, ip=172.16.26.6) Epoch 3/3 | loss=2.1964 | epoch_time=16.98s | tokens/s=30884.6
(RayTrainWorker pid=257, ip=172.16.26.6) Checkpoint successfully created at: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
Training completed! Final metrics: {'epoch': 3, 'loss': 2.1964, 'epoch_time_s': 16.98, 'tokens_per_s': 30884.6}
Saved checkpoint: Checkpoint(filesystem=gcs, path=my-project-ray-torchtpu/ray_train_torchtpu/torchtpu_gemma4_fsdp2_sft_1791867004/checkpoint_2026-10-13_04-55-01.508219)
------------------------------------------
Job 'raysubmit_zSNnCXhHkwKiSmJh' succeeded
------------------------------------------

פלט האימון מציג את הפרטים הבאים:

תקופה 1 כוללת הידור ראשוני של תרשים XLA. בתקופות 2 ו-3 נעשה שימוש חוזר בקובצי ההפעלה שנשמרו במטמון מ-/dev/shm/torch_tpu_cache, והן מגיעות ליותר מ-20,000 טוקנים לשנייה ב-TPU Trillium (v6e) ול-31,000 טוקנים לשנייה ב-Ironwood (TPU7x).
האובדן הגלובלי של אנטרופיה צולבת, שמסונכרן בין כל 16 הדרגות של ה-worker ב-FSDP2, יורד מ-3.14 בתקופה 1 ל-2.20 בתקופה 3, בשלוש תקופות הכוונון בקבוצת הנתונים של ההוראות ב-databricks/databricks-dolly-15k.
בסוף כל תקופה, דירוג 0 אוסף את מילון המצבים המלא של המודל ומעלה את נקודת הבדיקה לקטגוריה של Cloud Storage, תוך שמירה על נקודת הבדיקה עם האובדן הנמוך ביותר, כפי שהוגדר באובייקט CheckpointConfig(num_to_keep=1, checkpoint_score_attribute="loss", checkpoint_score_order="min").

בדיקת הביצועים של עומסי העבודה

כדי לעקוב אחרי השימוש ב-TPU TensorCore, ב-HBM ובמדדים של אשכול Ray ב-Cloud Monitoring, קראו את המאמרים איסוף והצגה של יומנים ומדדים של אשכולות Ray ב-GKE והצגת מדדי TPU במרכז הבקרה של Ray.

גישה ללוח הבקרה של Ray

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

אם העברה ליציאה אחרת לשירות הראשי של Ray עדיין לא פועלת, מפעילים אותה ברקע:
```
kubectl port-forward svc/ray-torchtpu-cluster-head-svc 8265:8265 >/dev/null 2>&1 &
```
ב-Cloud Shell, לוחצים על Web Preview ובוחרים באפשרות Preview on port 8265, או פותחים את http://localhost:8265 בדפדפן המקומי.
לוחצים על הכרטיסיות Jobs,‏ Cluster ו-Metrics כדי לבדוק את יומני ה-worker, את המצב של קבוצת המיקום, את התוויות של צומת ה-TPU ואת התרשים TPU Utilization by Node.

הסרת המשאבים

כדי להימנע מחיובים בחשבון Google Cloud על המשאבים שבהם השתמשתם במדריך הזה, אתם צריכים למחוק את המשאב המותאם אישית של RayCluster, את אשכול GKE, את הקטגוריה של Cloud Storage ואת המאגר של Artifact Registry:

עוצרים את תהליך העברת היציאה ברקע ומוחקים את המשאב בהתאמה אישית של RayCluster ואת הסוד של Kubernetes:
```
pkill -f "kubectl port-forward" || true
kubectl delete raycluster ray-torchtpu-cluster
kubectl delete secret hf-secret --ignore-not-found
```

מוחקים את אשכול ה-GKE:

טייס אוטומטי

gcloud container clusters delete ${CLUSTER_NAME} \
    --region=${REGION}

רגילה

gcloud container clusters delete ${CLUSTER_NAME} \
    --zone=${ZONE}

מוחקים את הקטגוריה של Cloud Storage ואת התוכן של נקודת הבדיקה:
```
gcloud storage rm --recursive gs://${BUCKET_NAME}
```

מוחקים את המאגר ב-Artifact Registry:

gcloud artifacts repositories delete ${ARTIFACT_REGISTRY} \
    --location=${REGION}

כוונון עדין של מודלים במארחים מרובים של TPU באמצעות Ray ו-TorchTPU קל לארגן דפים בעזרת אוספים אפשר לשמור ולסווג תוכן על סמך ההעדפות שלך.

רקע

יחידות TPU

TorchTPU

‫KubeRay ו-Ray ב-TPU

מטרות

עלויות

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

הכנת הסביבה

יצירת אשכול GKE

טייס אוטומטי

רגילה

Trillium (v6e)

‫Ironwood‏ (TPU7x)

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

פריסת המשאב בהתאמה אישית של RayCluster

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

הפעלת מקביליות מבוזרת של טנזורים באמצעות Ray Core

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

כוונון עדין של מודל Gemma 4 באמצעות Ray Train ו-PyTorch FSDP2

Trillium (v6e)

‫Ironwood‏ (TPU7x)

Trillium (v6e)

‫Ironwood‏ (TPU7x)

בדיקת הביצועים של עומסי העבודה

גישה ללוח הבקרה של Ray

הסרת המשאבים

טייס אוטומטי

רגילה

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

כוונון עדין של מודלים במארחים מרובים של TPU באמצעות Ray ו-TorchTPU