DeepSpeed
DeepSpeed 是一个为速度和规模而设计的库,用于对具有数十亿参数的大型模型进行分布式训练。它的核心是零冗余优化器 (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 阶段 3 + QLoRA,请参阅下面的部分 使用 PEFT QLoRA 和 DeepSpeed 与 ZeRO3 在多个 GPU 上微调大型模型。
为了确认这些观察结果,我们运行了 Transformers 强化学习 (TRL) 库 的 SFT (监督微调) 官方示例脚本,使用 QLoRA + PEFT 和 这里 提供的 accelerate 配置。我们在 2 个 NVIDIA T4 GPU 上运行了这些实验。
使用 PEFT 和 DeepSpeed 与 ZeRO3 在多个设备和多个节点上微调大型模型
本指南的这一部分将帮助您了解如何使用我们的 DeepSpeed 训练脚本 执行 SFT。您将配置该脚本以对 Llama-70B 模型进行 SFT (监督微调),使用 LoRA 和 ZeRO-3 在一台机器上的 8 个 H100 80GB GPU 上运行。您可以通过更改 accelerate 配置将其配置为扩展到多台机器。
配置
首先运行以下命令以使用 🤗 Accelerate 创建 DeepSpeed 配置文件。--config_file
标志允许您将配置文件保存到特定位置,否则它将作为 🤗 Accelerate 缓存中的 default_config.yaml
文件保存。
配置文件用于在启动训练脚本时设置默认选项。
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 中找到它
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_ratio 0.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,
tokenizer=tokenizer,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
peft_config=peft_config,
packing=data_args.packing,
dataset_kwargs={
"append_concat_token": data_args.append_concat_token,
"add_special_tokens": data_args.add_special_tokens,
},
dataset_text_field=data_args.dataset_text_field,
max_seq_length=data_args.max_seq_length,
)
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 内存使用情况更多资源
您还可以参考这篇博客文章 使用 🤗 PEFT 和 DeepSpeed 微调 Falcon 180B 模型,了解如何在 2 台机器上的 16 个 A100 GPU 上微调 180B Falcon 模型。
使用 PEFT QLoRA 和 DeepSpeed 与 ZeRO3 在多个 GPU 上微调大型模型
在本节中,我们将研究如何使用 QLoRA 和 DeepSpeed 阶段 3 在 2 个 40GB GPU 上微调 70B llama 模型。为此,我们首先需要 bitsandbytes>=0.43.0
、accelerate>=0.28.0
、transformers>4.38.2
、trl>0.7.11
和 peft>0.9.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_ratio 0.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",
+ torch_dtype=quant_storage_dtype or torch.float32,
)
请注意,AutoModelForCausalLM
的 torch_dtype
与 bnb_4bit_quant_storage
数据类型相同。就是这样。所有其他操作都由 Trainer 和 TRL 处理。
内存使用情况
在上面的示例中,每个 GPU 占用的内存为 **36.6 GB**。因此,以前使用 DeepSpeed Stage 3+LoRA 的 8X80GB GPU 和使用 DDP+QLoRA 的几个 80GB GPU 所需的资源,现在只需要 2X40GB GPU 即可。这使得大型模型的微调更加容易。
使用 PEFT 和 DeepSpeed 结合 ZeRO3 和 CPU Offloading 在单个 GPU 上微调大型模型
本节指南将帮助您学习如何使用我们的 DeepSpeed 训练脚本。您将配置脚本以使用 ZeRO-3 和 CPU Offload 训练大型条件生成模型。
配置
首先运行以下命令以使用 🤗 Accelerate 创建 DeepSpeed 配置文件。--config_file
标志允许您将配置文件保存到特定位置,否则它将作为 🤗 Accelerate 缓存中的 default_config.yaml
文件保存。
配置文件用于在启动训练脚本时设置默认选项。
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 层的丢弃概率。如果您想使用其他 🤗 PEFT 方法,请确保将 LoraConfig
替换为适当的 类。
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
)
在整个脚本中,您将看到 main_process_first 和 wait_for_everyone 函数,它们有助于控制和同步何时执行进程。
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/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 Offloading 时,使用 PEFT 将优化器状态和梯度缩减为适配器权重的主要优势将在 CPU RAM 上实现,而不会节省 GPU 内存。
- 与禁用 CPU Offloading 相比,DeepSpeed Stage 3 和 qlora 与 CPU Offloading 一起使用会导致更高的 GPU 内存使用量。