使用 PyTorch FSDP 微调 Llama 2 70B

这篇具有很好参考价值的文章主要介绍了使用 PyTorch FSDP 微调 Llama 2 70B。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

引言

通过本文,你将了解如何使用 PyTorch FSDP 及相关最佳实践微调 Llama 2 70B。在此过程中,我们主要会用到 Hugging Face Transformers、Accelerate 和 TRL 库。我们还将展示如何在 SLURM 中使用 Accelerate。

完全分片数据并行 (Fully Sharded Data Parallelism,FSDP) 是一种训练范式,在该范式中优化器状态、梯度和模型参数都会被跨设备分片。前向传播时,每个 FSDP 单元执行 all gather 以获取完整的权重,然后用它们进行计算并在计算后丢弃掉其他设备的分片。随后是反向传播,然后就是损失计算。反向传播时,每个 FSDP 单元执行 all gather 操作以获取完整的权重,并执行计算以获得本地 batch 的梯度。这些梯度通过 reduce scatter 在设备上进行均值计算并分片,这样每个设备都可以更新其对应分片的参数。有关 PyTorch FSDP 的更多信息,请参阅此博文: ‍使用 PyTorch 完全分片数据并行技术加速大模型训练‍。

使用 PyTorch FSDP 微调 Llama 2 70B,pytorch,llama,人工智能,python,深度学习
FSDP 工作流

使用的硬件

节点数: 2,至少 1 个节点
每节点 GPU 数: 8
GPU 类型: A100
GPU 显存: 80GB
节点内互联: NVLink
每节点内存: 1TB
每节点 CPU 核数: 96
节点间互联: AWS 的 Elastic Fabric Adapter (EFA)

微调 LLaMa 2 70B 面临的挑战

在尝试使用 FSDP 微调 LLaMa 2 70B 时,我们主要遇到了三个挑战:

  1. FSDP 会先加载整个预训练模型,然后再对模型进行分片。这样就意味着节点内的每个进程 (即 rank) 都会加载整个 Llama-70B 模型,因此需要 7048 GB ~ 2TB 的 CPU 内存,这个算式中 4 是每个参数所需字节数,8 是每个节点的 GPU 数。这会导致 CPU 内存不足,进而导致进程终止。

  2. 使用 FULL_STATE_DICT 来保存完整中间检查点并将其卸载至 rank 0 的 CPU 内存中需要花费大量时间,且由于在此期间通信库需要无限期挂起等待保存完成,因此经常会导致 NCCL 超时错误。然而,完全关掉这个选项也不好,因为在训练结束时我们需要保存完整的模型状态字典,而不是 FSDP 式分片的状态字典。

  3. 我们需要提高速度并减少显存使用,以加快训练并节约计算成本。

下文,我们主要讨论如何一一解决上述挑战,最终微调出一个 70B 的模型!

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

  1. 代码库: https://github.com/pacman100/DHS-LLM-Workshop/tree/main/chat_assistant/training,代码中包含了使能 flash 注意力 V2 的热补丁

  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。然后,安装最新的 PyTorch nightly (CUDA ≥11.8)。接着,根据 此文件 安装其余依赖软件。在本文中,我们是从主分支安装 🤗 Accelerate 和 🤗 Transformers 的。

微调

应对挑战 1

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

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

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

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

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

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

下面是在 2 个 GPU 上加载 7B 模型的输出日志片段,它测量了各个阶段内存的消耗及其加载的模型参数量。我们可以观察到,在加载预训练模型时,rank 0 和 rank 1 的 CPU 峰值内存分别为 32744 MB1506 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

应对挑战 2

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

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

accelerate config --config_file "fsdp_config.yaml"
使用 PyTorch FSDP 微调 Llama 2 70B,pytorch,llama,人工智能,python,深度学习
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 上去。

应对挑战 3

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

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

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

使用 PyTorch FSDP 微调 Llama 2 70B,pytorch,llama,人工智能,python,深度学习
注意力机制的性能瓶颈

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

使用 PyTorch FSDP 微调 Llama 2 70B,pytorch,llama,人工智能,python,深度学习
算子融合

