使用PyTorch FSDP对Llama 2 70B进行微调

PyTorch FSDP微调Llama 2 70B

介绍

在本博客文章中，我们将学习如何使用PyTorch FSDP和相关的最佳实践来微调Llama 2 70B。我们将利用Hugging Face Transformers、Accelerate和TRL。我们还将学习如何在SLURM中使用Accelerate。

完全分片数据并行（FSDP）是一种范式，其中优化器状态、梯度和参数在设备之间进行分片。在前向传播过程中，每个FSDP单元执行一个全收集操作，以获取完整的权重，然后进行计算，然后丢弃其他设备的分片。在前向传播之后，计算损失，然后进行反向传播。在反向传播中，每个FSDP单元执行一个全收集操作，以获取完整的权重，然后进行计算以获取局部梯度。这些局部梯度通过减少散射操作进行平均并分片到设备上，以便每个设备可以更新其分片的参数。有关PyTorch FSDP的更多信息，请参阅此博客文章：使用PyTorch完全分片数据并行加速大模型训练。

（来源：链接）

• 认识PyGraft：一款开源的基于Python的人工智能工具，可以生成高度定制化的、与领域无关的模式和知识图谱
• 大型语言模型真的能做数学吗？这篇人工智能研究介绍了MathGLM：一个强大的模型，可以在没有计算器的情况下解决数学问题
• 低成本四足机器人能掌握跑酷吗？揭示了一种革命性的敏捷机器人运动学习系统

使用的硬件

节点数：2。最低要求为1。每个节点的GPU数：8。GPU类型：A100。GPU内存：80GB。节点内部连接：NVLink。每个节点的RAM：1TB。每个节点的CPU核心数：96。节点之间的连接：弹性织物适配器。

微调LLaMa 70B的挑战

在尝试使用FSDP对LLaMa 70B进行微调时，我们遇到了三个主要挑战：

1. FSDP在加载预训练模型后包装模型。如果节点内的每个进程/秩加载Llama-70B模型，将需要70*4*8 GB ~ 2TB的CPU RAM，其中4是每个参数的字节数，8是每个节点上的GPU数。这将导致CPU RAM不足导致进程被终止。

2. 使用FULL_STATE_DICT和在秩0上进行CPU卸载保存完整中间检查点需要很长时间，并且经常由于广播期间的无限挂起导致NCCL超时错误。然而，训练结束时，我们希望获得整个模型状态字典，而不是仅与FSDP兼容的分片状态字典。

3. 我们需要提高训练速度并减少VRAM使用，以便更快地训练并节省计算成本。

让我们看看如何解决上述挑战并微调一个70B模型！

在开始之前，这里是重现我们的结果所需的所有资源：

1. 代码库：https://github.com/pacman100/DHS-LLM-Workshop/tree/main/chat_assistant/training，带有flash-attn V2 monkey patch

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

3. SLURM脚本launch.slurm：https://gist.github.com/pacman100/1cb1f17b2f1b3139a63b764263e70b25

4. 模型：meta-llama/Llama-2-70b-chat-hf

5. 数据集：smangrul/code-chat-assistant-v1（LIMA+GUANACO的混合，格式正确，可直接训练）

先决条件

首先按照以下步骤安装Flash Attention V2：Dao-AILab/flash-attention: Fast and memory-efficient exact attention (github.com)。使用CUDA ≥11.8安装最新的PyTorch夜版。根据DHS-LLM-Workshop/code_assistant/training/requirements.txt安装其余要求。在这里，我们将从主分支安装🤗 Accelerate和🤗 Transformers。

微调

解决挑战1

PRs huggingface/transformers#25107和huggingface/accelerate#1777解决了第一个挑战，并且不需要用户进行任何代码更改。它执行以下操作：

1. 在所有排名上创建没有权重的模型（使用meta设备）。
2. 仅在排名为0时加载状态字典，并将模型权重设置为该状态字典上的排名0。
3. 对于所有其他排名，在meta设备上对每个参数执行torch.empty(*param.size(), dtype=dtype)。
4. 因此，排名为0的将加载具有正确状态字典的模型，而所有其他排名将具有随机权重。
5. 设置sync_module_states=True，以便FSDP对象在训练开始之前将它们广播到所有排名。

