使用 PyTorch 训练 DLC 和 SFT 在 GKE 上微调 Gemma 2B

Gemma 是一个轻量级、最先进的开源模型系列，它基于与创建 Gemini 模型相同的研发和技术构建而成，由 Google DeepMind 和 Google 的其他团队开发。TRL 是一个由 Hugging Face 开发的全栈库，用于微调和对齐大型语言模型 (LLM)。此外，Google Kubernetes Engine (GKE) 是 Google Cloud 中一个完全托管的 Kubernetes 服务，可以使用 GCP 的基础设施大规模部署和运行容器化应用程序。

此示例展示了如何在 GKE 集群的多 GPU 设置中，通过监督微调 (SFT) 使用 TRL 对 Gemma 2B 进行完整微调。

设置/配置

首先，您需要在本地机器上安装 gcloud 和 kubectl，它们分别是 Google Cloud 和 Kubernetes 的命令行工具，分别用于与 GCP 和 GKE 集群交互。

• 要安装 gcloud，请按照Cloud SDK 文档 - 安装 gcloud CLI中的说明进行操作。
• 要安装 kubectl，请按照Kubernetes 文档 - 安装工具中的说明进行操作。

可选地，为了简化本教程中命令的使用，您需要为 GCP 设置以下环境变量

export PROJECT_ID=your-project-id
export LOCATION=your-location
export CLUSTER_NAME=your-cluster-name

然后，您需要登录到您的 GCP 帐户并将项目 ID 设置为您要用于部署 GKE 集群的项目 ID。

gcloud auth login
gcloud auth application-default login  # For local development
gcloud config set project $PROJECT_ID

登录后，您需要在 GCP 中启用必要的服务 API，例如 Google Kubernetes Engine API、Google Container Registry API 和 Google Container File System API，这些 API 对于部署 GKE 集群和用于 TGI 的 Hugging Face DLC 是必需的。

gcloud services enable container.googleapis.com
gcloud services enable containerregistry.googleapis.com
gcloud services enable containerfilesystem.googleapis.com

此外，要将 kubectl 与 GKE 集群凭据一起使用，您还需要安装 gke-gcloud-auth-plugin，可以使用 gcloud 如下安装

gcloud components install gke-gcloud-auth-plugin

创建 GKE 集群

一切设置完成后，您可以继续创建 GKE 集群和节点池，在本例中将是一个单 GPU 节点，以便使用 GPU 加速器进行高性能推理，也遵循 TGI 基于其内部 GPU 优化建议。

要部署 GKE 集群，将使用“Autopilot”模式，因为它是大多数工作负载的推荐模式，因为底层基础设施由 Google 管理。或者，您也可以使用“Standard”模式。

gcloud container clusters create-auto $CLUSTER_NAME \
    --project=$PROJECT_ID \
    --location=$LOCATION \
    --release-channel=stable \
    --cluster-version=1.28 \
    --no-autoprovisioning-enable-insecure-kubelet-readonly-port

gcloud container get-server-config \
    --flatten="channels" \
    --filter="channels.channel=STABLE" \
    --format="yaml(channels.channel,channels.defaultVersion)" \
    --location=$LOCATION

创建 GKE 集群后，您可以使用以下命令获取通过 kubectl 访问它的凭据

gcloud container clusters get-credentials $CLUSTER_NAME --location=$LOCATION

配置 GCS 的 IAM

在 GKE 集群上运行 Hugging Face PyTorch DLC 的微调作业以进行训练之前，您需要为 GCS 存储桶设置 IAM 权限，以便 GKE 集群中的 Pod 可以访问该存储桶，该存储桶将挂载到正在运行的容器中并用于写入生成的工件，以便这些工件自动上传到 GCS 存储桶。为此，您需要在 GKE 集群中创建一个命名空间和一个服务帐户，然后为 GCS 存储桶设置 IAM 权限。

为方便起见，由于将在后续步骤中引用命名空间和服务帐户，因此将设置环境变量 NAMESPACE 和 SERVICE_ACCOUNT。

export NAMESPACE=hf-gke-namespace
export SERVICE_ACCOUNT=hf-gke-service-account

