使用 SFT + LoRA 在 GKE 上通过 PyTorch 训练 DLC 微调 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 的基础设施大规模部署和操作容器化应用程序。

本示例展示了如何使用 TRL 通过监督微调 (SFT) 和低秩适应 (LoRA) 在 GKE 集群的单个 GPU 上微调 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 集群的项目。

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，这些都是部署 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 中设置 Secret

由于 mistralai/Mistral-7B-v0.3 是一个受限模型，您需要通过 kubectl 设置一个包含 Hugging Face Hub 令牌的 Kubernetes secret。

要为 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 secret

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 secret 的更多信息，请访问 https://cloud.google.com/secret-manager/docs/secret-manager-managed-csi-component。

定义作业配置

在通过 Hugging Face PyTorch DLC 在 Kubernetes 中部署训练批处理作业之前，您需要首先定义作业成功运行所需的配置，即哪种 GPU 能够在 bfloat16 模式下使用 LoRA 微调 mistralai/Mistral-7B-v0.3。

粗略计算，您可以假设以半精度微调模型所需的 GPU 显存量约为模型大小的四倍（更多信息请参阅 Eleuther AI - Transformer Math 101）。

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

运行作业

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

git clone https://github.com/huggingface/Google-Cloud-Containers
kubectl apply -f Google-Cloud-Containers/examples/gke/trl-lora-fine-tuning/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 Autopilot 模式部署的默认 GKE 集群只运行一个 e2-small 实例。