使用 JAX 在 GKE 上通过 TPU 微调 LLM

本教程介绍如何使用 Google Kubernetes Engine (GKE) 上的张量处理单元 (TPU) 和 JAX 来对大语言模型 (LLM) 进行微调。借助微调，您可以调整基础模型（例如 Gemma 3），使其适合特定领域或任务。此过程通过使用您自己的专业数据集更新模型参数，从而提高模型的精确度和准确度。

如果您在微调 AI/机器学习工作负载时需要利用托管式 Kubernetes 的精细控制、自定义、可伸缩性、弹性、可移植性和成本效益，那么本指南是一个很好的起点。

背景

通过在 GKE 上使用 TPU 和 Jax 对 LLM 进行微调，您可以构建一个可用于生产用途的强大微调解决方案，具备托管式 Kubernetes 的所有优势。

Gemma

Gemma 是一组公开提供的轻量级生成式 AI/机器学习多模态模型（根据开放许可发布）。这些 AI 模型可以在应用、硬件、移动设备或托管服务中运行。Gemma 3 引入了多模态功能，支持视觉语言输入和文本输出。它可处理最多 128,000 个 token 的上下文窗口，并支持 140 多种语言。Gemma 3 还提供改进的数学、推理和聊天功能，包括结构化输出和函数调用。

您可以使用 Gemma 模型生成文本，也可以针对专门任务对这些模型进行调优。

如需了解详情，请参阅 Gemma 文档。

TPU

TPU 是 Google 定制开发的应用专用集成电路 (ASIC)，用于加速使用 TensorFlow、PyTorch 和 JAX 等框架构建的机器学习和 AI 模型。

使用 GKE 中的 TPU 之前，我们建议您完成以下学习路线：

  1. 了解 Cloud TPU 系统架构中的当前 TPU 版本可用性。
  2. 了解 GKE 中的 TPU。

JAX

JAX 是一种高性能机器学习框架，旨在与 TPU 和 GPU 搭配使用。JAX 提供了一个用于构建和训练机器学习模型的 API。

如需了解详情，请参阅 JAX 代码库。

目标

本教程介绍以下步骤：

  1. 根据模型特征创建一个具有推荐 TPU 拓扑的 GKE Autopilot 或 Standard 集群。 在本教程中，您将在单主机节点池上执行微调。
  2. 将数据添加到 Cloud Storage 存储桶，并通过 Cloud Storage FUSE 将其装载到容器。
  3. 在 GKE 上部署 LLM 微调作业。
  4. 监控微调作业并查看日志。

准备工作

  • 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.

    Roles required to enable APIs

    To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

    Enable the API

  • 确保您在项目中拥有以下一个或多个角色： roles/container.admin、roles/iam.serviceAccountAdmin、roles/storage.admin

    检查角色

    1. 在 Google Cloud 控制台中，前往 IAM 页面。

      转到 IAM
    2. 选择项目。

    3. 在主账号列中，找到标识您或您所属群组的所有行。如需了解您属于哪些群组，请与您的管理员联系。

    4. 对于指定或包含您的所有行，请检查角色列以查看角色列表是否包含所需的角色。

    授予角色

    1. 在 Google Cloud 控制台中，前往 IAM 页面。

      转到 IAM
    2. 选择项目。
    3. 点击 授予访问权限。

    4. 在新的主账号字段中，输入您的用户标识符。 这通常是员工身份池中的用户的标识符。如需了解详情，请参阅在 IAM 政策中表示员工池用户，或与您的管理员联系。

    5. 点击选择角色，然后搜索相应角色。
    6. 如需授予其他角色，请点击 添加其他角色，然后添加其他各个角色。
    7. 点击 Save（保存）。
  • 确保您有足够的配额用于 16 个 TPU Trillium (v6e) 芯片。在本教程中，您将使用需要 16 个芯片和按需实例的节点池配置。
  • 确保您拥有 Docker 代码库。如果您没有，请在 Artifact Registry 中创建一个标准代码库。

准备环境

在本教程中，您将使用 Cloud Shell 来管理 Google Cloud上托管的资源。Cloud Shell 中预安装了本教程所需的软件，包括 kubectl 和 Google Cloud CLI。

