使用 Vertex AI 上的 SFT 对 Mistral 7B v0.3 进行微调，并使用 PyTorch 训练 DLC

设置/配置

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

然后，您还需要安装 google-cloud-aiplatform Python SDK，该 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"),
)

在继续定义 CustomContainerTrainingJob 之前，您需要定义在运行 trl sft 命令时要使用的 accelerate 配置文件，这是因为您处于多 GPU 环境中，否则将使用默认配置，并且可能无法在运行多个 GPU 上的微调作业时获得最佳性能。

您需要通过在本地创建以下 deepspeed.yaml 文件来定义 DeepSpeed Zero3 配置，该文件包含将在分布式环境中使用多个 GPU 运行 SFT 微调所需的配置。以下配置文件中定义的一些值是

mixed_precision=bf16，因为微调将在 bfloat16 中进行
num_processes=4，因为微调将在 4 个 A100 GPU 上运行
num_machines=1 和 same_network=true，因为 GPU 位于同一个单一实例中

请注意，DeepSpeed Zero3 已被选定为 accelerate 的分布式配置，但任何其他配置都可以使用，并且可以通过 accelerate config 命令进行配置，该命令将提示不同的配置；或者只是浏览 Accelerate 配置库中的一些预定义配置文件。

%%writefile "./assets/deepspeed.yaml"
compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  offload_optimizer_device: none
  offload_param_device: none
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

现在您需要定义一个 CustomContainerTrainingJob，该作业在 Hugging Face PyTorch 训练 DLC 上运行，需要运行以下顺序步骤

创建 $HF_HOME/accelerate 路径（如果尚未创建），因为 accelerate 配置将转储到该路径。
将 deepspeed.yaml 配置文件的内容写入缓存中的 default_config.yaml 名称（因为这是 accelerate 的默认路径，即用于微调作业的配置）。
添加 trl sft 命令，捕获运行作业时将提供的参数。

CustomContainerTrainingJob 将覆盖提供的容器 URI 中提供的默认 ENTRYPOINT，因此如果 ENTRYPOINT 已经适合接收参数，那么就没有必要定义自定义 command。

job = aiplatform.CustomContainerTrainingJob(
    display_name="trl-full-sft",
    container_uri=os.getenv("CONTAINER_URI"),
    command=[
        "sh",
        "-c",
        " && ".join(
            (
                "mkdir -p $HF_HOME/accelerate",
                f"echo \"{open('./assets/deepspeed.yaml').read()}\" > $HF_HOME/accelerate/default_config.yaml",
                'exec trl sft "$@"',
            )
        ),
        "--",
    ],
)

定义 CustomContainerTrainingJob 需求

在通过 Hugging Face PyTorch 训练 DLC 继续进行 CustomContainerTrainingJob 之前，您需要先定义作业成功运行所需的配置，即哪种 GPU 能够在 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'

运行 CustomContainerTrainingJob

如前所述，该任务将在 mistralai/Mistral-7B-v0.3 之上使用 TRL CLI 运行监督微调 (SFT)，使用 timdettmers/openassistant-guanaco 以 bfloat16 格式进行，它是一个来自 OpenAssistant/oasst1 的子集，包含约 10k 个样本。

确定要使用哪些资源运行任务后，需要根据这些资源定义超参数，以确保选定的实例能够运行该任务。为了避免出现 OOM 错误，您可能需要查看一些超参数：

优化器：默认情况下会使用 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 命令中列出。

有关 TRL CLI 的更多信息，请查看 https://huggingface.co/docs/trl/en/clis。

由于 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",
    # TRAINER
    "--bf16",
    "--max_seq_length=1024",
    "--per_device_train_batch_size=2",
    "--gradient_accumulation_steps=4",
    "--gradient_checkpointing",
    "--gradient_checkpointing_use_reentrant",
    "--learning_rate=0.00002",
    "--lr_scheduler_type=cosine",
    "--optim=adamw_bnb_8bit",
    "--num_train_epochs=1",
    "--logging_steps=10",
    "--do_eval",
    "--eval_steps=100",
    "--save_strategy=epoch",
    "--report_to=none",
    f"--output_dir={os.getenv('BUCKET_URI').replace('gs://', '/gcs/')}/Mistral-7B-v0.3-SFT-Guanaco",
    "--overwrite_output_dir",
    "--seed=42",
    "--log_level=info",
]

然后，您需要在 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）；它们之间是相互关联的，因此您需要选择支持您正在使用的加速器的实例，反之亦然。有关不同实例的更多信息，请查看 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 分别定义在中断作业执行或作业创建请求（分配所需资源并开始执行所需的时间）之前的超时时间（以秒为单位）。
(可选) boot_disk_size 定义引导磁盘的大小（以 GiB 为单位），它不仅用于存储模型权重，还用于存储所有中间检查点（如果有）；否则，它默认为 100 GiB，在某些情况下可能不足够。

!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="a2-highgpu-4g",
    accelerator_type="NVIDIA_TESLA_A100",
    accelerator_count=4,
    base_output_dir=f"{os.getenv('BUCKET_URI')}/Mistral-7B-v0.3-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)
    boot_disk_size_gb=250,
)

Pipeline created in Vertex AI

Vertex AI Pipeline successfully completed

Vertex AI Pipeline logs

Vertex AI Run

GCS Bucket with uploaded artifacts

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

📍 在 GitHub 上找到完整示例此处！

< > 更新在 GitHub 上