然后，您可以在 GKE 集群中创建命名空间和服务帐户，从而启用在该命名空间中为 Pod 创建 IAM 权限，以便在使用该服务帐户时访问 GCS 存储桶。

kubectl create namespace $NAMESPACE
kubectl create serviceaccount $SERVICE_ACCOUNT --namespace $NAMESPACE

然后，您需要按如下方式将 IAM 策略绑定添加到存储桶

gcloud storage buckets add-iam-policy-binding \
    gs://$BUCKET_NAME \
    --member "principal://iam.googleapis.com/projects/$(gcloud projects describe $PROJECT_ID --format="value(projectNumber)")/locations/global/workloadIdentityPools/$PROJECT_ID.svc.id.goog/subject/ns/$NAMESPACE/sa/$SERVICE_ACCOUNT" \
    --role "roles/storage.objectUser"

可选：在 GKE 中设置密钥

由于 google/gemma-2b 是一个受限模型，您需要通过 kubectl 设置一个包含 Hugging Face Hub 令牌的 Kubernetes 密钥。

要为 Hugging Face Hub 生成自定义令牌，您可以按照 https://huggingface.co/docs/hub/en/security-tokens 中的说明进行操作；推荐的方式是安装 huggingface_hub Python SDK，如下所示

pip install --upgrade --quiet huggingface_hub

然后使用生成的令牌登录，该令牌对受限/私有模型具有读取权限

huggingface-cli login

最后，您可以使用 huggingface_hub Python SDK 检索令牌，如下所示，使用生成的 Hugging Face Hub 令牌创建 Kubernetes 密钥

kubectl create secret generic hf-secret \
    --from-literal=hf_token=$(python -c "from huggingface_hub import get_token; print(get_token())") \
    --dry-run=client -o yaml \
    --namespace $NAMESPACE | kubectl apply -f -

或者，您也可以直接设置令牌，如下所示

kubectl create secret generic hf-secret \
    --from-literal=hf_token=hf_*** \
    --dry-run=client -o yaml \
    --namespace $NAMESPACE | kubectl apply -f -

有关如何在 GKE 集群中设置 Kubernetes 密钥的更多信息，请访问 https://cloud.google.com/secret-manager/docs/secret-manager-managed-csi-component。

定义作业配置

在通过 Hugging Face PyTorch DLC 进行批处理作业的 Kubernetes 部署之前，您需要首先定义作业成功运行所需的配置，即哪个 GPU 能够使用 bfloat16 对 google/gemma-2b 进行微调。

粗略估算，您可以假设使用半精度微调模型所需的 GPU VRAM 量大约是模型大小的四倍（有关详细信息，请参阅 Eleuther AI - Transformer Math 101）。

或者，如果您的模型已上传到 Hugging Face Hub，您可以在社区空间 Vokturz/can-it-run-llm 中查看这些数字，它会根据要微调的模型和可用硬件为您计算这些数字。

运行作业

现在，您可以通过 job.yaml 配置文件中的 kubectl 在 GKE 集群上运行 Hugging Face PyTorch DLC 的 Kubernetes 作业，该文件包含运行 trl sft 命令的作业规范，该命令由 TRL CLI 提供，用于使用 timdettmers/openassistant-guanaco 对 google/gemma-2b 进行 bfloat16 的 SFT 全微调，timdettmers/openassistant-guanaco 是 OpenAssistant/oasst1 的一个子集，包含大约 10k 个样本，在 4 个 A100 40GiB GPU 上运行，并将生成的工件存储到 /data 下的卷挂载中，该卷挂载链接到一个 GCS 存储桶。

kubectl apply -f job.yaml

此外，您可以使用 kubectl 流式传输作业日志，如下所示

kubectl logs -f job/trl-full-sft --container trl-container --namespace $NAMESPACE

最后，作业完成后，Pod 将缩减到 0，并且工件将在作业中挂载的 GCS 存储桶中可见。

删除 GKE 集群

最后，微调作业完成后，您可以安全地删除 GKE 集群以避免产生不必要的费用。

gcloud container clusters delete $CLUSTER_NAME --location=$LOCATION

或者，您可能决定在作业完成后继续运行 GKE 集群，因为使用 GKE 自动驾驶仪模式部署的默认 GKE 集群只运行一个 e2-small 实例。