使用 PyTorch 训练 DLC 和 SFT 在 Vertex AI 上微调 Mistral 7B v0.3

设置/配置

首先，您需要在本地机器上安装 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"),
)

在继续定义 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 Config Zoo 中一些预定义的配置文件。

%%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

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

创建 $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 进行训练之前，您需要首先定义成功运行作业所需的配置，即哪些 GPU 能够以 bfloat16 精度微调 mistralai/Mistral-7B-v0.3。

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

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

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

运行 CustomContainerTrainingJob

如前所述，该作业将使用 TRL CLI 在 mistralai/Mistral-7B-v0.3 上以 bfloat16 精度运行监督微调 (SFT)，使用 timdettmers/openassistant-guanaco，这是 OpenAssistant/oasst1 的一个子集，包含约 1 万个样本。

一旦您决定了用于运行作业的资源，您需要相应地定义超参数，以确保所选实例能够运行该作业。您可能需要查看以下一些超参数以避免出现 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` 命令获取。

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

由于 GCS FUSE 用于将存储桶作为目录挂载到正在运行的容器作业中，因此挂载路径遵循 /gcs/ 格式。更多信息请参阅 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）。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 分别定义了中断作业执行或作业创建请求（分配所需资源并开始执行的时间）之前的超时时间（以秒为单位）。
（可选）boot_disk_size 定义了启动磁盘的大小（以 GiB 为单位），为了存储模型权重以及所有中间检查点（如果有），它会增加；否则，它默认为 100GiB，这在某些情况下可能不足。

!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 中的 pipeline）或通过 Hugging Face TGI DLC（因为模型已针对文本生成进行微调）对其进行推理。

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

< > 在 GitHub 上更新