使用 PyTorch 训练 DLC 在 Vertex AI 上对 Gemma 2B 进行 SFT + LoRA 微调

设置/配置

首先，您需要在本地机器上安装 gcloud，这是 Google Cloud 的命令行工具，请按照Cloud SDK 文档 - 安装 gcloud CLI中的说明进行操作。

然后，您还需要安装 google-cloud-aiplatform Python SDK，它用于以编程方式创建 Vertex AI 模型、注册模型、创建端点并在 Vertex AI 上部署模型。

!pip install --upgrade --quiet google-cloud-aiplatform

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

%env PROJECT_ID=your-project-id
%env LOCATION=your-location
%env BUCKET_URI=gs://hf-vertex-pipelines
%env CONTAINER_URI=us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-pytorch-training-cu121.2-3.transformers.4-42.ubuntu2204.py310

然后，您需要登录到您的 GCP 帐户并将项目 ID 设置为您要在 Vertex AI 上注册和部署模型的项目 ID。

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

登录后，您需要在 GCP 中启用必要的服务 API，例如 Vertex AI API、Compute Engine API 和 Google Container Registry 相关的 API。

!gcloud services enable aiplatform.googleapis.com
!gcloud services enable compute.googleapis.com
!gcloud services enable container.googleapis.com
!gcloud services enable containerregistry.googleapis.com
!gcloud services enable containerfilesystem.googleapis.com

可选：在 GCS 中创建存储桶

如果您已经有存储桶，则可以使用现有的存储桶来存储微调工件，请随时跳过此步骤并进入下一步。

由于 Vertex AI 作业将生成工件，因此您需要指定一个 Google Cloud Storage (GCS) 存储桶来转储这些工件。因此，您需要使用以下 gcloud storage buckets create 子命令创建 GCS 存储桶

!gcloud storage buckets create $BUCKET_URI --project $PROJECT_ID --location=$LOCATION --default-storage-class=STANDARD --uniform-bucket-level-access

准备 CustomContainerTrainingJob

配置好环境并创建 GCS 存储桶（如果适用）后，您可以继续定义 CustomContainerTrainingJob，它是在 Vertex AI 上运行容器的标准容器作业，该容器是用于训练的 Hugging Face PyTorch DLC。

import os
from google.cloud import aiplatform

aiplatform.init(
    project=os.getenv("PROJECT_ID"),
    location=os.getenv("LOCATION"),
    staging_bucket=os.getenv("BUCKET_URI"),
)

现在，您需要定义一个在 Hugging Face PyTorch DLC 上运行的 CustomContainerTrainingJob 进行训练，该训练需要设置 trl sft 命令来捕获每次作业运行时提供的参数。

CustomContainerTrainingJob 将覆盖提供的容器 URI 中提供的默认 ENTRYPOINT，因此，如果 ENTRYPOINT 已经准备好接收参数，则无需定义自定义 command。

job = aiplatform.CustomContainerTrainingJob(
    display_name="trl-lora-sft",
    container_uri=os.getenv("CONTAINER_URI"),
    command=[
        "sh",
        "-c",
        'exec trl sft "$@"',
        "--",
    ],
)

定义 CustomContainerTrainingJob 要求

在通过 Hugging Face PyTorch DLC 进行 CustomContainerTrainingJob 训练之前，您需要首先定义作业成功运行所需的配置，即哪个 GPU 能够使用 LoRA 适配器以 bfloat16 格式对 mistralai/Mistral-7B-v0.3 进行微调。

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

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

'Vokturz/can-it-run-llm' for 'mistralai/Mistral-7B-v0.3'

运行自定义容器训练作业

如前所述，该作业将在 mistralai/Mistral-7B-v0.3 的基础上，使用 timdettmers/openassistant-guanaco（来自 OpenAssistant/oasst1 的一个包含约 10k 个样本的子集）以 bfloat16 格式，通过 TRL CLI 运行 LoRA 有监督微调 (SFT)。

一旦您决定使用哪些资源来运行作业，您需要相应地定义超参数，以确保所选实例能够运行该作业。为了避免出现 OOM 错误，您可能需要关注以下一些超参数。

LoRA / QLoRA：您可能需要调整秩，用 r 表示，它定义了每个线性层中可训练参数的比例，这意味着秩越低，内存消耗越少。
优化器：默认情况下将使用 AdamW 优化器，但也可以使用低精度优化器来减少内存，例如 adamw_bnb_8bit（有关 8 位优化器的更多信息，请查看 https://huggingface.co/docs/bitsandbytes/main/en/optimizers）。
批次大小：在遇到 OOM 时，您可以调整批次大小以使用更小的批次，或者调整梯度累积步数来模拟类似的批次大小以更新梯度，但每次提供较少的输入，例如 batch_size=8 和 gradient_accumulation=1 等效于 batch_size=4 和 gradient_accumulation=2。

由于 CustomContainerTrainingJob 定义了命令 trl sft，因此要提供的参数要么在 Python 参考中列出（trl.SFTConfig），要么通过 trl sft --help 命令列出。