下面是在2个GPU上测量7B模型在各个阶段消耗的内存和模型参数的输出片段。我们可以观察到，在加载预训练模型时，排名为0和排名1的CPU总峰值内存分别为32744 MB和1506 MB。因此，只有排名为0的正在加载预训练模型，从而有效地使用了CPU内存。完整的日志可以在此处找到。

accelerator.process_index=0 进入加载之前的GPU内存：0
accelerator.process_index=0 加载结束时的GPU内存消耗（结束-开始）：0
accelerator.process_index=0 加载期间的GPU峰值内存消耗（最大-开始）：0
accelerator.process_index=0 加载期间的GPU总峰值内存消耗（最大）：0
accelerator.process_index=0 进入加载之前的CPU内存：926
accelerator.process_index=0 加载结束时的CPU内存消耗（结束-开始）：26415
accelerator.process_index=0 加载期间的CPU峰值内存消耗（最大-开始）：31818
accelerator.process_index=0 加载期间的CPU总峰值内存消耗（最大）：32744

accelerator.process_index=1 进入加载之前的GPU内存：0
accelerator.process_index=1 加载结束时的GPU内存消耗（结束-开始）：0
accelerator.process_index=1 加载期间的GPU峰值内存消耗（最大-开始）：0
accelerator.process_index=1 加载期间的GPU总峰值内存消耗（最大）：0
accelerator.process_index=1 进入加载之前的CPU内存：933
accelerator.process_index=1 加载结束时的CPU内存消耗（结束-开始）：10
accelerator.process_index=1 加载期间的CPU峰值内存消耗（最大-开始）：573
accelerator.process_index=1 加载期间的CPU总峰值内存消耗（最大）：1506

解决挑战2

通过在创建FSDP配置时选择SHARDED_STATE_DICT状态字典类型来解决。 SHARDED_STATE_DICT单独保存每个GPU的碎片，使得从中间检查点保存或恢复训练变得更快。当使用FULL_STATE_DICT时，第一个进程（排名0）将整个模型聚集在CPU上，然后以标准格式保存。

通过以下命令创建加速配置：

accelerate config --config_file "fsdp_config.yaml"

生成的配置文件在此处可用：fsdp_config.yaml。这里，分片策略是FULL_SHARD。我们使用TRANSFORMER_BASED_WRAP作为自动包装策略，它使用_no_split_module来查找嵌套FSDP自动包装的Transformer块名称。我们使用SHARDED_STATE_DICT以这种格式保存中间检查点和优化器状态，这是PyTorch团队推荐的格式。请确保按照上面解决挑战1的段落中所提到的，在开始时启用从排名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) # 或者，如果整个 ckpt 文件小于 50GB（LFS 每个文件的限制为 50GB），可以使用 trainer.push_to_hub()

解决挑战 3

为了加快训练速度、减少 VRAM 使用量、实现微调并节省计算成本，需要使用 Flash Attention 并启用梯度检查点。该代码库目前使用了猴子补丁，实现位于 chat_assistant/training/llama_flash_attn_monkey_patch.py。

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness 提出了一种在计算精确注意力时更快、更节省内存的方法，通过利用底层硬件/GPU 的内存层次结构知识，具有更高的带宽/速度的内存容量更小，因为它变得更昂贵。

如果我们遵循博客 Making Deep Learning Go Brrrr From First Principles，我们可以发现当前硬件上的 Attention 模块是内存限制/带宽限制。原因是 Attention 主要由逐元素操作组成，如下图左侧所示。我们可以观察到掩码、softmax 和 dropout 操作占用了大部分时间，而不是矩阵乘法，后者占据了大部分 FLOPs。

(来源: 链接)

这正是 Flash Attention 解决的问题。其思想是通过将所有中间步骤保持在 SRAM 中，在执行所有中间步骤后，才将最终结果写回 HBM，也被称为核融合。下图说明了这如何克服内存限制瓶颈。

(来源: 链接)

在正向传播和反向传播过程中使用分块计算 NxN softmax/scores，以克服 SRAM 内存大小的限制，实现分块计算使用了在线 softmax 算法。在反向传播过程中使用重计算，以避免在正向传播过程中存储整个 NxN softmax/score 矩阵，从而大大减少了内存消耗。

要深入了解 Flash Attention 的简化和深入理解，请参考博客 ELI5: FlashAttention 和 Making Deep Learning Go Brrrr From First Principles 以及原始论文 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness。

将一切融合在一起

使用 SLURM 和 Accelerate 启动器运行训练，请参考这个 gist launch.slurm。下面是一个等效命令的示例，展示如何使用 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小时，下面是训练损失图。让我们计算训练过程中的模型Flops利用率（MFU）。

1. A100 GPU每秒执行约3.12e14个FLOPS（float32或bfloat16格式）
2. 上述实验中训练的标记数量=序列长度*批次大小*训练步骤数=(2048*16*
   500. = 16,384,000 = 1.64e7
3. 上述实验的计算量约为6 * P（参数数量）* D（标记数量）=6 * 7e10 * 1.64e7 = 6.89e18个FLOPS
4. 每秒训练FLOPS=近似计算量/训练时间=6.89e17 /（13.5 * 3600）=6.89e18 / 4.86e4 = 1.42e14个FLOPS
5. MFU=每秒训练FLOPS / A100峰值性能=1.42e14 / 3.12e14 = 0.4551 = 峰值性能的45.51%

下面是使用上述模型的示例对话：

系统提示：您是一个乐于助人、尊重和诚实的助手。始终以最大的帮助为目标，同时保持安全。您的回答不应包含任何有害、不道德、种族主义、性别歧视、有毒、危险或非法内容。请确保您的回答在社会上没有偏见，并且具有积极的性质。

如果某个问题没有意义，或者在事实上不连贯，请解释原因，而不是回答不正确的内容。如果您不知道某个问题的答案，请不要分享错误的信息。

- 人类: 什么是深度学习？用海盗的方式解释一下。

+ 助手: 啊，小水手！你是在寻找一种对深度学习的恰当解释吗？
+ 好吧，听好了，伙计！深度学习是一种使用人工神经网络来分析数据的机器学习方法。
+ 是的，伙计们，你没听错！人工神经网络是人脑的模型，但是它们使用相互连接的节点或“人工神经元”来处理信息。
+ 这些节点一起形成处理单元的层，可以从大量的数据中进行学习。
+ 这些系统接收的数据越多，它们在识别模式和做出决策方面就越好。
+ 所以，和深度学习一起，扬起帆，开启冒险之旅吧，伙计们！这是一个等待被发现的知识宝库。
+ 顺风顺水！

- 人类: 现在用厨师的方式解释一下。

+ 助手: 当然可以！这是一个从烹饪角度解释深度学习的例子：

+ 将食谱看作是将原材料转化为美味佳肴的一系列步骤。
+ 类似地，在深度学习中，有多个层次的“成分”（或特征），通过各种操作进行组合和转化，以产生最终的输出或预测。
+ 每一层都建立在前一层的基础上，不断细化和增强输入数据的表示，直到达到所需的结果。
+ 正如不同的香料增强了食物的味道，每一层都为模型对输入的理解增添了深度和复杂性，使其能够做出更准确的预测。
+ 就像厨师必须仔细平衡食谱中的盐、糖和其他调味品的用量一样，在深度学习中的训练过程中，需要调整网络的权重和偏差，以实现准确性和效率的完美融合。
+ 经过足够的实践和尝试，任何人都可以成为一位厨师大师......我是说，深度学习大师！祝您用餐愉快！

整个对话的格式如下：

<|系统|> 系统消息 <|结束文本|> <|提示者|> Q1 <|结束文本|> <|助手|> A1 <|结束文本|> ...

结论

我们成功地使用PyTorch FSDP在多节点多GPU环境中对70B Llama模型进行了微调，并解决了各种挑战。我们看到🤗 Transformers和🤗 Accelerates现在支持使用FSDP初始化大型模型的高效方式，以克服CPU RAM内存不足的问题。接下来，我们介绍了保存/加载中间检查点的推荐方法，以及如何以可直接使用的方式保存最终模型。为了实现更快的训练和减少GPU内存使用量，我们概述了Flash Attention和Gradient Checkpointing的重要性。总的来说，我们可以看到，使用🤗 Accelerate的简单配置可以在多节点多GPU环境中微调如此大型的模型。