LLMs之Llama2 70B：使用 PyTorch FSDP 微调 Llama 2 70B实现全部过程讲解之详细攻略

地址

文章地址：https://huggingface.co/blog/ram-efficient-pytorch-fsdp

时间

2023年9月13日

作者

Sourab Mangrulkar

Sylvain Gugger

Lewis Tunstall

Philipp Schmid

总结

这篇文章主要介绍了如何利用PyTorch实现的全分片数据并行Fully Sharded Data Parallel(FSDP)进行模型精调，以实现高效利用资源训练大模型。

背景痛点：

>> 直接加载70B模型会占用过多CPU内存。如果每个进程加载完整模型，将需要消耗2TB以上的CPU内存，导致OOM错误。

>> 如果使用FULL_STATE_DICT保存中间检查点，将整个模型状态都保存到CPU会花费很长时间。使用FSDP保存中间检查点或者最终模型时，将整个模型状态从rank0串行广播到其他 ranks 需要很长时间，容易发生NCCL Timeout错误。

>> 需要提高训练速度，减少GPU内存占用来降低训练成本

具体解决方案：

>> 升级Transformers和Accelerate，支持在所有进程使用虚拟设备进行初始化，只在rank0加载模型权重，避免重复加载消耗巨量CPU内存。

>> 利用FSDP只在进程0加载模型权重，其他进程初始化随机权重，通过广播 synchronization解决内存问题。

>> 使用FSDP的分片状态字典格式保存检查点，可以快速并行保存/加载，避免中间广播问题。训练结束使用全状态字典保存最终模型。

>> 使用SHARDED_STATE_DICT格式保存检查点，每个GPU单独保存分片状态，节省时间。

>> 使用Flash Attention提高速度，实现内存高效的精确注意力计算。使用Flash Attention V2加速注意力操作，使用梯度累积减少GPU内存开销，加速训练过程。

>> 开启梯度累积减少GPU内存占用。

核心特点：

>> 充分利用多节点多卡资源，支持TB级别模型在分布式环境下高效训练。

>> 检查点和模型保存加载方式的优化，保证训练过程和结果的稳定性。

>> FSDP实现跨GPU、节点数据并行训练。

>> 优化加载/保存检查点流程。

>> Flash Attention通过内存层次知识提高训练速度。Flash Attention技术提升了注意力运算的效率，在保证精度的同时提高训练速度，节约计算成本。

主要优势：

>> 高效利用计算资源训练大模型

>> 更快训练速度，低内存占用降低成本

>> 联合多种技术完美解决模型精调中的内存和效率问题

总之，这篇文章通过实践案例，系统总结了如何利用PyTorch FSDP等技术高效训练大模型，从根本上解决了内存和速度方面的障碍。

本文介绍了如何利用Torch嵌入式FSDP实施并行学习策略，为大规模深度学习模型在多节点多GPU环境下提供高效训练的实践方案。它解决了内存、检查点和速度三大问题，在保障结果质量的同时最大限度提升资源利用率。

配置

节点数： 2，至少 1 个节点

每节点 GPU 数： 8

GPU 类型： A100

GPU 显存： 80GB

节点内互联： NVLink

每节点内存： 1TB

每节点 CPU 核数： 96

节点间互联： AWS 的 Elastic Fabric Adapter (EFA)

所有资源

先列出重现结果所需的所有资源：

代码库： https://github.com/pacman100/DHS-LLM-Workshop/tree/main/chat_assistant/training

，代码中包含了使能 flash 注意力 V2 的热补丁

FSDP 配置文件： https://github.com/pacman100/DHS-LLM-Workshop/blob/main/chat_assistant/training/configs/fsdp_config.yaml

SLURM 启动脚本 - launch.slurm : https://gist.github.com/pacman100/1cb1f17b2f1b3139a63b764263e70b25

模型: meta-llama/Llama-2-70b-chat-hf

数据集: smangrul/code-chat-assistant-v1 (混合了 LIMA 和 GUANACO 数据集，且已转换为训练所需的格式)

修改

PR 25107 和 PR 1777 解决了第一个挑战，且无需用户侧更改任何代码。主要做的事情如下:

>> 在所有 rank 上创建无权重的空模型 (使用 meta 设备)

>> 仅在 rank 0 上将状态字典加载至模型

>> 其他 rank 仅对 meta 设备上的参数执行 torch.empty(*param.size(), dtype=dtype)

>> 因此，只有 rank 0 上加载了完整的模型及权重，而所有其他 rank 上的权重是空的

>> 设置 sync_module_states=True ，以便 FSDP 实例在训练开始之前将权重广播到各 rank

解释

