使用 PyTorch Training DLC 在 GKE 上通过 SFT + LoRA 微调 Mistral 7B v0.3

Mistral 是 Mistral AI 团队创建的一系列不同大小的模型；Mistral 7B v0.3 LLM 是 Mistral 7B v0.2 的扩展词汇版本。TRL 是 Hugging Face 开发的一个用于微调和对齐大型语言模型 (LLM) 的完整栈库。并且，Google Kubernetes Engine (GKE) 是 Google Cloud 中一个完全托管的 Kubernetes 服务，可用于使用 GCP 的基础设施大规模部署和运行容器化应用程序。

此示例展示了如何在 GKE 集群上的单个 GPU 上，通过监督微调 (SFT) 和低秩适应 (LoRA) 使用 TRL 微调 Mistral 7B v0.3。

设置/配置

首先，您需要在本地机器上安装 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 创建访问 GCS 存储桶的 IAM 权限。

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 中设置密钥

由于mistralai/Mistral-7B-v0.3 是一个受限模型，因此您需要通过 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 能够使用 LoRA 以 bfloat16 格式微调 mistralai/Mistral-7B-v0.3。

粗略计算，您可以假设以半精度微调模型所需的 GPU VRAM 量大约是模型大小的四倍（在Eleuther AI - Transformer Math 101中了解更多信息）。

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

运行作业

现在，您已经可以在 GKE 集群上使用 Hugging Face PyTorch DLC 进行训练，通过 job.yaml 配置文件中的 kubectl 运行 Kubernetes 作业，该文件包含运行 TRL CLI 提供的 trl sft 命令的作业规范，用于使用 timdettmers/openassistant-guanaco（来自 OpenAssistant/oasst1 的一个子集，包含约 10k 个样本）以 bfloat16 格式对 mistralai/Mistral-7B-v0.3 进行 SFT LoRA 微调，在一个 L4 24GiB GPU 上运行，并将生成的工件存储到 /data 下的卷挂载点中，该挂载点链接到一个 GCS 存储桶。

kubectl apply -f job.yaml

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

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

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

删除 GKE 集群

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

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

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