Perfeziona un LLM utilizzando le TPU su GKE con JAX

Creare e configurare risorse Google Cloud

In questa sezione, creerai e configurerai le risorse Google Cloud .

Crea un cluster GKE

Puoi eseguire il fine tuning di un LLM sulle TPU in un cluster GKE Autopilot o Standard. Ti consigliamo di utilizzare un cluster Autopilot per un'esperienza Kubernetes completamente gestita. Per scegliere la modalità operativa GKE più adatta ai tuoi workload, consulta Scegliere una modalità operativa GKE.

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

Installa JobSet

1. Configura kubectl per comunicare con il cluster:

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

2. Installa l'ultima versione rilasciata di JobSet:

   kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml
   

   Sostituisci JOBSET_VERSION con l'ultima versione rilasciata di JobSet. Ad esempio: v0.11.0.

3. Verifica l'installazione di JobSet:

   kubectl get pods -n jobset-system
   

   L'output è simile al seguente:

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

   Potresti dover aggiungere altri nodi se JobSet è in attesa di risorse.

Configurare Cloud Storage FUSE

Per ottimizzare l'LLM, devi fornire dati di addestramento. In questo tutorial, utilizzerai il set di dati TinyStories di Hugging Face. Questo set di dati contiene racconti generati sinteticamente da GPT-3.5 e GPT-4, che utilizzano un vocabolario limitato.

Questa sezione descrive i passaggi per configurare Cloud Storage FUSE in modo da leggere i dati da un bucket Cloud Storage.

1. Scarica il set di dati:

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

2. Carica i dati in un nuovo bucket 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}
   