如需使用 Cloud Shell 设置您的环境，请按照以下步骤操作：

  1. 在 Google Cloud 控制台中，启动 Cloud Shell 会话，然后点击 Cloud Shell 激活图标 激活 Cloud Shell。此操作会在 Google Cloud 控制台的底部窗格中启动会话。

  2. 设置默认环境变量：

    gcloud config set project PROJECT_ID
gcloud config set billing/quota_project PROJECT_ID
export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=CLUSTER_NAME
export REGION=CONTROL_PLANE_LOCATION
export ZONE=ZONE
export GCS_BUCKET_NAME=BUCKET_NAME

    替换以下值：

    • PROJECT_ID：您的 Google Cloud 项目 ID。
    • CLUSTER_NAME：GKE 集群的名称。
    • CONTROL_PLANE_LOCATION：GKE 集群和 TPU 节点所在的 Compute Engine 区域。相应区域必须包含提供 TPU Trillium (v6e) 机器类型的可用区。

    • ZONE：所选 CONTROL_PLANE_LOCATION 区域内可使用 TPU Trillium (v6e) 机器类型的可用区。如需列出提供 TPU Trillium (v6e) TPU 的地区，请运行以下命令：

        gcloud compute accelerator-types list --filter="name~ct6e" --format="value(zone)"

    • BUCKET_NAME：包含训练数据的 Cloud Storage 存储桶的名称。

  3. 克隆示例代码库：

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

  4. 导航到工作目录：

    cd ai-ml/llm-training-jax-tpu-gemma3

创建和配置 Google Cloud 资源

在本部分中，您将创建和配置 Google Cloud 资源。

创建 GKE 集群

您可以在 GKE Autopilot 或 Standard 集群中的 TPU 上对 LLM 进行微调。我们建议您使用 Autopilot 集群获得全托管式 Kubernetes 体验。如需选择最适合您的工作负载的 GKE 操作模式，请参阅选择 GKE 操作模式。

Autopilot

创建使用 适用于 GKE 的工作负载身份联合并已启用 Cloud Storage FUSE 的 GKE Autopilot 集群。

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

集群创建可能需要几分钟的时间。

标准

  1. 创建使用Workload Identity Federation for GKE并已启用 Cloud Storage FUSE 的区域级 GKE Standard 集群。

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

    集群创建可能需要几分钟的时间。

  2. 创建单主机节点池：

    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 会创建一个具有 1x1 拓扑和一个节点的 TPU Trillium 节点池。--workload-metadata=GKE_METADATA 标志将节点池配置为使用 GKE 元数据服务器。

安装 JobSet

  1. 配置 kubectl 以与您的集群通信：

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

  2. 安装最新发布的 JobSet 版本：

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

    将 JOBSET_VERSION 替换为最新发布的 JobSet 版本。例如 v0.11.0。

  3. 验证 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 进行微调，您需要提供训练数据。在本教程中，您将使用 Hugging Face 中的 TinyStories 数据集。此数据集包含由 GPT-3.5 和 GPT-4 合成生成的短篇故事，这些故事使用有限的词汇。

本部分介绍了配置 Cloud Storage FUSE 以从 Cloud Storage 存储桶读取数据的步骤。

  1. 下载数据集：

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

  2. 将数据上传到新的 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. 如需允许工作负载通过 Cloud Storage FUSE 读取数据，请创建 Kubernetes 服务账号 (KSA) 并添加所需权限。运行 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}'."

    运行此脚本后，您的Google Cloud 项目和 GKE 集群中会配置以下资源：

    • 系统会在您的项目中创建一个名为 gcs-fuse-sa 的新 IAM 服务账号。
    • 创建的 Google Cloud 服务账号 (GSA) (gcs-fuse-sa) 会被授予 ${GCS_BUCKET_NAME} 指定的 Cloud Storage 存储桶的 roles/storage.objectViewer 角色。此权限允许 GSA 从存储桶中读取对象。
    • 系统会在 GKE 集群的 default 命名空间中创建一个名为 jaxserviceaccount 的新 KSA。
    • 更新 GSA 的 IAM 政策，以向 KSA 授予 roles/iam.workloadIdentityUser 角色。此权限允许 KSA 模拟 GSA。

    • KSA 已添加注释，可将其与 GSA 相关联。此注解会告知 GKE，KSA 应使用 Workload Identity 模拟哪个 GSA。

      现在，在 GKE 集群的 default 命名空间中使用 jaxserviceaccount 服务账号运行的任何 Pod 都将能够以 gcs-fuse-sa GSA 的身份进行身份验证。这些 Pod 将拥有对存储在 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 将文本数据转换为模型可以处理的 token。
  • load_and_preprocess_data 函数从文件中读取训练数据，将其拆分为各个故事，并使用分词器将文本转换为填充后的词法单元序列。
  • generate_text 函数接受模型、其参数和提示，以生成文本。
  • train_step 函数定义了一次训练迭代，其中包括前向传递、损失计算（使用交叉熵）、梯度计算和参数更新。
  • train_model 函数会按指定的周期数遍历数据集，并针对每个批次调用 train_step 函数。
  • run_training 函数可协调整个流程，以加载数据、初始化 Gemma 3 模型 (Gemma3_270M) 和优化器、加载预训练的参数、设置用于并行处理的数据分片、运行测试生成、执行训练循环，并执行最终的文本生成来演示微调的效果。
  • 该脚本使用 argparse 库来接受 maxlen、batch_size 和 datacount 参数的命令行实参。

