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

איך מבצעים כוונון עדין של מודל LLM באמצעות TPUs ב-GKE עם JAX

טייס אוטומטי רגילה

במדריך הזה נסביר איך לכוונן מודל גדול של שפה (LLM) באמצעות יחידות לעיבוד טנסורים (TPU) ב-Google Kubernetes Engine ‏ (GKE) עם JAX. באמצעות כוונון עדין אפשר להתאים מודל בסיס כמו Gemma 3 לדומיין או למשימה ספציפיים. התהליך הזה משפר את הדיוק של המודל על ידי עדכון הפרמטרים שלו באמצעות מערך נתונים מיוחד משלכם.

המדריך הזה הוא נקודת התחלה טובה אם אתם צריכים שליטה גרנולרית, התאמה אישית, יכולת הרחבה, עמידות, ניידות ויעילות כלכלית של Kubernetes מנוהל, כדי לכוונן את עומסי העבודה שלכם ב-AI/ML.

רקע

באמצעות שימוש ב-TPU ב-GKE עם Jax כדי לבצע כוונון עדין של LLM, אתם יכולים לבנות פתרון כוונון עדין חזק ומוכן לייצור עם כל היתרונות של Kubernetes מנוהל.

‏Gemma

‫Gemma היא סדרה של מודלים קלי משקל של AI גנרטיבי/ML מולטימודאלי שזמינים לשימוש חופשי ומופצים ברישיון קוד פתוח. מודלים של AI זמינים להרצה באפליקציות, בחומרה, במכשירים ניידים או בשירותים מתארחים. ‫Gemma 3 מציג יכולות מולטימודאליות, והוא תומך בקלט של שילוב בין ראייה לשפה ופלט של טקסט. הוא יכול לטפל בחלונות הקשר של עד 128,000 טוקנים, ויש לו תמיכה ביותר מ-140 שפות. בנוסף, מודל Gemma 3 מציע יכולות משופרות במתמטיקה, בחשיבה רציונלית ובצ'אט, כולל פלט מובנה וקריאה להפעלת פונקציות.

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

מידע נוסף זמין במסמכי התיעוד של Gemma.

TPUs

יחידות TPU הן מעגלים משולבים לאפליקציות ספציפיות (ASIC) שפותחו על ידי Google כדי להאיץ מודלים של למידת מכונה ו-AI שנבנים באמצעות frameworks כמו TensorFlow,‏ PyTorch ו-JAX.

לפני שמשתמשים ב-TPU ב-GKE, מומלץ להשלים את תוכנית הלימודים הבאה:

מידע על הזמינות של גרסאות TPU עדכניות זמין בארכיטקטורת המערכת של Cloud TPU.
מידע על TPU ב-GKE

JAX

‫JAX היא מסגרת ללמידת מכונה עם ביצועים גבוהים, שמיועדת לשימוש עם מעבדי TPU ו-GPU. ‫JAX מספקת API ליצירה ולאימון של מודלים של למידת מכונה.

מידע נוסף זמין במאגר JAX.

מטרות

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

יוצרים אשכול GKE Autopilot או אשכול רגיל עם טופולוגיית TPU מומלצת, על סמך מאפייני המודל. במהלך המדריך הזה, תבצעו כוונון עדין של מאגרי צמתים עם מארח יחיד.
הוספת נתונים לקטגוריה של Cloud Storage וטעינתה במאגר באמצעות Cloud Storage FUSE.
פורסים את משימת הכוונון העדין של ה-LLM ב-GKE.
עוקבים אחרי משימת הכוונון העדין וצופים ביומנים.

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

נכנסים לחשבון Google Cloud . אנחנו ממליצים למשתמשים חדשים ב- Google Cloud ליצור חשבון כדי שיוכלו להעריך את הביצועים של המוצרים שלנו בתרחישים מהעולם האמיתי. לקוחות חדשים מקבלים בחינם גם קרדיט בשווי 300$ להרצה, לבדיקה ולפריסה של עומסי העבודה.

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API, if it is not already enabled.

Roles required to enable APIs

To enable APIs, you need the serviceusage.services.enable permission. If you created the project, then you likely already have this permission through the Owner role (roles/owner). Otherwise, you can get this permission through the Service Usage Admin role (roles/serviceusage.serviceUsageAdmin). Learn how to grant roles.

Enable the API

In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

Roles required to select or create a project

Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Go to project selector

Verify that billing is enabled for your Google Cloud project.

Enable the required API, if it is not already enabled.

Roles required to enable APIs

Enable the API

צריך לוודא שיש לכם בפרויקט את התפקיד או התפקידים הבאים: roles/container.admin,roles/iam.serviceAccountAdmin,roles/storage.admin
בדיקת התפקידים
1. נכנסים לדף IAM במסוף Google Cloud .
  כניסה לדף IAM