在前向和反向传播过程中我们还使用了 平铺 (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_ipmain_process_portmachine_ranknum_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 小时,下图给出了训练损失曲线。

使用 PyTorch FSDP 微调 Llama 2 70B,pytorch,llama,人工智能,python,深度学习
训练损失曲线

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

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 上微调大模型。

🤗 宝子们可以戳 阅读原文 查看文中所有的外部链接哟!


英文原文: https://hf.co/blog/ram-efficient-pytorch-fsdp

原文作者: Sourab Mangrulkar,Sylvain Gugger,Lewis Tunstall,Philipp Schmid

译者: Matrix Yao (姚伟峰),英特尔深度学习工程师,工作方向为 transformer-family 模型在各模态数据上的应用及大规模模型的训练推理。

FSDP MFU (Model FLOPS Utilization) 相关讨论:https://github.com/huggingface/blog/issues/1649文章来源地址https://www.toymoban.com/news/detail-762246.html

到了这里,关于使用 PyTorch FSDP 微调 Llama 2 70B的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 【大模型系列 06】LLaMA-7B/13B for PyTorch 昇腾迁移

    https://gitee.com/ascend/ModelZoo-PyTorch/tree/master/PyTorch/built-in/foundation/LLaMA-13B LLaMA是由Meta AI发布的大语言系列模型,完整的名字是Large Language Model Meta AI。LLaMA按照参数量的大小分为四个型号:LLaMA-7B、LLaMA-13B、LLaMA-30B与LLaMA-65B。LLaMA 模型的效果极好,LLaMA-13B在大多数基准测试中的表现

    2024年02月12日
    浏览(49)
  • 使用LLaMA-Factory微调ChatGLM3

    略 (1)下载LLaMA-Factory https://github.com/hiyouga/LLaMA-Factory (2)安装依赖 (3)启动LLaMA-Factory的web页面 得到如下页面: 设置如下参数,点击开始即可: 点击“预览命令”,可以看到要执行的python脚本,如下所示:

    2024年02月03日
    浏览(31)
  • 使用QLoRA对Llama 2进行微调的详细笔记

    使用QLoRA对Llama 2进行微调是我们常用的一个方法,但是在微调时会遇到各种各样的问题,所以在本文中,将尝试以详细注释的方式给出一些常见问题的答案。这些问题是特定于代码的,大多数注释都是针对所涉及的开源库以及所使用的方法和类的问题。 对于大模型,第一件事

    2024年02月08日
    浏览(30)
  • 使用 1 个 GPU 在 Colab 上微调 LLama 2.0

    语言模型彻底改变了自然语言处理任务,使计算机能够生成连贯的文本、回答问题,甚至进行对话。像 GPT-3.5 这样的预训练模型已经取得了显着的成果,但研究人员和开发人员正在不断突破这些模型的极限。在这篇博文中,我们将深入探讨最新的 Meta 模型 LLMA 2(Meta 于 2023 年

    2024年02月09日
    浏览(26)
  • Python - 深度学习系列30 - 使用LLaMA-Factory微调模型

    最实用的一种利用大语言模型的方式是进行微调。预训练模型与我们的使用场景一定会存在一些差异,而我们又不可能重头训练。 微调的原理并不复杂,载入模型,灌新的数据,然后运行再训练,保留checkpoints。但是不同项目的代码真的不太一样,每一个都要单独去看的话比

    2024年04月12日
    浏览(33)
  • 使用 LoRA 技术对 LLaMA 65B 大模型进行微调及推理

    前几天,Meta 发布了 LIMA 大模型,在LLaMA-65B的基础上,无需使用 RLHF,只用了 1000 个精心准备的样本数据进行微调,就达到了和 GPT-4 相媲美的程度。这激发了我探索 LLaMA 65B 大模型的兴趣。 之前的一系列大模型相关文章都是在LLaMA 7B/13B模型参数上面进行微调,文本使用 LoRA 技

    2024年02月08日
    浏览(41)
  • 小白也能微调大模型:LLaMA-Factory使用心得

    大模型火了之后,相信不少人都在尝试将预训练大模型应用到自己的场景上,希望得到一个垂类专家,而不是通用大模型。 目前的思路,一是RAG(retrieval augmented generation),在模型的输入prompt中加入尽可能多的“目标领域”的相关知识,引导模型在生成时尽量靠拢目标领域,运

    2024年04月13日
    浏览(42)
  • LLaMA Factory+ModelScope实战——使用 Web UI 进行监督微调

    文章原始地址:https://onlyar.site/2024/01/14/NLP-LLaMA-Factory-web-tuning/ 大语言模型微调一直都是一个棘手的问题,不仅因为需要大量的计算资源,而且微调的方法也很多。在尝试每种方法过程中,配置环境和第三方库也颇为麻烦。。而 LLaMA Factory 1 是一个高效的大语言模型训练和推理

    2024年04月10日
    浏览(52)
  • 使用 QLoRA 进行微调Llama 2 和 Mistral的初学者指南

    本指南适用于任何想要为自己的项目定制强大的语言模型(如 Llama 2 和 Mistral)的人。使用 QLoRA,我们将逐步完成微调这些大型语言模型 (LLMs) 的步骤,即使您没有可供使用的超级计算机。 关键点:一个好的模型需要好的数据。我们将介绍对现有数据的训练以及如何创建自

    2024年02月22日
    浏览(52)
  • [NLP]使用Alpaca-Lora基于llama模型进行微调教程

    Stanford Alpaca 是在 LLaMA 整个模型上微调,即对预训练模型中的所有参数都进行微调(full fine-tuning)。但该方法对于硬件成本要求仍然偏高且训练低效。 [NLP]理解大型语言模型高效微调(PEFT) 因此, Alpaca-Lora 则是利用 Lora 技术,在冻结原模型 LLaMA 参数的情况下,通过往模型中加

    2024年02月15日
    浏览(43)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包