dstack:您的LLM启动平台——从微调到部署,一切简化

Dstack是一个开源项目,它使开发者能够无缝管理虚拟机实例,从而使大型语言模型 (LLM) 的多节点微调变得更加易于访问。通过将 dstack 与Google Cloud Platform (GCP)、Amazon Web Services (AWS)、Microsoft Azure、Oracle Cloud Infrastructure (OCI) 等各种云服务平台结合使用,您可以简化设置和管理分布式训练流程。本指南将引导您在GCP上使用dstack微调 Gemma 7B,并结合Hugging Face alignment-handbook 的最佳实践,然后使用Hugging Face的 Text Generation Inference (TGI) 部署模型以供服务。
注意:本博客文章的实验是在3个节点上进行的,每个节点都配备(2个A10)GPU,并选择Gemma 7B模型作为要微调的基础模型。
在GCP上设置dstack
通过四个简单的步骤,我们可以通过dstack与GCP进行交互。首先,我们需要安装dstack Python包。由于dstack支持多个云提供商,我们可以将范围缩小到GCP。
$ pip install dstack[gcp]
接下来,我们需要在 ~/.dstack/server/config.yml 中配置GCP特定的凭据。以下假设您已获得应用程序默认凭据。此外,您需要通过 gcloud auth application-default login
进行应用程序默认凭据验证,然后才能继续。有关服务帐户凭据,请参阅 dstack官方文档。
projects:
- name: main
backends:
- type: gcp
project_id: <your-gcp-project-id>
creds:
type: default
接下来,我们可以按如下方式启动dstack服务器。
$ dstack server
INFO Applying ~/.dstack/server/config.yml
INFO Configured the main project in ~/.dstack/config.yml
INFO The admin token is ab6e8759-9cd9-4e84-8d47-5b94ac877ebf
INFO The dstack server 0.18.4 is running at http://127.0.0.1:3000
获取管理员令牌、运行IP地址和端口。然后,在您选择的文件夹中初始化您的dstack项目,这样您就一切准备就绪了!从这一点开始的每个步骤都与任何其他基础架构设置相同。
# inside a project folder
$ dstack init
$ dstack config –url http://127.0.0.1:3000 \
--project main \
--token ab6e8759-9cd9-4e84-8d47-5b94ac877ebf
使用dstack微调LLM
让我们深入探讨LLM微调的实际步骤。以下是使用dstack启动微调作业的命令:
$ ACCEL_CONFIG_PATH=fsdp_qlora_full_shard.yaml \
FT_MODEL_CONFIG_PATH=qlora_finetune_config.yaml \
HUGGING_FACE_HUB_TOKEN=xxxx \
WANDB_API_KEY=xxxx \
dstack apply . -f ft.task.dstack.yml
FT_MODEL_CONFIG_PATH、ACCEL_CONFIG_PATH、HUGGING_FACE_HUB_TOKEN 和 WANDB_API_KEY 是在 dstack 运行脚本中定义的环境变量。这些变量最终将在 GCP 上预置的虚拟机中设置为环境变量。dstack apply . 在 GCP 上运行在 ft.task.dstack.yml 中定义的作业,它还会复制当前目录的内容并将其设置为工作目录。
注意:如果这些环境变量在您本地机器的当前终端会话中已定义,则无需明确设置它们。
让我们逐一查看每个 YAML 文件。本博客文章中只重点介绍重要部分,但它们的完整内容可在 此仓库 中找到。首先,我们来看一下 qlora_finetune_config.yaml 和 fsdp_qlora_full_shard.yaml,它们分别定义了我们如何微调 LLM 以及如何利用底层 GPU 基础设施进行微调。
qlora_finetune_config.yaml 是 alignment-handbook 可以理解的关于您如何微调 LLM 的配置。
# Model arguments
model_name_or_path: google/gemma-7b
tokenizer_name_or_path: philschmid/gemma-tokenizer-chatml
torch_dtype: bfloat16
bnb_4bit_quant_storage: bfloat16
# LoRA arguments
load_in_4bit: true
use_peft: true
lora_r: 8
lora_alpha: 16
lora_dropout: 0.05
lora_target_modules:
- q_proj
- k_proj
...
# Data training arguments
dataset_mixer:
chansung/mental_health_counseling_conversations: 1.0
dataset_splits:
- train
- test
...
模型参数
- model_name_or_path: 选择 Google 的 Gemma 7B 作为基础模型
- tokenizer_name_or_path: alignment-handbook 使用所选分词器的 apply_chat_template() 方法。本教程使用 ChatML 模板,而不是 Gemma 7B 的标准对话模板。
- torch_dtype 和 bnb_4bit_quant_storage:如果我们要利用 FSDP+QLoRA 微调方法,这两个值应该定义相同。由于 Gemma 7B 很难适应单个 A10 GPU,本博客文章使用 FSDP+QLoRA 将模型分片到 2 个 A10 GPU 上,同时利用 QLoRA 技术。
LoRA 参数: LoRA 特定的配置。由于本博客文章利用 FSDP+QLoRA 技术,load_in_4bit 设置为 true。其他配置可能因实验而异。
数据训练参数: 我们准备了一个基于 Amod 的心理健康咨询对话数据集。由于 alignment-handbook 只理解 [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}, ….] 形式的数据,这些数据可以通过分词器的 apply_chat_template() 方法进行解释,因此 准备好的数据集 基本上是将原始数据集转换为 apply_chat_template() 兼容格式。
现在,有了微调配置,是时候定义底层基础设施应该如何进行微调作业了,这就是 fsdp_qlora_full_shard.yaml 发挥作用的地方。
compute_environment: LOCAL_MACHINE
distributed_type: FSDP # Use Fully Sharded Data Parallelism
fsdp_config:
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
fsdp_backward_prefetch: BACKWARD_PRE
fsdp_cpu_ram_efficient_loading: true
fsdp_use_orig_params: false
fsdp_offload_params: true
fsdp_sharding_strategy: FULL_SHARD
# ... (other FSDP configurations)
# ... (other configurations)
- distributed_type: FSDP 表示使用 Fully Sharded Data Parallel (FSDP),这是一种支持训练大型模型(否则无法适应单个 GPU)的技术。
- fsdp_config: 这些设置了 FSDP 的操作方式,例如模型的切分方式 (fsdp_sharding_strategy) 以及参数是否卸载到 CPU (fsdp_offload_params)。
通过 distributed_type 的 FSDP 和 fsdp_config 的 fsdp_sharding_strategy 的 FULL_SHARD,模型将在多个节点上的多个 GPU 之间进行分片。如果您想在一个节点中分片模型并在其他节点中拥有相同的分片模型实例,则 fsdp_sharding_strategy 的值应设置为 HYBRID_SHARD。在这种情况下,如果有多个节点,每个节点将拥有自身内部多个 GPU 之间分片的相同模型。这意味着每个节点中的每个分片模型实例将学习给定数据集的不同部分/批次。
还有其他重要变量,例如 machine_rank、num_machines 和 num_processes,但这些变量的值最好在运行时根据您的目标环境注入,因为我们可以轻松切换到不同的基础设施规格。
dstack 的强大之处:轻松配置
除了 fsdp_qlora_full_shard.yaml,我们的多节点设置的核心是 ft.task.dstack.yaml。dstack 简化了分布式训练环境复杂配置的定义。
type: task
python: "3.11"
nodes: 3
env:
- ACCEL_CONFIG_PATH
- FT_MODEL_CONFIG_PATH
- HUGGING_FACE_HUB_TOKEN
- WANDB_API_KEY`
commands:
# ... (setup steps, cloning repo, installing requirements)
- ACCELERATE_LOG_LEVEL=info accelerate launch \
--config_file recipes/custom/accel_config.yaml \
--main_process_ip=$DSTACK_MASTER_NODE_IP \
--main_process_port=8008 \
--machine_rank=$DSTACK_NODE_RANK \
--num_processes=$DSTACK_GPUS_NUM \
--num_machines=$DSTACK_NODES_NUM \
scripts/run_sft.py recipes/custom/config.yaml
ports:
- 6006
resources:
gpu: 1..2
shm_size: 24GB
重点如下:
- 无缝集成: dstack 毫不费力地与 Hugging Face 的开源生态系统集成。特别是,您可以简单地使用 accelerate 库,以及我们在 fsdp_qlora_full_shard.yaml 中定义的配置,就像正常使用一样。
- 自动配置: $DSTACK_MASTER_NODE_IP、$DSTACK_NODE_RANK、$DSTACK_GPUS_NUM 和 $DSTACK_NODES_NUM 变量由 dstack 自动管理,减少了手动设置。
- 资源分配: dstack 可以轻松指定微调作业的节点数和 GPU(gpu:1..2)。因此,对于本博客文章,有三个节点,每个节点都配备 2 个 A10(24GB)GPU。
部署您微调的模型
一旦您的模型经过微调,dstack 就可以轻松地使用 Hugging Face 的 Text Generation Inference (TGI) 框架将其部署到 GCP 上。以下是如何在 dstack 中定义服务以安全地部署您的微调模型的示例:
type: service
image: ghcr.io/huggingface/text-generation-inference:latest`
env:
- HUGGING_FACE_HUB_TOKEN
- MODEL_ID=chansung/mental_health_counseling_merged_v0.1
commands:
- text-generation-launcher \
--max-input-tokens 512 --max-total-tokens 1024 \
--max-batch-prefill-tokens 512 --port 8000
port: 8000
resources:
gpu:
memory: 48GB
# (Optional) Enable the OpenAI-compatible endpoint
model:
format: tgi
type: chat
name: chansung/mental_health_counseling_merged_v0.1
这种方法的主要优点
- 安全HTTPS网关: Dstack简化了通过网关设置安全HTTPS连接的过程,这是生产级模型服务的一个关键方面。
- 为推理优化: TGI框架专为高效的文本生成推理而设计,确保您的模型提供响应迅速且可靠的结果。
此时,您可以通过标准 curl 命令和 Python 的 requests、OpenAI SDK 和 Hugging Face 的 InferenceClient 库与服务实例进行交互。例如,以下代码片段显示了 curl 的示例。
curl \-X POST [https://black-octopus-1.deep-diver-main.sky.dstack.ai/generate](https://black-octopus-1.deep-diver-main.sky.dstack.ai/generate) \\
\-H "Authorization: Bearer \<dstack-token\>" \\
\-H 'Content-Type: application/json' \\
\-d '{"inputs":"I feel bad .....","parameters":{"max\_new\_tokens":128}}'
此外,如果您正在使用 dstack Sky,您可以直接从 dstack 聊天 UI 与部署的模型进行交互。dstack Sky 是一个完全托管的云平台,允许您免费管理自己的云资源。或者,dstack.ai 可以以有竞争力的市场价格提供各种云服务提供商的资源配额。
图1. dstack Sky上的聊天UI
结论
通过遵循本指南中概述的步骤,您已经解锁了一种强大的方法,可以使用 dstack、GCP 和 Hugging Face 生态系统的组合功能来微调和部署 LLM。您现在可以利用 dstack 友好的用户界面有效管理您的 GCP 资源,从而简化设置分布式训练环境的过程,以用于您的 LLM 项目。此外,与 Hugging Face 的 alignment-handbook 和 TGI 框架的集成使您能够无缝微调和部署模型,确保它们针对性能进行优化并为实际应用做好准备。我们鼓励您进一步探索可能性,并尝试不同的模型和配置,以在自然语言处理领域实现您期望的结果。
额外提示
dstack fleet 使您能够在云端和本地环境中预置资源,让您甚至在执行任何任务之前就拥有所需的可用资源。当您需要 dstack 的好处而无需直接访问云资源时,或在任何需要高效的跨云端和本地环境资源管理和预置的情况下,这都特别有用。
dstack volume 允许您为开发环境、任务和服务创建并附加持久卷。卷目前处于实验阶段,并支持 aws、gcp 和 runpod 后端。它们允许您在运行之间持久化数据。您可以定义一个配置文件来创建新卷,或注册一个现有卷。一旦创建了卷,您可以将其附加到开发环境、任务和服务。这使您可以在不同运行之间共享数据。