2. בוחרים את הפרויקט.
3. בעמודה Principal, מחפשים את כל השורות שמזהות אתכם או קבוצה שאתם נכללים בה. כדי לדעת באילו קבוצות אתם נכללים, פנו לאדמין.
4. בודקים את העמודה Role בכל השורות שבהן מצוין או מופיע השם שלכם, כדי לראות אם רשימת התפקידים כוללת את התפקידים הנדרשים.
מתן התפקידים
1. נכנסים לדף IAM במסוף Google Cloud .
  כניסה לדף IAM
2. בוחרים את הפרויקט.
3. לוחצים על Grant access.
4. בשדה New principals, מזינים את מזהה המשתמש. ‫ בדרך כלל מזהה המשתמש הוא כתובת האימייל של חשבון Google.
  ‫
5. לוחצים על Select a role ומחפשים את התפקיד.
6. כדי לתת עוד תפקידים, לוחצים על Add another role ומוסיפים אותם.
7. לוחצים על Save.

מוודאים שיש לכם מספיק מכסת שימוש בשבבי TPU Trillium ‏ (v6e)‎. במדריך הזה משתמשים בהגדרת מאגר צמתים שדורשת 16 שבבים ומופעים על פי דרישה.
מוודאים שיש לכם מאגר Docker. אם אין לכם מאגר, אתם יכולים ליצור מאגר רגיל ב-Artifact Registry.

הכנת הסביבה

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

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

במסוף Google Cloud , מפעילים סשן של Cloud Shell ולוחצים על Activate Cloud Shell. הפעולה הזו מפעילה סשן בחלונית התחתונה של מסוף Google Cloud .
מגדירים את משתני הסביבה שמוגדרים כברירת מחדל:
```
gcloud config set project PROJECT_ID
gcloud config set billing/quota_project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export CLUSTER_NAME=CLUSTER_NAME
export REGION=CONTROL_PLANE_LOCATION
export ZONE=ZONE
export GCS_BUCKET_NAME=BUCKET_NAME
```
מחליפים את הערכים הבאים:
- ‫PROJECT_ID: מזהה הפרויקט ב- Google Cloud .
- ‫CLUSTER_NAME: השם של אשכול GKE.
- ‫CONTROL_PLANE_LOCATION: האזור ב-Compute Engine שבו נמצאים אשכול GKE וצומתי TPU. האזור צריך להכיל אזורים שבהם זמינים סוגי מכונות TPU Trillium ‏ (v6e).
- ‫ZONE: אזור בתוך אזור CONTROL_PLANE_LOCATION שבחרתם שבו זמינים סוגי מכונות TPU Trillium‏ (v6e). כדי להציג את רשימת האזורים שבהם זמינים מכשירי TPU Trillium ‏ (v6e), מריצים את הפקודה הבאה:
```
  gcloud compute accelerator-types list --filter="name~ct6e" --format="value(zone)"
```
- ‫BUCKET_NAME: השם של הקטגוריה של Cloud Storage שמכילה את נתוני האימון.

משכפלים את המאגר לדוגמה:

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples.git
cd kubernetes-engine-samples

מנווטים לספריית העבודה:
```
cd ai-ml/llm-training-jax-tpu-gemma3
```

יצירה והגדרה של Google Cloud משאבים

בקטע הזה יוצרים ומגדירים Google Cloud משאבים.

יצירת אשכול GKE

אפשר לכוונן מודל שפה גדול (LLM) ב-TPU באשכול GKE Autopilot או באשכול רגיל. מומלץ להשתמש באשכול Autopilot כדי ליהנות מחוויית Kubernetes מנוהלת באופן מלא. כדי לבחור את מצב הפעולה של GKE שהכי מתאים לעומסי העבודה שלכם, אפשר לעיין במאמר בחירת מצב פעולה של GKE.

טייס אוטומטי

יוצרים אשכול GKE Autopilot שמשתמש ב איחוד זהויות של עומסי עבודה ל-GKE ומופעל בו Cloud Storage FUSE.

gcloud container clusters create-auto ${CLUSTER_NAME} \
    --location=${REGION}

יצירת האשכול עשויה להימשך כמה דקות.

רגילה

יוצרים אשכול GKE Standard אזורי שמשתמש באיחוד זהויות של עומסי עבודה ל-GKE ומופעל בו Cloud Storage FUSE.

gcloud container clusters create ${CLUSTER_NAME} \
    --enable-ip-alias \
    --addons GcsFuseCsiDriver \
    --machine-type=n2-standard-4 \
    --num-nodes=2 \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${REGION}

יצירת האשכול עשויה להימשך כמה דקות.

יוצרים מאגר צמתים עם מארח יחיד:

gcloud container node-pools create jax-tpu-nodepool \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct6e-standard-1t \
    --num-nodes=1 \
    --location=${REGION} \
    --node-locations=${ZONE} \
    --workload-metadata=GKE_METADATA

‫GKE יוצר מאגר צמתים של TPU Trillium עם טופולוגיה של 1x1 וצומת אחד. הדגל --workload-metadata=GKE_METADATA מגדיר את מאגר הצמתים לשימוש בשרת המטא-נתונים של GKE.