下面是在 2 个 GPU 上加载 7B 模型的输出日志片段，它测量了各个阶段内存的消耗及其加载的模型参数量。我们可以观察到，在加载预训练模型时，rank 0 和 rank 1 的 CPU 峰值内存分别为 32744 MB 和 1506 MB 。因此可知，仅有 rank 0 加载了预训练模型，这就实现了 CPU 内存的有效利用。你可在 此处 找到完整日志。

accelerator.process_index=0 GPU Memory before entering the loading : 0

accelerator.process_index=0 GPU Memory consumed at the end of the loading (end-begin): 0

accelerator.process_index=0 GPU Peak Memory consumed during the loading (max-begin): 0

accelerator.process_index=0 GPU Total Peak Memory consumed during the loading (max): 0

accelerator.process_index=0 CPU Memory before entering the loading : 926

accelerator.process_index=0 CPU Memory consumed at the end of the loading (end-begin): 26415

accelerator.process_index=0 CPU Peak Memory consumed during the loading (max-begin): 31818

accelerator.process_index=0 CPU Total Peak Memory consumed during the loading (max): 32744

accelerator.process_index=1 GPU Memory before entering the loading : 0

accelerator.process_index=1 GPU Memory consumed at the end of the loading (end-begin): 0

accelerator.process_index=1 GPU Peak Memory consumed during the loading (max-begin): 0

accelerator.process_index=1 GPU Total Peak Memory consumed during the loading (max): 0

accelerator.process_index=1 CPU Memory before entering the loading : 933

accelerator.process_index=1 CPU Memory consumed at the end of the loading (end-begin): 10

accelerator.process_index=1 CPU Peak Memory consumed during the loading (max-begin): 573

accelerator.process_index=1 CPU Total Peak Memory consumed during the loading (max): 1506

修改

该挑战可以通过在配置 FSDP 时将状态字典类型设为 SHARDED_STATE_DICT 来解决。设为 SHARDED_STATE_DICT 后，每个 rank 各自保存各自 GPU 所需要的分片，这使得用户可以快速保存中间检查点并快速从其恢复训练。而当使用 FULL_STATE_DICT 时，第一个进程 (rank 0) 会用 CPU 收集整个模型，然后将其保存为标准格式。

我们可以用以下命令创建相应的 accelerte 配置文件:

accelerate config --config_file "fsdp_config.yaml"

解释

fsdp 配置

你可以从此处获取生成的配置文件: fsdp_config.yaml。在该配置文件中，分片策略是 FULL_SHARD 。我们使用 TRANSFORMER_BASED_WRAP 作为自动模型包装策略，它使用 _no_split_module 来搜索 transformer 块名并自动进行嵌套 FSDP 包装。我们使用 SHAARDED_STATE_DICT 把中间检查点和优化器状态保存为 PyTorch 官方推荐的格式。同时，如上一节 应对挑战 1 中所述，我们还需要确保训练开始时用 rank 0 来广播参数。从配置文件中你还可以看到我们用的是 bf16 混合精度训练。

那么，在保存最终检查点时，如果将其保存成单个文件呢？我们使用的是以下代码段:

if trainer.is_fsdp_enabled:

    trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")

trainer.save_model(script_args.output_dir) # 或者 , 如果整个模型小于 50 GB (即 LFS 单文件的最大尺寸)，你还可以使用 trainer.push_to_hub() 把模型推到 hub 上去。

修改

为了加快训练速度并减少显存占用，我们可以使用 flash 注意力并开启梯度检查点优化，从而在微调的同时节省计算成本。当前，我们用了一个热补丁来实现 flash 注意力，具体代码可见 这儿。

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness 一文基于对底层硬件 (即 GPU) 的内存层次结构的深刻理解而引入了一种更快、更节省内存的无损注意力加速算法。底层硬件在设计内存层次结构时，遵循的实践原则是: 带宽/速度越高的内存，其容量越小，因为它更贵。

解释

根据博文 根据第一性原理让深度学习性能起飞，我们可以发现，当前硬件上的注意力模块是 内存带宽受限 的。原因是注意力机制 主要由逐元素操作 组成，如下左图所示。我们可以观察到，掩码、softmax 和 dropout 操作占用了大部分时间，而非需要大量 FLOP 的矩阵乘法。

注意力机制的性能瓶颈

这正是 flash 注意力解决的问题，其想法是 去除冗余的 HBM 读/写操作。该算法通过将所有内容保留在 SRAM 中，待执行完所有中间步骤后再将最终结果写回到 HBM，即 算子融合 来实现这一目的。下图简要描述了算子融合是如何克服内存瓶颈的。

算子融合

在前向和反向传播过程中我们还使用了 平铺 (Tiling) 优化技巧，将 NxN 大小的 softmax 分数计算切成块，以克服 SRAM 内存大小的限制。在使用平铺技巧时，我们会使用在线 softmax 算法。同时，我们还在反向传播中使用了 重计算 技巧，以大大降低在前向传播过程中存储整个 NxN softmax 分数矩阵所带来的内存消耗。