现在，您已经探索了微调脚本，接下来将其容器化，以便在 GKE 上运行。

将微调脚本容器化

在 GKE 集群中运行微调脚本之前，您需要将其容器化。本教程使用 JAX AI 映像作为基础映像。

  1. 打开与 Gemma3LLMTrain.py 文件位于同一目录中的 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

WORKDIR /app

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

    此 Dockerfile 会安装必要的依赖项，并将 Gemma3LLMTrain.py 文件复制到容器中。

  2. 构建 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 制品库的名称。

  3. 向服务账号添加角色绑定：

    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." \
    --role="roles/artifactregistry.reader"

将映像放入仓库后，您现在可以将微调作业部署到 GKE 集群中。

部署 LLM 微调作业

本部分介绍了如何将 LLM 微调作业部署到 GKE 集群。

  1. 打开 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. 应用清单：

    envsubst < training_singlehost.yaml | kubectl apply -f -

GKE 会创建一个作业，该作业在 TPU Trillium (v6e) 节点上启动一个 Pod。此 Pod 运行 Python 微调脚本，该脚本使用 Cloud Storage FUSE 从装载在 /data 路径的指定 Cloud Storage 存储桶中访问微调数据。然后，脚本会对 Gemma 模型进行微调。

监控训练作业

在本部分中，您将监控微调作业的进度及其性能。

查看微调进度

  1. 列出 Pod：

    # Find the Pods
kubectl get pods

  2. 按照日志输出操作：

    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

  3. 分析输出内容：

    • Global device count: 1 线表示使用的 TPU 核心数。
    • 在运行此微调之前，模型会生成合理的文本，因为它会从预训练的检查点加载。
    • 微调后生成的输出更像短篇故事的开头，这表明模型正在从新数据集中学习。
    • 在完整数据集上进行微调应能生成更精细的输出。

观察指标

通过检查 TPU 和 CPU 指标，查看微调作业的性能。如需查看集群的可观测性指标，请按照查看集群和工作负载可观测性指标中的步骤操作。

其他微调配置

本部分概述了微调工作负载的替代配置。

模型选择

本教程使用了 Gemma3_270M 模型，这是一个小型模型，可放入单主机 TPU Trillium (v6e) 节点池中。对于需要更多内存和计算资源才能进行微调的较大模型，您可以使用多主机或多切片节点池配置。

如需查看可用模型的完整列表，请参阅 Gemma 文档。

节点池配置

本教程使用了单主机节点池。您还可以根据需要创建多主机 TPU 切片节点池或多切片节点池。

以下标签页展示了如何为多主机和多切片节点池创建节点池：

多主机

  1. 在 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 会创建一个具有 2x4 拓扑和两个节点的 TPU Trillium 节点池。

  2. 打开 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"

  3. 部署微调作业：

    envsubst < training_multihost_jobset.yaml | kubectl apply -f -

多切片

  1. 在 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 拓扑和两个节点。

  2. 打开 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"

  3. 部署微调作业：

    envsubst < training_multislice_jobset.yaml | kubectl apply -f -