התקנת JobSet

מגדירים את kubectl כדי לתקשר עם האשכול:

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}

מתקינים את הגרסה האחרונה של JobSet:
```
kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml
```
מחליפים את JOBSET_VERSION בגרסה העדכנית של JobSet. לדוגמה, v0.11.0.

מאמתים את ההתקנה של JobSet:

kubectl get pods -n jobset-system

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

NAME                                         READY   STATUS    RESTARTS   AGE
jobset-controller-manager-6c56668494-l4dhc   1/1     Running   0          4m45s

יכול להיות שתצטרכו להוסיף עוד צמתים אם JobSet ממתין למשאבים.

הגדרת Cloud Storage FUSE

כדי לבצע התאמה עדינה של מודל שפה גדול (LLM), צריך לספק נתונים לאימון. במדריך הזה השתמשנו במערך הנתונים TinyStories מ-Hugging Face. קבוצת הנתונים הזו מכילה סיפורים קצרים שנוצרו באופן סינתטי על ידי GPT-3.5 ו-GPT-4, והם משתמשים באוצר מילים מוגבל.

בקטע הזה מוסבר איך מגדירים את Cloud Storage FUSE לקריאת נתונים מקטגוריה של Cloud Storage.

מורידים את מערך הנתונים:

wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStories-train.txt?download=true -O TinyStories-train.txt

מעלים את הנתונים לקטגוריה חדשה של Cloud Storage:

gcloud storage buckets create gs://${GCS_BUCKET_NAME} \
    --location=${REGION} \
    --enable-hierarchical-namespace \
    --uniform-bucket-level-access
gcloud storage cp TinyStories-train.txt gs://${GCS_BUCKET_NAME}

כדי לאפשר לעומס העבודה לקרוא נתונים באמצעות Cloud Storage FUSE, צריך ליצור Kubernetes ServiceAccount ‏ (jaxserviceaccount) במרחב השמות default:
```
kubectl create serviceaccount jaxserviceaccount --namespace default
```
מקצים לחשבון השירות jaxserviceaccount ב-Kubernetes את התפקיד Storage Object Viewer‏ (roles/storage.objectViewer) בקטגוריה של Cloud Storage:
```
gcloud storage buckets add-iam-policy-binding gs://${GCS_BUCKET_NAME} \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/jaxserviceaccount" \
    --role="roles/storage.objectViewer" \
    --condition=None
```
לכל Pod שפועל במרחב השמות default של אשכול GKE שמשתמש ב-ServiceAccount‏ jaxserviceaccount יש עכשיו גישת קריאה לאובייקטים שמאוחסנים בקטגוריה gs://${GCS_BUCKET_NAME}. הגישה הזו נדרשת כדי שמשימת הכוונון העדין תוכל לגשת למערך הנתונים באמצעות Cloud Storage FUSE.

יצירת סקריפט לכוונון עדין

בקטע הזה נבדוק את סקריפט האימון שמבצע פעולת כוונון עדין במודל Gemma 3. הסקריפט הזה משתמש ב-Gemma3Tokenizer.

כדאי לעיין בסקריפט הבא של Gemma3LLMTrain.py כוונון עדין:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import grain.python as pygrain
import jax
import jax.numpy as jnp
import optax
import pandas as pd
import time
import argparse

from dataclasses import dataclass
from functools import partial
from gemma import gm
from flax.training import train_state
from jax.sharding import Mesh, PartitionSpec, NamedSharding

jax.distributed.initialize()
print("Global device count:", jax.device_count())
print("jax version:", jax.__version__)

tokenizer = gm.text.Gemma3Tokenizer()
num_epochs = 1
learning_rate = 2e-5

@dataclass
class TextDataset:
    data: list
    maxlen: int
    tokenizer: gm.text.Gemma3Tokenizer

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx: int):
        encoding = self.tokenizer.encode(self.data[idx])[:self.maxlen]  # Tokenize and truncate
        return encoding + [0] * (self.maxlen - len(encoding))  # Pad to maxlen

def load_and_preprocess_data(file_path, batch_size, maxlen, datacount, tokenizer):

    with open(file_path, 'r') as f:
      text = f.read()

    stories = text.split('<|endoftext|>')
    stories = [story for story in stories if story.strip()][:datacount]
    df = pd.DataFrame({'text': stories})
    data = df['text'].dropna().tolist()
    dataset = TextDataset(data, maxlen, tokenizer)

    sampler = pygrain.IndexSampler(
        len(dataset),
        shuffle=False,
        seed=42,
        shard_options=pygrain.NoSharding(),
        num_epochs=num_epochs,
    )

    dataloader = pygrain.DataLoader(
        data_source=dataset,
        sampler=sampler,
        operations=[pygrain.Batch(batch_size=batch_size, drop_remainder=True)],
    )

    return dataloader