如欲深入理解 flash 注意力，请参考博文 ELI5: FlashAttention、根据第一性原理让深度学习性能起飞 以及原始论文 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness。

解释

你可参考 此脚本，以在 SLURM 中用 Accelerate 启动器运行训练。下面还给出了一个等效命令，展示了如何使用 Accelerate 启动器来运行训练。请注意，该命令会覆盖 fsdp_config.yaml 中的 main_process_ip 、 main_process_port 、 machine_rank 、 num_processes 以及 num_machines 配置。另一个需要重点注意的是，这里的存储是所有节点共享的。

脚本

accelerate launch \

    --config_file configs/fsdp_config.yaml \

    --main_process_ip $MASTER_ADDR \

    --main_process_port $MASTER_PORT \

    --machine_rank \$MACHINE_RANK \

    --num_processes 16 \

    --num_machines 2 \

    train.py \

    --model_name "meta-llama/Llama-2-70b-chat-hf" \

    --dataset_name "smangrul/code-chat-assistant-v1" \

    --max_seq_len 2048 \

    --max_steps 500 \

    --logging_steps 25 \

    --eval_steps 100 \

    --save_steps 250 \

    --bf16 True \

    --packing True \

    --output_dir "/shared_storage/sourab/experiments/full-finetune-llama-chat-asst" \

    --per_device_train_batch_size 1 \

    --gradient_accumulation_steps 1 \

    --dataset_text_field "content" \

    --use_gradient_checkpointing True \

    --learning_rate 5e-5 \

    --lr_scheduler_type "cosine" \

    --weight_decay 0.01 \

    --warmup_ratio 0.03 \

    --use_flash_attn True

整个微调过程需要约 13.5 小时，下图给出了训练损失曲线。

案例

下例给出了使用上述模型完成的一段对话:

System Prompt: You are a helpful, respectful and honest assistant. Always answer as helpfully \

as possible, while being safe. Your answers should not include any harmful, \

unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that \

your responses are socially unbiased and positive in nature.

If a question does not make any sense, or is not factually coherent, explain why \

instead of answering something not correct. If you don’t know the answer to a \

question, please don’t share false information.

- Human: What is Deep Learning? Explain like a Pirate.

+ Assistant: Arrr Grumete! Are ye lookin' fer a proper explanation o' Deep Learnin'?

+ Alright then, listen up me hearty! Deep Learnin' be a type o' machine learnin' that uses artificial neural networks to analyze data.

+ Aye, ye heard that right, mateys! Artificial neural networks be models o' the human brain, but instead o' using neurons, they use interconnected nodes or "artificial neurons" to process information.

+ These nodes work together to form layers o' processing units that can learn from large amounts o' data.

+ The more data these systems receive, the better they get at recognizing patterns and makin' decisions on their own.

+ So hoist the sails and set course for adventure with deep learnin', me hearties! It's a treasure trove o' knowledge just waitin' to be discovered.

+ Fair winds and following seas!

- Human: Now explain it like a chef.

+ Assistant: Certainly! Here's an explanation of deep learning from a culinary perspective:

+ Think of a recipe as a sequence of steps used to transform raw ingredients into a delicious dish.

+ Similarly, in deep learning, there are multiple layers of "ingredients" (or features) that are combined and transformed through various operations to produce a final output or prediction.

+ Each layer builds upon the previous one, refining and enhancing the representation of the input data until it reaches the desired outcome.

+ Just as different spices enhance the flavors of food, each layer adds depth and complexity to the model's understanding of the input, allowing it to make more accurate predictions.

+ And much like how chefs must carefully balance the amount of salt, sugar, and other seasonings in a recipe, the training process in deep learning involves adjusting the weights and biases of the network to achieve the perfect blend of accuracy and efficiency.

+ With enough practice and experimentation, anyone can become a master chef... I mean, a master of deep learning! Bon appétit!

格式如下

整个对话使用的格式如下:

<|system|> system message <|endoftext|> <|prompter|> Q1 <|endoftext|> <|assistant|> A1 <|endoftext|> ...

总结

我们在多节点多 GPU 上使用 PyTorch FSDP 成功微调了一个 70B Llama 模型，并在此过程中解决了各种挑战。我们看到了当前在  Transformers 和  Accelerates 中应如何初始化大模型从而有效克服 CPU 内存不足的问题。我们还给出了如何高效地保存/加载中间检查点，同时又能以易于使用的方式保存最终模型的最佳实践。为了加速训练并减少 GPU 显存使用，我们还强调了 flash 注意力和梯度检查点机制的重要性。最后，我们向大家展示了在  Accelerate 上仅需要简单的配置就可以在多节点多 GPU 上微调大模型。