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

创建和配置 Google Cloud 资源

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

创建 GKE 集群

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

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

安装 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,