def generate_text(model, params, tokenizer, prompt):
    sampler = gm.text.Sampler(
        model=model,
        params=params,
        tokenizer=tokenizer,
    )
    print("Generating response for: " + prompt)
    out = sampler.sample(prompt, max_new_tokens=32)
    print("Reponse: \n" + out + "\n")
    return out

prep_target_batch = jax.vmap(lambda tokens: jnp.concatenate((tokens[1:], jnp.array([0]))))

@partial(jax.jit, donate_argnums=(0,))
def train_step(state, batch):
    """Performs one supervised fine-tuning step."""

    def loss_fn(params):
        # Run the forward pass. The model returns logits.
        logits = state.apply_fn({'params': params}, batch[0]).logits

        # Calculate the cross-entropy loss.
        loss = optax.softmax_cross_entropy_with_integer_labels(
            logits=logits, labels=batch[1]
        ).mean()

        return loss

    # Compute gradients
    grad_fn = jax.value_and_grad(loss_fn)
    loss, grads = grad_fn(state.params)

    # Update the model state
    state = state.apply_gradients(grads=grads)

    metrics = {'loss': loss}
    return state, metrics

def train_model(state, text_dl, num_epochs, sharding):
    batchCount = 0
    start_time = time.time()
    for epoch in range(num_epochs):
        start_time = time.time()
        for batch in text_dl:
            if len(batch) % len(jax.devices()) != 0:
              continue  # skip the remaining elements
            input_batch = jnp.array(jnp.array(batch).T)
            target_batch = prep_target_batch(input_batch)
            state, metrics = train_step(state, jax.device_put((input_batch, target_batch), sharding))

            if batchCount % 10 == 0:
                print(f"Loss after batch {batchCount}: {metrics['loss']}")
            batchCount += 1

    end_time = time.time()
    print(f"Completed training model. Total time for training {end_time - start_time} seconds \n")
    return state

def run_training(maxlen, batch_size, datacount):
    print(f"Batch size: {batch_size}, Max length: {maxlen}, Data count: {datacount}")
    #Load the training data
    tiny_stories_dl = load_and_preprocess_data('/data/TinyStories-train.txt', batch_size, maxlen, datacount, tokenizer)
    # Get the Gemma3 model
    model = gm.nn.Gemma3_270M()
    # Load the pretrained parameters
    params = gm.ckpts.load_params(gm.ckpts.CheckpointPath.GEMMA3_270M_PT)
    # Create an optimizer
    optimizer = optax.adamw(learning_rate=learning_rate)
    # Define sharding for data parallel training
    mesh = Mesh(jax.devices(), ('batch',))
    sharding = NamedSharding(mesh, PartitionSpec('batch', None))

    # Testing out current state of the model
    test_prompt = "Once upon a time, there was a girl named Amy."
    generate_text(model, params, tokenizer, test_prompt)

    state = train_state.TrainState.create(
        apply_fn=model.apply,
        params=params,
        tx=optimizer
    )

    # Perform post training
    print("Start training model")
    state = train_model(state, tiny_stories_dl, num_epochs, sharding)

    # Final text generation
    generate_text(model, state.params, tokenizer, test_prompt)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train Gemma model with custom parameters.')
    parser.add_argument('--maxlen', type=int, default=256, help='Maximum sequence length')
    parser.add_argument('--batch_size', type=int, default=128, help='Batch size')
    parser.add_argument('--datacount', type=int, default=296000, help='Number of data samples to use')
    args = parser.parse_args()

    run_training(maxlen=args.maxlen, batch_size=args.batch_size, datacount=args.datacount)

בסקריפט הזה, התנאים הבאים חלים:

Gemma3Tokenizer ממיר נתוני טקסט לטוקנים שהמודל יכול לעבד.
הפונקציה load_and_preprocess_data קוראת את נתוני האימון מקובץ, מפצלת אותם לסיפורים נפרדים ומשתמשת ב-tokenizer כדי להמיר את הטקסט לרצפים מרופדים של טוקנים.
הפונקציה generate_text מקבלת את המודל, את הפרמטרים שלו ואת ההנחיה ליצירת טקסט.
הפונקציה train_step מגדירה איטרציה אחת של אימון שכוללת את המעבר קדימה, חישוב ההפסד (באמצעות אנטרופיה צולבת), חישוב הגרדיאנט ועדכוני הפרמטרים.
הפונקציה train_model מבצעת איטרציה במערך הנתונים למספר מוגדר של תקופות, ומפעילה את הפונקציה train_step לכל אצווה.
הפונקציה run_training מתזמנת את כל התהליך של טעינת נתונים, אתחול של מודל Gemma 3 ‏ (Gemma3_270M) ואופטימיזציה, טעינה של פרמטרים שאומנו מראש, הגדרה של חלוקת נתונים לעיבוד מקביל, הפעלה של יצירת טקסט לבדיקה, ביצוע של לולאת האימון וביצוע של יצירת טקסט סופית כדי להדגים את ההשפעה של כוונון עדין.
הסקריפט משתמש בספריית argparse כדי לקבל ארגומנטים בשורת הפקודה עבור הפרמטרים maxlen, ‏batch_size ו-datacount.