3. Per consentire al tuo workload di leggere i dati tramite Cloud Storage FUSE, crea un account di servizio Kubernetes (KSA) e aggiungi le autorizzazioni richieste. Esegui lo script permissionsetup.sh:

   # 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.
   
   #!/bin/bash
   
   # --- Configuration Variables ---
   # Kubernetes Service Account details
   export KSA_NAME="jaxserviceaccout"
   export NAMESPACE="default"
   
   # Google Cloud IAM Service Account details
   export GSA_NAME="<GSA_NAME>"
   # Automatically get the current project ID
   export PROJECT_ID=$(gcloud config get-value project)
   export  GSA_DESCRIPTION="GKE Service Account to read GCS bucket for ${KSA_NAME}"
   
   # GCS Bucket details
   export GCS_BUCKET_NAME="<GCS_BUCKET_NAME>" # <--- IMPORTANT: Update this to your bucket name
   
   # Derived Variables
   export GSA_EMAIL="${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com"
   export WI_MEMBER="serviceAccount:${PROJECT_ID}.svc.id.goog[${NAMESPACE}/${KSA_NAME}]"
   
   # --- Check if PROJECT_ID is set ---
   if [ -z "${PROJECT_ID}" ]; then
     echo "Error: PROJECT_ID is not set. Please set it using 'gcloud config set project YOUR_PROJECT_ID'"
     exit 1
   fi
   
   echo "--- Configuration ---"
   echo "KSA_NAME:      ${KSA_NAME}"
   echo "NAMESPACE:     ${NAMESPACE}"
   echo "GSA_NAME:      ${GSA_NAME}"
   echo "PROJECT_ID:    ${PROJECT_ID}"
   echo "GSA_EMAIL:     ${GSA_EMAIL}"
   echo "GCS_BUCKET_NAME:   ${GCS_BUCKET_NAME}"
   echo "WI_MEMBER:     ${WI_MEMBER}"
   echo "--------------------"
   read -p "Press enter to continue..."
   
   # --- Command Execution ---
   
   echo "[1/5] Creating Google Cloud IAM Service Account (GSA): ${GSA_NAME}"
   gcloud iam service-accounts create "${GSA_NAME}" \
       --project="${PROJECT_ID}" \
       --description="${GSA_DESCRIPTION}" \
       --display-name="${GSA_NAME}"
   
   echo "[2/5] Granting GSA '${GSA_EMAIL}' read access (roles/storage.objectViewer) to bucket 'gs://${GCS_BUCKET_NAME}'"
   gcloud storage buckets add-iam-policy-binding "gs://${GCS_BUCKET_NAME}" \
       --member="serviceAccount:${GSA_EMAIL}" \
       --role="roles/storage.objectViewer" \
       --project="${PROJECT_ID}"
   
   echo "[3/5] Creating Kubernetes Service Account (KSA): ${KSA_NAME} in namespace ${NAMESPACE}"
   kubectl create serviceaccount "${KSA_NAME}" --namespace "${NAMESPACE}"
   
   echo "[4/5] Allowing KSA to impersonate GSA (Workload Identity Binding): ${GSA_EMAIL}"
   gcloud iam service-accounts add-iam-policy-binding "${GSA_EMAIL}" \
       --role roles/iam.workloadIdentityUser \
       --member "${WI_MEMBER}" \
       --project="${PROJECT_ID}"
   
   echo "[5/5] Annotating KSA '${KSA_NAME}' to link with GSA '${GSA_EMAIL}'"
   kubectl annotate serviceaccount "${KSA_NAME}" \
       --namespace "${NAMESPACE}" \
       iam.gke.io/gcp-service-account="${GSA_EMAIL}"
   
   echo "--- Setup Complete ---"
   echo "Pods in namespace '${NAMESPACE}' using serviceAccount '${KSA_NAME}' can now authenticate as '${GSA_EMAIL}' and have read access to 'gs://${GCS_BUCKET_NAME}'."
   

   Dopo aver eseguito questo script, le seguenti risorse vengono configurate nel tuo progettoGoogle Cloud e nel cluster GKE:

   • Nel tuo progetto viene creato un nuovo account di servizio IAM denominato gcs-fuse-sa.
   • Al service account Google (GSA) creato Google Cloud (gcs-fuse-sa) viene concesso il ruolo roles/storage.objectViewer sul bucket Cloud Storage specificato da ${GCS_BUCKET_NAME}. Questa autorizzazione consente al GSA di leggere gli oggetti dal bucket.
   • Nel tuo cluster GKE viene creato un nuovo KSA denominato jaxserviceaccount nello spazio dei nomi default.
   • Il criterio IAM del service account Google è aggiornato per concedere il ruolo roles/iam.workloadIdentityUser al service account Kubernetes. Questa autorizzazione consente al KSA di rappresentare il GSA.

   • Il KSA è annotato per essere collegato al GSA. Questa annotazione indica a GKE quale account di servizio Google deve essere rappresentato dal KSA utilizzando Workload Identity.

     Qualsiasi pod in esecuzione nello spazio dei nomi default del cluster GKE che utilizza il account di servizio jaxserviceaccount ora potrà autenticarsi come GSA gcs-fuse-sa. Questi pod avranno accesso in lettura agli oggetti archiviati nel bucket gs://${GCS_BUCKET_NAME}, il che è essenziale per consentire al job di perfezionamento di accedere al set di dati utilizzando Cloud Storage FUSE.

Crea lo script di ottimizzazione

In questa sezione esplorerai lo script di addestramento che esegue un'operazione di perfezionamento su un modello Gemma 3. Questo script utilizza Gemma3Tokenizer.

Esamina il seguente script di ottimizzazione 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)

In questo script, vale quanto segue:

• Un Gemma3Tokenizer converte i dati di testo in token che il modello può elaborare.
• La funzione load_and_preprocess_data legge i dati di addestramento da un file, li suddivide in singole storie e utilizza il tokenizer per convertire il testo in sequenze di token con padding.
• La funzione generate_text accetta il modello, i relativi parametri e un prompt per generare il testo.
• La funzione train_step definisce una singola iterazione di addestramento che include il passaggio in avanti, il calcolo della perdita (utilizzando l'entropia incrociata), il calcolo del gradiente e gli aggiornamenti dei parametri.
• La funzione train_model scorre il set di dati per un numero specificato di epoche, che chiama la funzione train_step per ogni batch.
• La funzione run_training orchestra l'intero processo di caricamento dei dati, inizializzazione del modello Gemma 3 (Gemma3_270M) e dell'ottimizzatore, caricamento dei parametri preaddestrati, configurazione dello sharding dei dati per l'elaborazione parallela, esecuzione di una generazione di test, esecuzione del ciclo di addestramento ed esecuzione di una generazione di testo finale per dimostrare l'effetto del fine tuning.
• Lo script utilizza la libreria argparse per accettare gli argomenti della riga di comando per i parametri maxlen, batch_size e datacount.

Ora che hai esplorato lo script di perfezionamento, inseriscilo in un container per eseguirlo su GKE.

Containerizza lo script di perfezionamento

Prima di eseguire lo script di perfezionamento in un cluster GKE, devi inserirlo in un container. Questo tutorial utilizza un'immagine AI JAX come immagine di base.

1. Apri Dockerfile nella stessa directory del file 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 .
   

   Questo Dockerfile installa le dipendenze necessarie e copia il file Gemma3LLMTrain.py nel container.

2. Crea l'immagine Docker ed eseguine il push in un repository di immagini:

   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}"
   

   Sostituisci REPOSITORY_NAME con il nome del tuo repository Artifact Registry.

3. Aggiungi associazioni di ruoli al account di servizio:

   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"
   

Con l'immagine nel repository, ora puoi eseguire il deployment del job di perfezionamento in un cluster GKE.

Esegui il deployment del job di ottimizzazione LLM

Questa sezione mostra come eseguire il deployment del job di perfezionamento dell'LLM nel cluster GKE.

1. Apri il file manifest 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

2. Applica il manifest:

   envsubst < training_singlehost.yaml | kubectl apply -f -
   

GKE crea un job che avvia un pod su un nodo TPU Trillium (v6e). Questo pod esegue lo script di perfezionamento Python, che accede ai dati di perfezionamento dal bucket Cloud Storage specificato montato nel percorso /data utilizzando Cloud Storage FUSE. Lo script esegue quindi il fine tuning del modello Gemma.

Monitora il job di addestramento

In questa sezione, monitori l'avanzamento del job di perfezionamento e le relative prestazioni.

Visualizzare l'avanzamento dell'ottimizzazione

1. Elenca i pod:

   # Find the Pods
   kubectl get pods
   

2. Segui l'output del log:

   kubectl logs -f pods/POD_NAME
   

   Sostituisci POD_NAME con il nome del tuo pod.

   L'output è simile al seguente:

   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
   

3. Analizza l'output:

   • La riga Global device count: 1 indica i core TPU utilizzati.
   • Il modello genera un testo ragionevole prima di questa esecuzione di ottimizzazione perché viene caricato da un checkpoint preaddestrato.
   • L'output generato dopo l'ottimizzazione mostra una maggiore somiglianza con l'inizio di un racconto, il che indica che il modello sta apprendendo dal nuovo set di dati.
   • Il perfezionamento sull'intero set di dati dovrebbe produrre risultati ancora più raffinati.

Osservare le metriche

Visualizza il rendimento del job di perfezionamento controllando le metriche TPU e CPU. Per visualizzare le metriche di osservabilità per il cluster, segui i passaggi descritti in Visualizzare le metriche di osservabilità di cluster e carichi di lavoro.

Configurazioni di ottimizzazione alternative

Questa sezione descrive configurazioni alternative per il tuo workload di perfezionamento.

Selezione del modello

Questo tutorial ha utilizzato il modello Gemma3_270M, un modello piccolo che rientra in un pool di nodi TPU Trillium (v6e) single-host. Per i modelli più grandi che richiedono più memoria e calcolo per il perfezionamento, puoi utilizzare configurazioni di pool di nodi multi-host o multislice.

Per un elenco completo dei modelli disponibili, consulta la documentazione di Gemma.

Configurazioni del node pool

Questo tutorial ha utilizzato un pool di nodi single-host. Puoi anche creare node pool di slice TPU multi-host o node pool multislice, a seconda delle tue esigenze.

Le seguenti schede mostrano come creare pool di nodi multihost e multislices: