DeepSpeed

DeepSpeed 是一个专为大规模分布式训练设计的库，旨在高效训练拥有数十亿参数的大型模型。其核心是 Zero Redundancy Optimizer（ZeRO），它会将优化器状态（ZeRO-1）、梯度（ZeRO-2）和参数（ZeRO-3）分片到不同的数据并行进程中。这大幅降低了内存占用，使您能够将训练规模扩展到数十亿参数的模型。为了进一步提升内存效率，ZeRO-Offload 通过在优化阶段利用 CPU 资源，减少了 GPU 的计算和内存开销。

这两项功能均受 🤗 Accelerate 支持，并且可以与 🤗 PEFT 配合使用。

与 bitsandbytes 量化 + LoRA 的兼容性

下表总结了 PEFT 的 LoRA、bitsandbytes 库以及 DeepSpeed Zero 各阶段在微调方面的兼容性。DeepSpeed Zero-1 和 2 在推理阶段不会产生任何影响，因为阶段 1 对优化器状态进行分片，阶段 2 对优化器状态和梯度进行分片：

DeepSpeed 阶段	是否兼容？
Zero-1	🟢
Zero-2	🟢
Zero-3	🟢

对于 DeepSpeed Stage 3 + QLoRA，请参阅下文中的使用 PEFT QLoRA 和 DeepSpeed ZeRO3 在多 GPU 上微调大型模型部分。

为了验证这些观察结果，我们使用 QLoRA + PEFT 以及可在此处获取的 accelerate 配置，运行了Transformers 强化学习（TRL）库的 SFT（监督微调）官方示例脚本。我们在 2x NVIDIA T4 GPU 上运行了这些实验。

使用 PEFT 和 DeepSpeed 的 ZeRO3 在多设备、多节点上微调大型模型

本指南部分将帮助你学习如何使用我们的 DeepSpeed 训练脚本来执行 SFT。你将配置该脚本，在单台机器上的 8 块 H100 80GB GPU 上，使用 LoRA 和 ZeRO-3 对 Llama-70B 模型进行 SFT（监督微调）。你可以通过修改 accelerate 配置，将其扩展至多台机器。

配置

首先运行以下命令，使用 🤗 Accelerate 创建 DeepSpeed 配置文件。--config_file 标志允许你将配置文件保存到特定位置，否则它会作为 default_config.yaml 文件保存在 🤗 Accelerate 缓存中。

配置文件用于在启动训练脚本时设置默认选项。

accelerate config --config_file deepspeed_config.yaml

系统会询问几个关于您环境配置的问题，请配置以下参数。在本示例中，您将使用 ZeRO-3，因此请确保选择相应的选项。

`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them. Pass the same value as you would pass via cmd argument else you will encounter mismatch error.
`gradient_clipping`: Enable gradient clipping with value. Don't set this as you will be passing it via cmd arguments.
`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2. Set this as `none` as don't want to enable offloading.
`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3. Set this as `none` as don't want to enable offloading.
`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3. Set this to `True`.
`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3. Set this to `True`.
`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training. Set this to `True`.

完成此操作后，相应的配置应如下所示，您可以在 deepspeed_config.yaml 所在的 config 文件夹中找到它：

compute_environment: LOCAL_MACHINE                                                                                                                                           
debug: false
deepspeed_config:
  deepspeed_multinode_launcher: standard
  gradient_accumulation_steps: 4
  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: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

启动命令

启动命令位于 run_peft_deepspeed.sh，如下所示：

accelerate launch --config_file "configs/deepspeed_config.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-lora-deepspeed" \
--per_device_train_batch_size 8 \
--per_device_eval_batch_size 8 \
--gradient_accumulation_steps 4 \
--gradient_checkpointing True \
--use_reentrant False \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization False

请注意，我们使用 LoRA，其中 rank=8、alpha=16，并针对所有线性层。我们传入 deepspeed 配置文件，并在 ultrachat 数据集的子集上微调 70B Llama 模型。

重要部分

让我们更深入地了解一下这个脚本，看看其中发生了什么，并理解其工作原理。

首先要了解的是，由于传入了 DeepSpeed 配置，该脚本使用 DeepSpeed 进行分布式训练。SFTTrainer 类负责处理使用传入的 peft 配置创建 PEFT 模型的所有繁重工作。之后，当你调用 trainer.train() 时，SFTTrainer 内部会使用 🤗 Accelerate 结合 DeepSpeed 配置来准备模型、优化器和训练器，从而创建 DeepSpeed 引擎，然后对其进行训练。主要代码片段如下：

# trainer
trainer = SFTTrainer(
    model=model,
    processing_class=tokenizer,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
)
trainer.accelerator.print(f"{trainer.model}")

# train
checkpoint = None
if training_args.resume_from_checkpoint is not None:
    checkpoint = training_args.resume_from_checkpoint
trainer.train(resume_from_checkpoint=checkpoint)

# saving final model
trainer.save_model()

内存使用

在上面的示例中，如下面的截图所示，每个 GPU 消耗的内存为 64 GB（80%）：

GPU memory usage for the training run

使用 PEFT QLoRA 和 DeepSpeed ZeRO3 在多 GPU 上微调大型模型

在本节中，我们将介绍如何使用 QLoRA 和 DeepSpeed Stage-3 在 2X40GB GPU 上微调 70B Llama 模型。为此，我们首先需要 bitsandbytes>=0.43.3、accelerate>=1.0.1、transformers>4.44.2、trl>0.11.4 和 peft>0.13.0。在使用 Accelerate 配置时，需要将 zero3_init_flag 设置为 true。以下是可在 deepspeed_config_z3_qlora.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: 2
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

启动命令如下，位于 run_peft_qlora_deepspeed_stage3.sh 中：

accelerate launch --config_file "configs/deepspeed_config_z3_qlora.yaml"  train.py \
--seed 100 \
--model_name_or_path "meta-llama/Llama-2-70b-hf" \
--dataset_name "smangrul/ultrachat-10k-chatml" \
--chat_template_format "chatml" \
--add_special_tokens False \
--append_concat_token False \
--splits "train,test" \
--max_seq_len 2048 \
--num_train_epochs 1 \
--logging_steps 5 \
--log_level "info" \
--logging_strategy "steps" \
--eval_strategy "epoch" \
--save_strategy "epoch" \
--push_to_hub \
--hub_private_repo True \
--hub_strategy "every_save" \
--bf16 True \
--packing True \
--learning_rate 1e-4 \
--lr_scheduler_type "cosine" \
--weight_decay 1e-4 \
--warmup_steps 0 \
--max_grad_norm 1.0 \
--output_dir "llama-sft-qlora-dsz3" \
--per_device_train_batch_size 2 \
--per_device_eval_batch_size 2 \
--gradient_accumulation_steps 2 \
--gradient_checkpointing True \
--use_reentrant True \
--dataset_text_field "content" \
--use_flash_attn True \
--use_peft_lora True \
--lora_r 8 \
--lora_alpha 16 \
--lora_dropout 0.1 \
--lora_target_modules "all-linear" \
--use_4bit_quantization True \
--use_nested_quant True \
--bnb_4bit_compute_dtype "bfloat16" \
--bnb_4bit_quant_storage_dtype "bfloat16"

注意新传入的参数 bnb_4bit_quant_storage_dtype，它指定了用于打包 4 位参数的数据类型。例如，当设置为 bfloat16 时，量化后 32/4 = 8 个 4 位参数会被打包在一起。

就训练代码而言，重要的代码更改如下：

...

bnb_config = BitsAndBytesConfig(
    load_in_4bit=args.use_4bit_quantization,
    bnb_4bit_quant_type=args.bnb_4bit_quant_type,
    bnb_4bit_compute_dtype=compute_dtype,
    bnb_4bit_use_double_quant=args.use_nested_quant,
+   bnb_4bit_quant_storage=quant_storage_dtype,
)

