1. 简介
在本实验中，您将学习如何使用 Google Kubernetes Engine (GKE) 为 Llama 2（一种热门的开源语言模型）构建完整的生产级微调流水线。您将了解架构决策、常见权衡取舍以及反映真实机器学习操作 (MLOps) 工作流的组件。
您将预配一个 GKE 集群，使用 LoRA（低秩自适应）构建容器化训练流水线，并在 GKE 上运行训练作业。
架构概览
以下是我们今天将构建的内容：
该架构包括：
- GKE 集群：管理我们的计算资源
- GPU 节点池：1 个 L4 GPU（竞价型），用于训练
- GCS 存储分区：存储模型和数据集
- Workload Identity：在 K8s 和 GCS 之间实现安全访问
学习内容
- 预配和配置具有针对机器学习工作负载优化的功能的 GKE 集群。
- 使用 Workload Identity 实现从 GKE 到其他 Google Cloud 服务的安全无密钥访问。
- 使用 Docker 构建容器化训练流水线。
- 使用 LoRA 通过参数高效微调 (PEFT) 来高效地微调开源模型。
2. 项目设置
Google 账号
如果您还没有个人 Google 账号，则必须先创建一个 Google 账号。
请使用个人账号，而不是工作账号或学校账号。
登录 Google Cloud 控制台
使用个人 Google 账号登录 Google Cloud 控制台。
创建项目（可选）
如果您没有要用于此实验的当前项目，请在此处创建一个新项目。
3. 打开 Cloud Shell Editor
- 点击此链接可直接前往 Cloud Shell 编辑器
- 如果系统在今天任何时间提示您进行授权，请点击授权继续。
- 如果终端未显示在屏幕底部，请打开它：
- 点击查看
- 点击终端
- 在终端中，使用以下命令设置项目：
gcloud config set project [PROJECT_ID]- 示例：
gcloud config set project lab-project-id-example - 如果您不记得自己的项目 ID，可以使用以下命令列出所有项目 ID：
gcloud projects list
- 示例：
- 您应会看到以下消息：
Updated property [core/project].
4. 启用 API
如需使用 GKE 和其他服务，您需要在 Google Cloud 项目中启用必要的 API。
- 在终端中，启用以下 API：
gcloud services enable container.googleapis.com \ artifactregistry.googleapis.com \ cloudbuild.googleapis.com \ iam.googleapis.com \ compute.googleapis.com \ iamcredentials.googleapis.com \ storage.googleapis.com
API 简介
- Google Kubernetes Engine API (
container.googleapis.com) 可让您创建和管理运行应用的 GKE 集群。 - Artifact Registry API (
artifactregistry.googleapis.com) 提供了一个安全的私有代码库来存储您的容器映像。 - Cloud Build API (
cloudbuild.googleapis.com) 由gcloud builds submit命令用于在云端构建容器映像。 - 借助 IAM API (
iam.googleapis.com)，您可以管理 Google Cloud 资源的访问权限控制和身份。 - Compute Engine API (
compute.googleapis.com) 可提供在 Google 基础设施上运行的安全且可自定义的虚拟机。 - IAM Service Account Credentials API (
iamcredentials.googleapis.com) 允许为服务账号创建短期有效的凭据。 - Cloud Storage API (
storage.googleapis.com) 可让您在云端存储和检索数据，在此处用于模型和数据集存储。
5. 设置项目环境
创建工作目录
- 在终端中，为您的项目创建一个目录，然后进入该目录。
mkdir llama-finetuning cd llama-finetuning
设置环境变量
- 在终端中，创建一个名为
env.sh的文件来存储环境变量。这样可确保您在会话断开连接后轻松重新加载这些标签页。cat <<EOF > env.sh export PROJECT_ID=$(gcloud config get-value project) export CLUSTER_NAME="ml-gke" export GPU_NODE_POOL_NAME="gpu-pool" export MACHINE_TYPE="e2-standard-4" export GPU_MACHINE_TYPE="g2-standard-16" export GPU_TYPE="nvidia-l4" export GPU_COUNT=1 export REGION="asia-southeast1" export NODE_LOCATIONS="asia-southeast1-a,asia-southeast1-b" EOF - 获取文件以将变量加载到当前会话中：
source env.sh
6. 预配 GKE 集群
- 在终端中，创建具有默认节点池的 GKE 集群。此过程大约需要 5 分钟。
gcloud container clusters create $CLUSTER_NAME \ --project=$PROJECT_ID \ --region=$REGION \ --release-channel=rapid \ --machine-type=$MACHINE_TYPE \ --workload-pool=${PROJECT_ID}.svc.id.goog \ --addons=GcsFuseCsiDriver,HttpLoadBalancing \ --enable-image-streaming \ --enable-ip-alias \ --num-nodes=1 \ --enable-autoscaling \ --min-nodes=1 \ --max-nodes=3 - 接下来，向集群添加 GPU 节点池。此节点池将用于训练模型。
gcloud container node-pools create $GPU_NODE_POOL_NAME \ --project=$PROJECT_ID \ --cluster=$CLUSTER_NAME \ --region=$REGION \ --machine-type=$GPU_MACHINE_TYPE \ --accelerator type=$GPU_TYPE,count=$GPU_COUNT,gpu-driver-version=latest \ --ephemeral-storage-local-ssd=count=1 \ --enable-autoscaling \ --enable-image-streaming \ --num-nodes=0 \ --min-nodes=0 \ --max-nodes=1 \ --location-policy=ANY \ --node-taints=nvidia.com/gpu=present:NoSchedule \ --node-locations=$NODE_LOCATIONS \ --spot - 最后，获取新集群的凭据，并验证您是否可以连接到该集群。
gcloud container clusters get-credentials $CLUSTER_NAME --region=$REGION kubectl get nodes
7. 配置 Hugging Face 访问权限
准备好基础设施后，您现在需要为项目提供必要的凭据，以便访问模型和数据。在此任务中，您将首先获取 Hugging Face 令牌。
获取 Hugging Face 令牌
- 如果您没有 Hugging Face 账号，请在新浏览器标签页中前往 huggingface.co/join，然后完成注册流程。
- 注册并登录后，前往 huggingface.co/meta-llama/Llama-2-7b-hf。
- 阅读许可条款，然后点击相应按钮以接受这些条款。
- 前往 Hugging Face 访问令牌页面 (huggingface.co/settings/tokens)。
- 点击新建令牌。
- 对于角色，选择读取。
- 在名称部分，输入一个描述性名称（例如 finetuning-lab）。
- 点击创建令牌。
- 将生成的令牌复制到剪贴板。下一步操作将会用到该地址。
更新环境变量
现在，我们将 Hugging Face 令牌和 GCS 存储分区的名称添加到 env.sh 文件中。将 [your-hf-token] 替换为您刚刚复制的令牌。
- 在终端中，将新变量附加到
env.sh并重新加载它们：cat <<EOF >> env.sh export HF_TOKEN="[your-hf-token]" export BUCKET_NAME="\${PROJECT_ID}-llama-fine-tuning" EOF source env.sh
8. 配置 Workload Identity
接下来，您将设置 Workload Identity，这是允许在 GKE 上运行的应用访问 Google Cloud 服务（而无需管理静态服务账号密钥）的推荐方法。如需了解详情，请参阅 Workload Identity 文档。
- 首先，创建 Google 服务账号 (GSA)。在终端中，运行以下命令：
cat <<EOF >> env.sh export GSA_NAME="llama-fine-tuning" EOF source env.sh gcloud iam service-accounts create $GSA_NAME \ --display-name="Llama Fine-tuning Service Account" - 接下来，创建 GCS 存储分区并向 GSA 授予访问权限：
gcloud storage buckets create gs://$BUCKET_NAME --project=$PROJECT_ID --location=$REGION gcloud storage buckets add-iam-policy-binding gs://$BUCKET_NAME \ --member=serviceAccount:${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role=roles/storage.admin - 现在，创建 Kubernetes 服务账号 (KSA)：
cat <<EOF >> env.sh export KSA_NAME="llama-workload-sa" export NAMESPACE="ml-workloads" EOF source env.sh kubectl create namespace $NAMESPACE kubectl create serviceaccount $KSA_NAME --namespace $NAMESPACE - 最后，在 GSA 和 KSA 之间创建 IAM 政策绑定：
gcloud iam service-accounts add-iam-policy-binding ${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/iam.workloadIdentityUser \ --member "serviceAccount:${PROJECT_ID}.svc.id.goog[${NAMESPACE}/${KSA_NAME}]" kubectl annotate serviceaccount $KSA_NAME --namespace $NAMESPACE \ iam.gke.io/gcp-service-account=${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com
9. 暂存基本模型
在生产环境机器学习流水线中，Llama 2（约 13 GB）等大型模型通常会在训练期间预先暂存到 Cloud Storage 中，而不是在训练期间下载。此方法可提供更好的可靠性、更快的访问速度，并避免网络问题。Google Cloud 在公开 GCS 存储分区中提供了预下载的热门模型版本，您将在本实验中使用这些版本。
- 首先，我们来验证您是否可以访问 Google 提供的 Llama 2 模型：
gcloud storage ls gs://vertex-model-garden-public-us-central1/llama2/llama2-7b-hf/ - 使用
gcloud storage命令将 Llama 2 模型从此公共存储分区复制到您自己的项目存储分区。此转移使用 Google 的高速内部网络，应该只需要一两分钟时间。gcloud storage cp -r gs://vertex-model-garden-public-us-central1/llama2/llama2-7b-hf \ gs://${BUCKET_NAME}/llama2-7b/ - 通过列出存储分区的内容，验证模型文件是否已正确复制。
gcloud storage ls --recursive --long gs://${BUCKET_NAME}/llama2-7b/llama2-7b-hf/
10. 准备训练代码
现在，您将构建用于微调模型的容器化应用。此任务使用 LoRA（低秩自适应），这是一种参数高效微调 (PEFT) 技术，通过仅训练小型“适配器”层而非整个模型，可大幅降低内存要求。
现在，为训练流水线创建 Python 脚本。
- 在终端中，运行以下命令以打开
train.py文件：cloudshell edit train.py - 将以下代码粘贴到
train.py文件中：
#!/usr/bin/env python3
"""Fine-tune Llama 2 with LoRA on American Stories dataset """
import os
import torch
import logging
from pathlib import Path
from datasets import load_dataset, concatenate_datasets
from transformers import (
AutoTokenizer,
AutoModelForCausalLM,
Trainer,
TrainingArguments,
DataCollatorForLanguageModeling
)
from peft import get_peft_model, LoraConfig
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["NCCL_DEBUG"] = "INFO"
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
class SimpleTextDataset(torch.utils.data.Dataset):
def __init__(self, input_ids, attention_mask):
self.input_ids = input_ids
self.attention_mask = attention_mask
def __len__(self):
return len(self.input_ids)
def __getitem__(self, idx):
return {
'input_ids': self.input_ids[idx],
'attention_mask': self.attention_mask[idx],
'labels': self.input_ids[idx].clone()
}
def get_lora_config():
config = {
"r": 16,
"lora_alpha": 32,
"target_modules": [
"q_proj", "k_proj", "v_proj", "o_proj",
"gate_proj", "up_proj", "down_proj"
],
"lora_dropout": 0.05,
"task_type": "CAUSAL_LM",
}
return LoraConfig(**config)
def load_model_and_tokenizer(model_path):
logger.info("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_path)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
logger.info("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
model_path,
torch_dtype=torch.float16,
device_map="auto",
trust_remote_code=True,
use_cache=False
)
return model, tokenizer
def prepare_dataset(tokenizer, max_length=512):
logger.info("Loading American Stories dataset...")
We recommend using o
dataset = load_dataset(
"dell-research-harvard/AmericanStories",
"subset_years",
year_list=["1809", "1810", "1811", "1812", "1813", "1814", "1815"],
trust_remote_code=True
)
all_articles = []
for year_data in dataset.values():
all_articles.extend(year_data["article"])
logger.info(f"Total articles collected: {len(all_articles)}")
batch_size = 1000
all_input_ids = []
all_attention_masks = []
for i in range(0, len(all_articles), batch_size):
batch_articles = all_articles[i:i+batch_size]
logger.info(f"Processing batch {i//batch_size + 1}/{(len(all_articles) + batch_size - 1)//batch_size}")
encodings = tokenizer(
batch_articles,
padding="max_length",
truncation=True,
max_length=max_length,
return_tensors="pt"
)
all_input_ids.append(encodings['input_ids'])
all_attention_masks.append(encodings['attention_mask'])
# Concatenate all batches
input_ids = torch.cat(all_input_ids, dim=0)
attention_mask = torch.cat(all_attention_masks, dim=0)
logger.info(f"Total tokenized examples: {len(input_ids)}")
# Create simple dataset
dataset = SimpleTextDataset(input_ids, attention_mask)
return dataset
def train_model(model, tokenizer, train_dataset, output_dir):
logger.info(f"Train dataset size: {len(train_dataset)}")
n_gpus = torch.cuda.device_count()
logger.info(f"Available GPUs: {n_gpus}")
# For multi-GPU, we can increase batch size
per_device_batch_size = 2 if n_gpus > 1 else 1
gradient_accumulation_steps = 2 if n_gpus > 1 else 4
# Training for 250 steps
max_steps = 250
training_args = TrainingArguments(
output_dir=output_dir,
max_steps=max_steps,
per_device_train_batch_size=per_device_batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
learning_rate=2e-4,
warmup_steps=20,
fp16=True,
gradient_checkpointing=True,
logging_steps=10,
evaluation_strategy="no",
save_strategy="no",
optim="adamw_torch",
ddp_find_unused_parameters=False,
dataloader_num_workers=0,
remove_unused_columns=False,
report_to=[],
disable_tqdm=False,
logging_first_step=True,
)
# Create trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
tokenizer=tokenizer,
data_collator=DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False,
pad_to_multiple_of=8
)
)
# Train
logger.info(f"Starting training on {n_gpus} GPU(s)...")
logger.info(f"Training for {max_steps} steps - approximately {(max_steps * 2.5 / 60):.1f} minutes")
try:
trainer.train()
logger.info("Training completed successfully!")
except Exception as e:
logger.error(f"Training error: {e}")
logger.info("Attempting to save model despite error...")
logger.info("Saving model...")
try:
trainer.save_model(output_dir)
tokenizer.save_pretrained(output_dir)
logger.info(f"Model saved successfully to {output_dir}!")
if not output_dir.startswith("/gcs-mount"):
logger.info("Copying artifacts to GCS bucket...")
gcs_target = "/gcs-mount/llama2-7b-american-stories"
os.makedirs(gcs_target, exist_ok=True)
return_code = os.system(f"cp -r {output_dir}/* {gcs_target}/")
if return_code != 0:
raise RuntimeError(f"Failed to copy model to GCS: cp command returned {return_code}")
logger.info(f"Copied to {gcs_target}")
except Exception as e:
logger.error(f"Error saving model or copying to GCS: {e}")
raise
def run_inference(model, tokenizer):
logger.info("Running inference test...")
prompt = "The year was 1812, and the"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
model.eval()
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=50,
do_sample=True,
temperature=0.7,
pad_token_id=tokenizer.pad_token_id
)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
logger.info("-" * 50)
logger.info(f"Input Prompt: {prompt}")
logger.info(f"Generated Text: {generated_text}")
logger.info("-" * 50)
def main():
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
logger.info(f"GPU {i}: {torch.cuda.get_device_name(i)}")
model_path = os.getenv('MODEL_PATH', '/gcs-mount/llama2-7b/llama2-7b-hf')
output_path = os.getenv('OUTPUT_PATH', '/gcs-mount/llama2-7b-american-stories')
# Load model and tokenizer
model, tokenizer = load_model_and_tokenizer(model_path)
model.enable_input_require_grads()
# Apply LoRA
logger.info("Applying LoRA configuration...")
lora_config = get_lora_config()
model = get_peft_model(model, lora_config)
model.train()
# Prepare dataset
train_dataset = prepare_dataset(tokenizer)
# Train
train_model(model, tokenizer, train_dataset, output_path)
# Run Inference
run_inference(model, tokenizer)
logger.info("Training and inference complete!")
if __name__ == "__main__":
main()
11. 了解训练代码
train.py 脚本用于编排微调过程。我们来了解一下其主要组成部分。
配置
该脚本使用 LoraConfig 定义低秩自适应设置。LoRA 可大幅减少可训练参数的数量，让您能够在较小的 GPU 上微调大型模型。
def get_lora_config():
config = {
"r": 16,
"lora_alpha": 32,
"target_modules": ["q_proj", "k_proj", "v_proj", "o_proj", ...],
"lora_dropout": 0.05,
"task_type": "CAUSAL_LM",
}
return LoraConfig(**config)
准备数据集
prepare_dataset 函数会加载“American Stories”数据集，并将其处理为标记化的块。它使用自定义 SimpleTextDataset 来高效处理输入张量。
def prepare_dataset(tokenizer, max_length=512):
dataset = load_dataset("dell-research-harvard/AmericanStories", ...)
# ... tokenization logic ...
return SimpleTextDataset(input_ids, attention_mask)
火车
train_model 函数使用针对此工作负载优化的特定实参设置 Trainer。关键参数包括：
gradient_accumulation_steps：有助于在不增加内存用量的情况下模拟更大的批次大小。fp16=True：使用混合精度训练来减少内存用量并提高速度。gradient_checkpointing=True：通过在反向传播期间重新计算激活值（而不是存储它们）来节省内存。optim="adamw_torch"：使用 PyTorch 中的标准 AdamW 优化器实现。
training_args = TrainingArguments(
per_device_train_batch_size=per_device_batch_size,
gradient_accumulation_steps=gradient_accumulation_steps,
fp16=True,
gradient_checkpointing=True,
optim="adamw_torch",
)
推理
run_inference 函数使用示例提示对微调后的模型执行快速测试。它可确保模型处于评估模式，并生成文本来验证适配器是否正常运行。
def run_inference(model, tokenizer):
prompt = "The year was 1812, and the"
# ... generation logic ...
logger.info(f"Generated Text: {generated_text}")
12. 将应用容器化
现在，使用 Docker 构建训练容器映像，并将其推送到 Google Artifact Registry。
- 在终端中，运行以下命令以打开
Dockerfile文件：cloudshell edit Dockerfile - 将以下代码粘贴到
Dockerfile文件中：
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime
WORKDIR /app
# Install required packages
RUN pip install --no-cache-dir \
transformers==4.46.0 \
datasets==3.1.0 \
pyarrow==15.0.0 \
peft==0.13.2 \
accelerate==1.1.0 \
tensorboard==2.18.0 \
nvidia-ml-py==12.535.161 \
scipy==1.13.1
# Copy training scripts
COPY train.py /app/
# Run training
CMD ["python", "train.py"]
构建和推送容器
- 创建 Artifact Registry 仓库：
gcloud artifacts repositories create gke-finetune \ --repository-format=docker \ --location=$REGION \ --description="Docker repository for Llama fine-tuning" - 使用 Cloud Build 构建并推送映像：
gcloud builds submit --tag ${REGION}-docker.pkg.dev/${PROJECT_ID}/gke-finetune/llama-trainer:latest .
13. 部署微调作业
- 创建 Kubernetes 作业清单以启动微调作业。在终端中，运行以下命令：
cloudshell edit training_job.yaml - 将以下代码粘贴到
training_job.yaml文件中：
apiVersion: batch/v1
kind: Job
metadata:
name: llama-fine-tuning
namespace: ml-workloads
spec:
template:
metadata:
annotations:
gke-gcsfuse/volumes: "true"
gke-gcsfuse/memory-limit: "4Gi"
spec:
serviceAccountName: llama-workload-sa
restartPolicy: OnFailure
tolerations:
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
- key: cloud.google.com/gke-spot
operator: Exists
effect: NoSchedule
nodeSelector:
cloud.google.com/gke-accelerator: nvidia-l4
containers:
- name: training
image: ${REGION}-docker.pkg.dev/${PROJECT_ID}/gke-finetune/llama-trainer:latest
env:
- name: MODEL_PATH
value: "/gcs-mount/llama2-7b/llama2-7b-hf"
- name: OUTPUT_PATH
value: "/tmp/llama2-7b-american-stories"
- name: NCCL_DEBUG
value: "INFO"
resources:
requests:
nvidia.com/gpu: 1
cpu: "8"
memory: "32Gi"
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: gcs-fuse
mountPath: /gcs-mount
- name: shm
mountPath: /dev/shm
volumes:
- name: gcs-fuse
csi:
driver: gcsfuse.csi.storage.gke.io
volumeAttributes:
bucketName: ${BUCKET_NAME}
mountOptions: "implicit-dirs"
- name: shm
emptyDir:
medium: Memory
sizeLimit: 32Gi
- 最后，应用 Kubernetes 作业清单，以在 GKE 集群上启动微调作业。
envsubst < training_job.yaml | kubectl apply -f -
14. 监控训练作业
您可以在 Google Cloud 控制台中监控训练作业的进度。
- 前往 Kubernetes Engine > 工作负载页面。
查看 GKE 工作负载 - 点击
llama-fine-tuning作业以查看其详细信息。 - 默认情况下，系统会显示详细信息标签页。您可以在资源部分中查看 GPU 利用率指标。
- 点击日志标签页以查看训练日志。您应该会看到训练进度，包括损失和学习速率。
15. 清理
为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除包含这些资源的项目，或者保留项目但删除各个资源。
删除 GKE 集群
gcloud container clusters delete $CLUSTER_NAME --region $REGION --quiet
删除 Artifact Registry 代码库
gcloud artifacts repositories delete gke-finetune --location $REGION --quiet
删除 GCS 存储分区
gcloud storage rm -r gs://${BUCKET_NAME}
16. 恭喜！
您已在 GKE 上成功对开源 LLM 进行了微调！
回顾
您此实验中，您将执行以下操作：
- 预配了具有 GPU 加速功能的 GKE 集群。
- 配置了 Workload Identity，以便安全地访问 Google Cloud 服务。
- 使用 Docker 和 Artifact Registry 将 PyTorch 训练作业容器化。
- 部署了使用 LoRA 的微调作业，以使 Llama 2 适应新数据集。
后续步骤
- 详细了解 GKE 上的 AI。
- 探索 Vertex AI Model Garden。
- 加入 Google Cloud 社区，与其他开发者互动交流。
Google Cloud 学习路线
本实验是可用于生产用途的 AI 与 Google Cloud 学习路线的组成部分。探索完整课程，弥合从原型设计到生产的差距。
使用 #ProductionReadyAI 标签分享您的进度。