在本文中，您将学习如何使用 Pytorch FSDP 微调 LLAMA 2 70B 以及相关的最佳实践。 在此过程中，我们将主要使用 Hugging Face Transformers、Accelerate 和 TRL 库。 我们还将向您展示如何在 Slurm 中使用 Accelerate。
完全分片数据并行性 （FSDP） 是一种训练范式，其中优化器状态、梯度和模型参数都跨设备分片。 向前传播时，每个 FSDP 单元执行 All Gather 以获取完整权重，然后使用它们进行计算，并在计算后丢弃其他设备的分片。 接下来是反向传播，然后是损耗计算。 反向传播时，每个 FSDP 单元执行 All Gather 操作以获取完整权重，并执行计算以获取本地批次的梯度。 这些梯度通过减少散点在设备上进行平均和分片，以便每个设备都可以更新其相应分片的参数。 有关 PyTorch FSDP 的详细信息，请参阅此博客文章：使用 PyTorch 完全分片数据并行加速大型模型训练。
FSDP 工作流。
节点数：2个，至少1个节点。
每个节点的 GPU 数：8
GPU类型：A100
GPU显存：80GB
节点内互联：nvlink
每个节点的内存：1TB
每个节点的 CPU 核心数：96
节点之间的互连：AWS 的 Elastic Fabric Adapter （EFA）。
在尝试使用 FSDP 微调 LLAMA 2 70B 时，我们遇到了三个主要挑战：
FSDP 在对模型进行分片之前加载整个预训练模型。 这意味着节点内的每个进程（即 Rank）都会加载整个 LLAMA-70B 模型，因此需要 7048 GB 的 2TB CPU 内存，其中 4 是每个参数的字节数，8 是每个节点的 GPU 数。 这会导致 CPU 内存不足，进而导致进程终止。 用full_state_dict保存完整的中间检查点并将其卸载到秩为 0 的 CPU 内存需要花费大量时间，并且通常会导致 NCCL 超时错误，因为通信库需要无限期挂起才能完成保存。 但是，完全关闭此选项并不是一个好主意，因为在训练结束时，我们需要保存完整的模型状态字典，而不是 fsdp 样式分片的状态字典。 我们需要提高速度并减少内存使用，以加快训练速度并节省计算成本。 在下文中，我们将讨论如何解决这些挑战，并最终微调 70b 模型！
首先列出重现结果所需的所有资源：
库：包含启用 Flash Attention v2 的热补丁。
FSDP 配置文件：
slurm 启动脚本 -launch.slurm
型：meta-llama/llama-2-70b-chat-hf
数据集：Smangrul code-chat-assistant-v1（Lima 和 Guanaco 数据集的混合，已转换为训练所需的格式）。
首先按照以下步骤安装 Flash Attention v2。 然后，安装最新的 PyTorch Nightly （cuda 11）。8)。接下来，根据此文件安装其余的相关软件。 在本文中，我们将从 main 分支安装