在 https://huggingface.co/docs/trl/en/clis 中了解更多关于 TRL CLI 的信息。

需要注意的是，由于 GCS FUSE 用于将存储桶作为运行容器作业中的目录挂载，因此挂载的路径遵循 /gcs/<BUCKET_NAME> 格式。更多信息请参见 https://cloud.google.com/vertex-ai/docs/training/code-requirements。因此，output_dir 需要设置为已挂载的 GCS 存储桶，这意味着 SFTTrainer 写入其中的任何内容都会自动上传到 GCS 存储桶。

args = [
    # MODEL
    "--model_name_or_path=mistralai/Mistral-7B-v0.3",
    "--torch_dtype=bfloat16",
    "--attn_implementation=flash_attention_2",
    # DATASET
    "--dataset_name=timdettmers/openassistant-guanaco",
    "--dataset_text_field=text",
    # PEFT
    "--use_peft",
    "--lora_r=16",
    "--lora_alpha=32",
    "--lora_dropout=0.1",
    "--lora_target_modules=all-linear",
    # TRAINER
    "--bf16",
    "--max_seq_length=1024",
    "--per_device_train_batch_size=2",
    "--gradient_accumulation_steps=8",
    "--gradient_checkpointing",
    "--learning_rate=0.0002",
    "--lr_scheduler_type=cosine",
    "--optim=adamw_bnb_8bit",
    "--num_train_epochs=1",
    "--logging_steps=10",
    "--do_eval",
    "--eval_steps=100",
    "--report_to=none",
    f"--output_dir={os.getenv('BUCKET_URI').replace('gs://', '/gcs/')}/Mistral-7B-v0.3-LoRA-SFT-Guanaco",
    "--overwrite_output_dir",
    "--seed=42",
    "--log_level=debug",
]

然后，您需要在 aiplatform.CustomContainerTrainingJob 上调用 submit 方法，这是一个非阻塞方法，它将安排作业而不会阻塞执行。

下面列出了提供给 submit 方法的参数。

args 定义要提供给 trl sft 命令的参数列表，以 trl sft --arg_1=value ... 的格式提供。
replica_count 定义要运行作业的副本数，对于训练，此值通常设置为 1。
machine_type、accelerator_type 和 accelerator_count 分别定义机器（即 Compute Engine 实例）、加速器（如果有）和加速器数量（范围从 1 到 8）。machine_type 和 accelerator_type 是相互关联的，因此您需要选择一个支持您正在使用的加速器的实例，反之亦然。有关不同实例的更多信息，请参见 Compute Engine 文档 - GPU 机器类型，以及有关 accelerator_type 命名的更多信息，请参见 Vertex AI 文档 - MachineSpec。
base_output_dir 定义从 GCS 存储桶中挂载到运行容器中的基本目录，以最初提供给 aiplatform.init 的 staging_bucket 参数为条件。
(可选) environment_variables 定义在运行容器中定义的环境变量。由于您正在微调一个受限模型，例如 mistralai/Mistral-7B-v0.3，因此您需要设置 HF_TOKEN 环境变量。此外，还定义了一些其他环境变量来设置缓存路径 (HF_HOME) 并确保日志消息正确地流式传输到 Google Cloud Logs Explorer (TRL_USE_RICH、ACCELERATE_LOG_LEVEL、TRANSFORMERS_LOG_LEVEL 和 TQDM_POSITION)。
(可选) timeout 和 create_request_timeout 分别定义在中断作业执行或作业创建请求（分配所需资源并开始执行所需的时间）之前的超时时间（以秒为单位）。

!pip install --upgrade --quiet huggingface_hub

from huggingface_hub import interpreter_login

interpreter_login()

from huggingface_hub import get_token

job.submit(
    args=args,
    replica_count=1,
    machine_type="g2-standard-12",
    accelerator_type="NVIDIA_L4",
    accelerator_count=1,
    base_output_dir=f"{os.getenv('BUCKET_URI')}/Mistral-7B-v0.3-LoRA-SFT-Guanaco",
    environment_variables={
        "HF_HOME": "/root/.cache/huggingface",
        "HF_TOKEN": get_token(),
        "TRL_USE_RICH": "0",
        "ACCELERATE_LOG_LEVEL": "INFO",
        "TRANSFORMERS_LOG_LEVEL": "INFO",
        "TQDM_POSITION": "-1",
    },
    timeout=60 * 60 * 3,  # 3 hours (10800s)
    create_request_timeout=60 * 10,  # 10 minutes (600s)
)

Pipeline created in Vertex AI

Vertex AI Pipeline successfully completed

Vertex AI Pipeline logs

GCS Bucket with uploaded artifacts

最后，您可以将微调后的模型上传到 Hugging Face Hub，或者将其保留在 Google Cloud Storage (GCS) 存储桶中。稍后，您将能够在合并适配器后，在其上运行推理。合并适配器可以通过 Hugging Face PyTorch DLC 进行推理（通过 transformers 中的管道），或者通过 Hugging Face DLC 进行 TGI（因为模型是针对文本生成进行微调的）。

📍 在 GitHub 上查看完整示例此处！

< > GitHub 更新