...

model = AutoModelForCausalLM.from_pretrained(
    args.model_name_or_path,
    quantization_config=bnb_config,
    trust_remote_code=True,
    attn_implementation="flash_attention_2" if args.use_flash_attn else "eager",
+   dtype=quant_storage_dtype or torch.float32,
)

请注意，AutoModelForCausalLM 的 dtype 与 bnb_4bit_quant_storage 的数据类型相同。就是这样。其余部分均由 Trainer 和 TRL 处理。

内存使用

在上面的示例中，每个 GPU 消耗的内存为 36.6 GB。因此，使用 DeepSpeed Stage 3+LoRA 需要 8 块 80GB GPU，使用 DDP+QLoRA 需要几块 80GB GPU 才能完成的任务，现在只需 2 块 40GB GPU 即可。这使得对大型模型进行微调变得更加容易实现。

使用 PEFT 和 DeepSpeed 的 ZeRO3 及 CPU Offloading 在单 GPU 上微调大型模型

本指南部分将帮助你学习如何使用我们的 DeepSpeed 训练脚本。你将配置该脚本，使用 ZeRO-3 和 CPU Offload 来训练用于条件生成的大模型。

💡 为了帮助您快速上手，请查看我们针对因果语言建模和条件生成的示例训练脚本。您可以将这些脚本适配到自己的应用中，如果您的任务与脚本中的任务相似，甚至可以直接使用它们。

配置

配置文件用于在启动训练脚本时设置默认选项。

accelerate config --config_file ds_zero3_cpu.yaml

系统会询问几个关于您环境配置的问题，请配置以下参数。在本示例中，您将使用 ZeRO-3 并配合 CPU-Offload，因此请确保选择这些选项。

`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
`gradient_clipping`: Enable gradient clipping with value.
`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training.

一个示例配置文件可能如下所示。需要注意的是，zero_stage 被设置为 3，并且 offload_optimizer_device 和 offload_param_device 被设置为 cpu。

compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
mixed_precision: 'no'
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
use_cpu: false

重要部分

让我们更深入地了解一下这个脚本，看看其中发生了什么，并理解其工作原理。

在 main 函数中，该脚本创建一个 Accelerator 类，以初始化分布式训练所需的所有配置。

💡 您可以随意更改 main 函数中的模型和数据集。如果您的数据集格式与脚本中的不同，可能还需要编写自己的预处理函数。

该脚本还会为您所使用的 🤗 PEFT 方法（本例中为 LoRA）创建一个配置。LoraConfig 用于指定任务类型以及重要参数，例如低秩矩阵的维度、矩阵缩放因子和 LoRA 层的 dropout 概率。如果您想使用其他 🤗 PEFT 方法，请确保将 LoraConfig 替换为相应的 class。

 def main():
+    accelerator = Accelerator()
     model_name_or_path = "facebook/bart-large"
     dataset_name = "twitter_complaints"
+    peft_config = LoraConfig(
         task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
     )

Throughout the script, you’ll see the main_process_first and wait_for_everyone functions which help control and synchronize when processes are executed.

get_peft_model() 函数接受一个基础模型以及你之前准备好的 peft_config，用于创建一个 PeftModel：

  model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
+ model = get_peft_model(model, peft_config)

将所有相关的训练对象传递给 🤗 Accelerate 的 prepare，以确保一切就绪，可以开始训练：

model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler = accelerator.prepare(
    model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler
)

下一段代码检查 Accelerator 中是否使用了 DeepSpeed 插件，如果该插件存在，则检查是否正在使用 ZeRO-3。此条件标志用于在推理期间调用 generate 函数时，在模型参数被分片的情况下同步 GPU：

is_ds_zero_3 = False
if getattr(accelerator.state, "deepspeed_plugin", None):
    is_ds_zero_3 = accelerator.state.deepspeed_plugin.zero_stage == 3

在训练循环中，常规的 loss.backward() 被替换为 🤗 Accelerate 的 backward，它会根据您的配置使用正确的 backward() 方法：

  for epoch in range(num_epochs):
      with TorchTracemalloc() as tracemalloc:
          model.train()
          total_loss = 0
          for step, batch in enumerate(tqdm(train_dataloader)):
              outputs = model(**batch)
              loss = outputs.loss
              total_loss += loss.detach().float()
+             accelerator.backward(loss)
              optimizer.step()
              lr_scheduler.step()
              optimizer.zero_grad()

就是这样！脚本的其余部分会处理训练循环、评估，甚至帮你将其推送到 Hub。

训练

运行以下命令以启动训练脚本。之前，您将配置文件保存为 ds_zero3_cpu.yaml，因此您需要使用 --config_file 参数将路径传递给启动器，如下所示：

accelerate launch --config_file ds_zero3_cpu.yaml examples/conditional_generation/peft_lora_seq2seq_accelerate_ds_zero3_offload.py

您会看到一些输出日志，用于跟踪训练过程中的内存使用情况。训练完成后，脚本会返回准确率，并将预测结果与标签进行对比：

GPU Memory before entering the train : 1916
GPU Memory consumed at the end of the train (end-begin): 66
GPU Peak Memory consumed during the train (max-begin): 7488
GPU Total Peak Memory consumed during the train (max): 9404
CPU Memory before entering the train : 19411
CPU Memory consumed at the end of the train (end-begin): 0
CPU Peak Memory consumed during the train (max-begin): 0
CPU Total Peak Memory consumed during the train (max): 19411
epoch=4: train_ppl=tensor(1.0705, device='cuda:0') train_epoch_loss=tensor(0.0681, device='cuda:0')
100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:27<00:00,  3.92s/it]
GPU Memory before entering the eval : 1982
GPU Memory consumed at the end of the eval (end-begin): -66
GPU Peak Memory consumed during the eval (max-begin): 672
GPU Total Peak Memory consumed during the eval (max): 2654
CPU Memory before entering the eval : 19411
CPU Memory consumed at the end of the eval (end-begin): 0
CPU Peak Memory consumed during the eval (max-begin): 0
CPU Total Peak Memory consumed during the eval (max): 19411
accuracy=100.0
eval_preds[:10]=['no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint', 'no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint']
dataset['train'][label_column][:10]=['no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint', 'no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint']

注意事项

在使用 PEFT 和 DeepSpeed 时，目前不支持合并操作，并会引发错误。
在使用 CPU 卸载时，使用 PEFT 将优化器状态和梯度缩减至适配器权重大小所带来的主要收益将体现在 CPU 内存上，而不会节省 GPU 内存。
当 DeepSpeed Stage 3 和 qlora 与 CPU offloading 一起使用时，与禁用 CPU offloading 相比，会导致更多的 GPU 内存使用。

💡 当您的代码需要合并（以及取消合并）权重时，请尝试事先使用 DeepSpeed Zero-3 手动收集参数：

import deepspeed

is_ds_zero_3 = ... # check if Zero-3

with deepspeed.zero.GatheredParameters(list(model.parameters()), enabled= is_ds_zero_3):
    model.merge_adapter()
    # do whatever is needed, then unmerge in the same context if unmerging is required
    ...
    model.unmerge_adapter()

Update on GitHub

PEFT

DeepSpeed

与 bitsandbytes 量化 + LoRA 的兼容性

使用 PEFT 和 DeepSpeed 的 ZeRO3 在多设备、多节点上微调大型模型

配置

启动命令

重要部分

内存使用

更多资源

使用 PEFT QLoRA 和 DeepSpeed ZeRO3 在多 GPU 上微调大型模型

内存使用

使用 PEFT 和 DeepSpeed 的 ZeRO3 及 CPU Offloading 在单 GPU 上微调大型模型

配置

重要部分

训练

注意事项