性能分析和优化

如需分析和优化机器学习微调的性能，您可以使用 XProf。XProf 是一套工具，用于分析和检查使用 JAX、TensorFlow 或 PyTorch/XLA 构建的机器学习工作负载。通过显示执行轨迹、内存用量和其他数据，XProf 可让您微调模型和训练设置，以提高效率并加快训练速度。

如需使用 XProf 分析微调工作负载的性能，请完成本部分中的以下步骤：

  • 安装 xprof 软件包。 修改训练脚本以启动 XProf 服务器。
  • 修改 Kubernetes 作业清单，以包含 XProf 日志的卷装载。
  • 向服务账号授予将 XProf 日志写入 Cloud Storage 存储桶的权限。
  • 在 Pod 内运行 XProf，并设置端口转发以访问 XProf 信息中心。

安装 XProf 软件包

  1. 导航到包含 XProf 样本的目录：

      cd ai-ml/llm-training-jax-tpu-gemma3/xprof-enabled

  2. 构建 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 制品库的名称。

  3. 运行 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 作业清单，其中包含 XProf 日志所需的卷装载。

  1. 打开 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

  2. 应用清单：

    envsubst < training_singlehost.yaml | kubectl apply -f -

向服务账号授予写入 XProf 日志的权限

  1. 如需使服务账号能够写入和读取，请添加 "roles/storage.objectUser" 角色：

    export GSA_NAME="GSA_NAME" # Same as used in initial setup

# Automatically get the current project ID
export PROJECT_ID=$(gcloud config get-value project)

# Cloud Storage Bucket details
export XPROF_GCS_BUCKET_NAME="XPROF_GCS_BUCKET_NAME"

# Derived Variables
export GSA_EMAIL="${GSA_NAME}@${PROJECT_ID}.iam."

gcloud storage buckets add-iam-policy-binding "gs://${XPROF_GCS_BUCKET_NAME}" \
    --member="serviceAccount:${GSA_EMAIL}" \
    --role="roles/storage.objectUser" \
    --project="${PROJECT_ID}"

    替换以下内容：

    • GSA_NAME：要向其授予角色的 Google 服务账号的名称。
    • XPROF_GCS_BUCKET_NAME：要向其授予角色的存储桶的名称。

  2. 在 Pod 中运行 XProf：

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

    将 POD_NAME 替换为您的 Pod 名称。

访问 XProf 信息中心

  1. 设置到 Pod 中 XProf 服务器的端口转发：

    kubectl port-forward POD_NAME 9001:9001

  2. 在浏览器的地址栏中，输入以下内容：

    http://localhost:9001/

    XProf Trace Viewer 随即打开。

  3. 在 TensorBoard 窗口中，点击捕获性能剖析文件。

  4. 在配置文件服务网址或 TPU 名称字段中，输入 localhost:9002。

  5. 如需捕获更多详细信息，请在主机跟踪记录 (TraceMe) 级别中选择 verbose 并启用 Python 跟踪记录日志记录。

  6. 如需查看信息中心，请点击捕获。

    TensorBoard 会捕获性能分析数据，并让您分析训练脚本的性能。该图显示了 TPU 和 CPU 性能配置的执行时间线：

显示性能矩阵图的 XProf 轨迹查看器示例

如需了解更多用于分析训练工作负载性能的性能分析选项，请参阅有关计算性能分析的 JAX 文档。

在生产环境中进行微调

本教程介绍了如何在分布式环境中测试基于 JAX 的训练。如需在生产环境中优化 LLM 微调，请使用 Maxtext 库。如果您对扩散模型感兴趣，请使用 Maxdiffusion 实现。

对于生产环境中长时间运行的训练或微调工作负载，请设置工作负载检查点，以最大限度地减少故障期间的进度损失。如需详细了解如何设置多层级检查点，请参阅在 GKE 上使用多层级检查点机制训练大规模机器学习模型。

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。

逐个删除资源

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留该项目但运行以下命令来删除各个资源：

  1. 删除您在本教程中创建的资源：

    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

  2. 如果您不需要 XProf 生成的数据，请移除 XProf 使用的 Cloud Storage 存储桶：

    gcloud storage rm --recursive gs://${XPROF_GCS_BUCKET_NAME}

后续步骤