אחרי שבדקתם את סקריפט הכוונון העדין, אתם יכולים להכניס אותו לקונטיינר כדי להריץ אותו ב-GKE.

העברת סקריפט הכוונון המדויק למאגר

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

פותחים את Dockerfile באותה ספרייה שבה נמצא הקובץ Gemma3LLMTrain.py:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

FROM us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.7.2-rev1
RUN apt-get update && apt-get install -y wget && rm -rf /var/lib/apt/lists/*

RUN pip install --upgrade pip
RUN pip install gemma grain

WORKDIR /app

# Copy your training script into the container
COPY Gemma3LLMTrain.py .

קובץ Dockerfile הזה מתקין את יחסי התלות הנדרשים ומעתיק את הקובץ Gemma3LLMTrain.py לקונטיינר.

יוצרים את קובץ האימג' של Docker ומעבירים אותו בדחיפה למאגר אימג':

export REPOSITORY=REPOSITORY_NAME
export IMAGE_NAME="jax-gemma3-training"
export IMAGE_TAG="latest"
export DOCKERFILE_PATH="./Dockerfile"
export IMAGE_URI="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}/${IMAGE_NAME}:${IMAGE_TAG}"

docker build -t "${IMAGE_URI}" -f "${DOCKERFILE_PATH}" .
gcloud auth configure-docker "${REGION}-docker.pkg.dev" -q
docker push "${IMAGE_URI}"

מחליפים את REPOSITORY_NAME בשם המאגר שלכם ב-Artifact Registry.

מוסיפים קישורי תפקידים לחשבון השירות:

export PROJECT_NUMBER=$(gcloud projects describe $PROJECT_ID --format 'get(projectNumber)')
gcloud artifacts repositories add-iam-policy-binding ${REPOSITORY} \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --member="serviceAccount:${PROJECT_NUMBER}-compute@developer.gserviceaccount.com" \
    --role="roles/artifactregistry.reader"

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

פריסת משימת הכוונון העדין של ה-LLM

בקטע הזה מוסבר איך לפרוס את משימת הכוונון העדין של ה-LLM לאשכול GKE.

פותחים את קובץ המניפסט training_singlehost.yaml:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: batch/v1
kind: Job
metadata:
  name: jax-gemma3-train-singlehost
spec:
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
    spec:
      serviceAccountName: jaxserviceaccout
      containers:
      - name: training-container
        image: ${IMAGE_URI}
        imagePullPolicy: "Always"
        command: ["python", "Gemma3LLMTrain.py", "--maxlen", "256", "--batch_size", "64", "--datacount", "355120"]
        resources:
          limits:
            google.com/tpu: 1
        volumeMounts:
        - name: gcs-fuse-csi-ephemeral
          mountPath: /data
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/gke-tpu-topology: 1x1
      restartPolicy: Never
      volumes:
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${GCS_BUCKET_NAME}
            mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"
  backoffLimit: 1

החלת המניפסט:

envsubst < training_singlehost.yaml | kubectl apply -f -

‫GKE יוצר Job שמפעיל Pod בצומת TPU Trillium ‏ (v6e). ה-Pod הזה מריץ את סקריפט הכוונון העדין של Python, שמקבל גישה לנתוני הכוונון העדין מקטגוריית Cloud Storage שצוינה, שנטענה בנתיב /data באמצעות Cloud Storage FUSE. הסקריפט מבצע כוונון עדין של מודל Gemma.

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

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

הצגת ההתקדמות של הכוונון העדין

מציגים את רשימת ה-Pods:
```
# Find the Pods
kubectl get pods
```

עוקבים אחרי פלט היומן:

kubectl logs -f pods/POD_NAME

מחליפים את POD_NAME בשם ה-Pod.

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

Global device count: 1
Batch size: 128, Max length: 256, Data count: 96000
I1028 00:12:55.925999 1387 google_auth_provider.cc:181] Running on GCE, using service account ...
Generating response for: Once upon a time, there was a girl named Amy.
Response:
Amy lived in a small house. The house was in a big field. Amy liked to play in the big field. She
Start training model
Loss after batch 0: 10.25
Loss after batch 10: 4.3125
.
.
.
Loss after batch 740: 1.41406
Completed training model. Total time for training 294.6791355609894 seconds
Generating response for: Once upon a time, there was a girl named Amy.
Response:
She loved to play with her toys. One day, Amy's mom told her that she had to go to the store to

ניתוח הפלט:
- בשורה Global device count: 1 מצוין מספר ליבות ה-TPU שהיו בשימוש.
- המודל יוצר טקסט סביר לפני הרצת הכוונון העדין, כי הוא נטען מנקודת ביקורת שאומנה מראש.
- הפלט שנוצר אחרי כוונון עדין דומה יותר להתחלה של סיפור קצר, מה שמצביע על כך שהמודל לומד ממערך הנתונים החדש.
- כוונון עדין של מערך הנתונים המלא אמור להניב תוצאות מדויקות עוד יותר.

התבוננות במדדים

כדי לראות את הביצועים של משימת הכוונון העדין, בודקים את מדדי ה-TPU וה-CPU. כדי לראות את מדדי יכולת הצפייה של האשכול, מבצעים את השלבים שמפורטים במאמר הצגת מדדי יכולת הצפייה של האשכול ועומס העבודה.

הגדרות חלופיות לכוונון עדין

בקטע הזה מפורטות הגדרות חלופיות לעומס העבודה של הכוונון העדין.

בחירת מודל

במדריך הזה נעשה שימוש במודל Gemma3_270M, שהוא מודל קטן שמתאים למאגר צמתים של TPU Trillium (v6e) במארח יחיד. כדי לבצע כוונון עדין של מודלים גדולים יותר שדורשים יותר זיכרון ומשאבי מחשוב, אפשר להשתמש בהגדרות של מאגר צמתים מרובה-מארחים או מרובה-פרוסות.

רשימה מלאה של המודלים הזמינים מופיעה בתיעוד של Gemma.

הגדרות של מאגר צמתים

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

בכרטיסיות הבאות מוסבר איך ליצור מאגרי צמתים עם כמה מארחים וכמה פרוסות:

כמה מארחים

ב-Cloud Shell, מריצים את הפקודה הבאה:

gcloud container node-pools create jax-tpu-multihost1 \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct6e-standard-4t \
    --num-nodes=2 \
    --tpu-topology=2x4 \
    --location=${REGION} \
    --node-locations=${ZONE}

‫GKE יוצר מאגר צמתים של TPU Trillium עם טופולוגיה של 2x4 ושני צמתים.

פותחים את training_multihost_jobset.yaml הגדרת המשרה:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: jax-gemma3-train-multihost
spec:
  replicatedJobs:
    - name: trainers
      replicas: 1
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 1
          template:
            metadata:
              annotations:
                gke-gcsfuse/volumes: "true"
            spec:
              serviceAccountName: jaxserviceaccout
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 2x4
                cloud.google.com/gke-nodepool: jax-tpu-multihost1
              containers:
              - name: training-container
                image: ${IMAGE_URI} 
                imagePullPolicy: "Always"
                ports:
                  - containerPort: 8471
                command: ["python", "Gemma3LLMTrain.py", "--maxlen", "256", "--batch_size", "64", "--datacount", "5120"]
                resources:
                  limits:
                    google.com/tpu: 4
                volumeMounts:
                - name: gcs-fuse-csi-ephemeral
                  mountPath: /data
              volumes:
                - name: gcs-fuse-csi-ephemeral
                  csi:
                    driver: gcsfuse.csi.storage.gke.io
                    volumeAttributes:
                      bucketName: ${GCS_BUCKET_NAME}
                      mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"

פורסים את משימת הכוונון העדין:

envsubst < training_multihost_jobset.yaml | kubectl apply -f -

Multislice

ב-Cloud Shell, מריצים את הפקודה הבאה:

gcloud container node-pools create jax-tpu-multihost1 \
  --cluster=${CLUSTER_NAME} \
  --machine-type=ct6e-standard-4t \
  --num-nodes=2 \
  --tpu-topology=2x4 \
  --location=${REGION} \
  --node-locations=${ZONE}

gcloud container node-pools create jax-tpu-multihost2 \
  --cluster=${CLUSTER_NAME} \
  --machine-type=ct6e-standard-4t \
  --num-nodes=2 \
  --tpu-topology=2x4 \
  --location=${REGION} \
  --node-locations=${ZONE}

‫GKE יוצר שני מאגרי צמתים של TPU Trillium. לכל מאגר צמתים יש 2x4 טופולוגיה ושני צמתים.

פותחים את training_multislice_jobset.yaml הגדרת המשרה:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: jax-gemma3-train-multislice
spec:
  replicatedJobs:
    - name: trainers
      replicas: 2
      template:
        spec:
          parallelism: 2
          completions: 2
          backoffLimit: 1
          template:
            metadata:
              annotations:
                gke-gcsfuse/volumes: "true"
            spec:
              serviceAccountName: jaxserviceaccout
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 2x4
              containers:
              - name: training-container
                image: ${IMAGE_URI}
                imagePullPolicy: "Always"
                ports:
                  - containerPort: 8471
                command: ["python", "Gemma3LLMTrain.py", "--maxlen", "256", "--batch_size", "64", "--datacount", "5120"]
                resources:
                  limits:
                    google.com/tpu: 4
                volumeMounts:
                - name: gcs-fuse-csi-ephemeral
                  mountPath: /data
              volumes:
                - name: gcs-fuse-csi-ephemeral
                  csi:
                    driver: gcsfuse.csi.storage.gke.io
                    volumeAttributes:
                      bucketName: ${GCS_BUCKET_NAME}
                      mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:100"

פורסים את משימת הכוונון העדין:

envsubst < training_multislice_jobset.yaml | kubectl apply -f -

ניתוח ביצועים ואופטימיזציה

כדי לנתח את הביצועים של כוונון עדין של למידת מכונה ולבצע אופטימיזציה שלהם, אפשר להשתמש ב-XProf. ‫XProf הוא חבילת כלים שיוצרת פרופיל של עומסי עבודה של למידת מכונה שנבנו באמצעות JAX, ‏ TensorFlow או PyTorch/XLA, ובודקת אותם. בעזרת XProf אפשר לראות את עקבות הביצוע, את השימוש בזיכרון ונתונים אחרים, ולבצע כוונון עדין של המודלים ושל הגדרות האימון כדי לשפר את היעילות ולקצר את זמן האימון.

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

מתקינים את חבילת xprof. משנים את סקריפט ההדרכה כדי להפעיל את שרת XProf.
משנים את מניפסט העבודה של Kubernetes כך שיכלול טעינת נפח ליומני XProf.
נותנים לחשבון השירות של Kubernetes הרשאות לכתוב יומנים של XProf לקטגוריה של Cloud Storage.
מריצים את XProf בתוך ה-Pod ומגדירים העברת יציאות כדי לגשת ללוח הבקרה של XProf.

התקנת חבילת XProf

עוברים לספרייה שמכילה את הדוגמאות של XProf:
```
  cd ai-ml/llm-training-jax-tpu-gemma3/xprof-enabled
```

יוצרים את קובץ האימג' של Docker ומעבירים אותו בדחיפה למאגר אימג':

export REPOSITORY=REPOSITORY_NAME
export IMAGE_NAME="jax-gemma3-training-xp"
export IMAGE_TAG="latest"
export DOCKERFILE_PATH="./Dockerfile"
export IMAGE_URI="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}/${IMAGE_NAME}:${IMAGE_TAG}"

docker build -t "${IMAGE_URI}" -f "${DOCKERFILE_PATH}" .
gcloud auth configure-docker "${REGION}-docker.pkg.dev" -q
docker push "${IMAGE_URI}"

מחליפים את REPOSITORY_NAME בשם המאגר שלכם ב-Artifact Registry.

מריצים את הסקריפט Dockerfile:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

FROM us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:jax0.7.2-rev1
RUN apt-get update && apt-get install -y wget && rm -rf /var/lib/apt/lists/*

RUN pip install --upgrade pip
RUN pip install gemma grain equinox
RUN pip install xprof

WORKDIR /app

# Copy your training script into the container
COPY Gemma3LLMTrain.py .

קובץ ה-Dockerfile הזה מתקין את יחסי התלות של XProf.

מעתיקים את סקריפט הכוונון העדין אל מאגר התגים.

בקטע הזה, יוצרים ומחילים מניפסט של Kubernetes Job שכולל את נקודות הצירוף (mount) של נפחי האחסון שנדרשות ליומני XProf.

פותחים את training_singlehost.yaml הגדרת המשרה:

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: batch/v1
kind: Job
metadata:
  name: jax-gemma3-train-singlehost
spec:
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
    spec:
      serviceAccountName: jaxserviceaccout
      containers:
      - name: training-container
        image: ${IMAGE_URI}
        imagePullPolicy: "Always"
        command: ["python", "Gemma3LLMTrain.py", "--maxlen", "256", "--batch_size", "64", "--datacount", "851200"]
        resources:
          limits:
            google.com/tpu: 1
        volumeMounts:
        - name: gcs-fuse-csi-ephemeral
          mountPath: /data
        - name: gcs-fuse-csi-ephemeral2
          mountPath: /xprof
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
        cloud.google.com/gke-tpu-topology: 1x1
      restartPolicy: Never
      volumes:
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${GCS_BUCKET_NAME}
            mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"
      - name: gcs-fuse-csi-ephemeral2
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${XPROF_GCS_BUCKET_NAME}
            mountOptions: "implicit-dirs,file-cache:enable-parallel-downloads:true,file-cache:parallel-downloads-per-file:100,file-cache:max-parallel-downloads:-1,file-cache:download-chunk-size-mb:10,file-cache:max-size-mb:-1"
  backoffLimit: 1

החלת המניפסט:

envsubst < training_singlehost.yaml | kubectl apply -f -

מתן הרשאות לחשבון השירות של Kubernetes לכתוב יומנים של XProf

כדי לאפשר ל-jaxserviceaccount Kubernetes ServiceAccount לכתוב ולקרוא יומני XProf, צריך להעניק לו את התפקיד Storage Object User‏ (roles/storage.objectUser) בקטגוריה:

export XPROF_GCS_BUCKET_NAME="XPROF_GCS_BUCKET_NAME"

gcloud storage buckets add-iam-policy-binding "gs://${XPROF_GCS_BUCKET_NAME}" \
    --member="principal://iam.googleapis.com/projects/${PROJECT_NUMBER}/locations/global/workloadIdentityPools/${PROJECT_ID}.svc.id.goog/subject/ns/default/sa/jaxserviceaccount" \
    --role="roles/storage.objectUser" \
    --project="${PROJECT_ID}" \
    --condition=None

מחליפים את XPROF_GCS_BUCKET_NAME בשם הקטגוריה שרוצים להקצות לה את התפקיד.

מריצים את XProf בתוך ה-Pod:

kubectl exec POD_NAME -c training-container -it -- bash # exec into the container
xprof --port 9001 --logdir /xprof # start xprof

מחליפים את POD_NAME בשם ה-Pod.

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

מגדירים העברה ליציאה אחרת לשרת XProf ב-Pod:
```
kubectl port-forward POD_NAME 9001:9001
```
בסרגל הכתובות של הדפדפן, מזינים את הכתובת הבאה:
```
http://localhost:9001/
```
הכלי XProf Trace Viewer ייפתח.
בחלון TensorBoard, לוחצים על Capture profile (לכידת פרופיל).
בשדה Profile Service URL(s) or TPU name, מזינים localhost:9002.
כדי ללכוד פרטים נוספים, בקטע Host Trace (TraceMe) Level (רמת מעקב אחר המארח (TraceMe)), בוחרים באפשרות verbose (מפורט) ומפעילים את האפשרות Python trace logging (רישום מעקב של Python).
כדי להציג את לוח הבקרה, לוחצים על Capture (צילום).

‫TensorBoard מתעד את הפרופיל ומאפשר לכם לנתח את הביצועים של סקריפט האימון. בתרשים מוצג ציר הזמן של ההפעלה עבור פרופילי הביצועים של TPU ושל CPU:

דוגמה לכלי XProf trace viewer שמציג תרשים של מטריצת ביצועים

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

כוונון עדין בסביבות ייצור

במדריך הזה למדתם איך לבדוק אימון מבוסס JAX בסביבה מבוזרת. כדי לבצע אופטימיזציה של כוונון עדין של מודלים גדולים של שפה (LLM) בסביבת הייצור, צריך להשתמש בספריית Maxtext. אם אתם רוצים להשתמש במודלים של דיפוזיה, אתם יכולים להשתמש בהטמעות של Maxdiffusion.

כדי למזער את אובדן ההתקדמות במהלך תקלה, מומלץ להגדיר יצירת נקודות ביקורת (checkpointing) של עומסי עבודה ארוכי טווח של אימון או כוונון עדין בסביבת ייצור. מידע נוסף על הגדרת נקודות ביקורת רב-שכבתיות זמין במאמר אימון מודלים של למידת מכונה בקנה מידה גדול ב-GKE באמצעות נקודות ביקורת רב-שכבתיות.

הסרת המשאבים

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

מחיקת המשאבים הבודדים

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

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

gcloud container clusters delete ${CLUSTER_NAME} --location=${REGION}
gcloud storage rm --recursive gs://${GCS_BUCKET_NAME}
gcloud artifacts docker images delete ${IMAGE_URI} --delete-tags

אם אתם לא צריכים את הנתונים שנוצרו על ידי XProf, אתם יכולים להסיר את קטגוריית Cloud Storage שבה נעשה שימוש על ידי XProf:
```
gcloud storage rm --recursive gs://${XPROF_GCS_BUCKET_NAME}
```

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

מידע נוסף על TPU ב-GKE
כדאי לעיין במאגר JAX.

איך מבצעים כוונון עדין של מודל LLM באמצעות TPUs ב-GKE עם JAX קל לארגן דפים בעזרת אוספים אפשר לשמור ולסווג תוכן על סמך ההעדפות שלך.

רקע

‏Gemma

TPUs

JAX

מטרות

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

בדיקת התפקידים

מתן התפקידים

הכנת הסביבה

יצירה והגדרה של Google Cloud משאבים

יצירת אשכול GKE

טייס אוטומטי

רגילה

התקנת JobSet

הגדרת Cloud Storage FUSE

יצירת סקריפט לכוונון עדין

העברת סקריפט הכוונון המדויק למאגר

פריסת משימת הכוונון העדין של ה-LLM

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

הצגת ההתקדמות של הכוונון העדין

התבוננות במדדים

הגדרות חלופיות לכוונון עדין

בחירת מודל

הגדרות של מאגר צמתים

כמה מארחים

Multislice

ניתוח ביצועים ואופטימיזציה

התקנת חבילת XProf

מעתיקים את סקריפט הכוונון העדין אל מאגר התגים.

מתן הרשאות לחשבון השירות של Kubernetes לכתוב יומנים של XProf

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

כוונון עדין בסביבות ייצור

הסרת המשאבים

מחיקת המשאבים הבודדים

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

איך מבצעים כוונון עדין של מודל LLM באמצעות TPUs ב-GKE עם JAX