双卡训练70B大模型：百亿部署实战（上）

# 0x00 引言 自从《在家训练70B的大模型》那篇文章发布后，后台私信瞬间爆满，几乎每位读者都在追问同一个问题：“是不是写错了？应该是7B才对？”坦白讲，最近确实事务缠身，一直未能挤出完整时间进行验证。近日，我陆续完成了一些测试，根据目前的结果判断，70B大模型的训练确实是可行的。不过，截至现在，我本人尚未完全成功运行——原因很简单：新购的内存条还在快递路上。 回顾时间线：2023年初，能将6B参数的模型勉强跑通已属不易。然而仅到2024年，72B模型也能磕磕绊绊地运行了。更令人惊讶的是，仅仅一个月后，连70B大模型的微调都开始变得触手可及。 先直接说结论： **FSDP（减少内存占用） + QLoRA（压缩训练参数） + HQQ（提升训练处理速度）** ——这三项核心技术的组合，成功将70B大模型的训练显存需求压缩至35GB。 那么代价是什么？ 1. 内存消耗大幅增加。要训练70B的大模型，系统内存至少需要128GB。笔者新买的内存条还在路上，不过从目前的侧面测试来看，这条路是走得通的。 2. 加载时间显著变长，训练本身倒还在可接受范围。原因在于数据需要在显存和内存之间反复合并、腾挪，时间成本自然就上去了。 **接下来的内容安排如下：** - 先聊聊背景和原理，这部分主要来源于 **answer.ai** 原文的翻译和摘录。 - 然后是微调的具体代码。 - 最后是在7B模型上的测试结果。 至于70B模型的完整训练准备过程、训练流程以及最终成果，将在下一篇中详细展开。当然，也有可能我还没写完，别人就已经捷足先登了。最近实在是分身乏术……/(ㄒoㄒ)/~~ 备注：在《0x01 背景》和《0x02 原理》两个章节中，“我们”指的是answer.ai的研究团队。 # 0x01 背景 训练深度学习模型所采用的硬件，主要分为两大类，性能与成本差异极大。一类是数据中心级的设备——H100、A100等，单套价格动辄几十万美元。另一类则是我们普通人的桌面电脑，配置两张4090显卡，成本也能控制在1万美元以内。如果善用二手零件攒机，价格还能再砍一半。 关键问题在于：游戏GPU的性能其实与数据中心GPU相差无几，价格却便宜了十倍不止！如果能用这些性价比爆棚的显卡来训练大语言模型，那该多好。可惜现实很骨感——这些显卡的显存容量太小了。目前最好的数据中心卡能配备80GB显存，而游戏卡最多也只有24GB。偏偏只有最大的模型才能产出最好的结果，这就导致了一个尴尬的局面：绝大多数人实际上被挡在了“造好模型”的门槛之外。 仔细想想，这件事本就不该有什么内在的壁垒。高性能的硬件明明就在那里等着被用，我们缺的只不过是一套方法——能让模型和数据在有限的内存限制下流畅运行。那么问题来了：为什么之前没人这么干？答案其实很简单：所有大型工业实验室早就用上了贵十倍的硬件，他们压根儿没有动力去解决这个难题。 # 0x02 原理 ## QLoRA：在单张GPU上训练更大的模型 QLoRA巧妙融合了近年神经网络领域的两大重要进展：量化与LoRA。量化技术，简言之，就是将神经网络权重的存储格式从16位甚至32位压缩至4位（甚至更少）。尽管4位数字只有16种可能取值，但Dettmers与Zettlemoyer的研究表明，对于当前主流的大语言模型而言，这种精度已经绰绰有余。 当模型被量化到4位（即每个参数只需要0.5字节）后，一个70B的模型大约需要70/2 = 35GB，这仍然超出了我们计划使用的那张24GB游戏GPU的容量。 ## FSDP：将训练扩展到多个GPU 一个直观的解决方案是：既然一张卡不够用，那就用多张卡！开源社区里有个很常见的做法——把模型的不同层分别放到不同的显卡上。比如，你想训练一个70B（约140GB）的模型，可以将其分布在8张24GB的GPU上，每张卡分到17.5GB。Hugging Face的Transformers库里甚至提供了一个便利的 `device_map='auto'` 参数，你很可能已经用过——实际上，它在幕后就是这么处理的。 这个方法确实能跑起来，但有一个致命的缺陷：同一时刻只有一张GPU在工作，其他七张卡都在排队等待“轮到自己”。结果是，八分之七的计算能力被白白浪费了。 分布式数据并行（DDP）曾是跨多张GPU高效训练模型的金标准。它的思路是：每张GPU上都保留一份完整的模型副本。如果你的模型较小——比如一个2B的模型，只需要4GB——那完全可以把整个模型分别加载到每张GPU上，然后让它们并行处理不同的训练样本。假设你有4张GPU，训练速度就是原来的4倍。但问题是，如果模型太大，一张卡装不下，同时也腾不出足够的空间来存放训练过程中所需的数据，那么DDP就无法使用了。 因此，我们真正需要的是：既能跨GPU拆分模型（如同 `device_map='auto'` 那样），又能让所有GPU并行工作（如同DDP那样）的方案。Meta的完全分片数据并行（FSDP）库正是为此而生的。它通过将参数拆分到多张GPU上进行“分片”，让所有GPU可以同时参与计算。在训练过程中，当某一层需要在某张特定GPU上计算时，所有相关的分片都会被复制过来；计算完成后，再将这些复制的部分从该GPU上清掉。听起来好像效率很低？实际上，通过在当前层计算的同时提前把下一层所需的数据复制过来，这种做法在速度上完全可以做到不输DDP。 ## 其他技术 要实现这样的突破，光靠FSDP和QLoRA还不够，还得益于过去几年学术界与开源社区积累的一系列巧妙技术。具体来说，我们还用了： - **Gradient checkpointing（梯度检查点）**：也称为激活检查点。它不去存储完整的梯度，而是在模型的多个“检查点”位置保存激活值，需要用到梯度时，再重新跑一遍前向计算来获取。这本质上是用计算换内存。 - **CPU offloading（CPU卸载）**：当模型权重暂时不用时，将它们从GPU搬到CPU内存里。这样可以大幅降低GPU内存的需求。对于拥有H100等高端GPU的“显存富户”来说，这招用处不大——因为H100有高度优化的方式在GPU之间互相传递权重。但在我们的场景里，这绝对是必不可少的，毕竟消费级GPU和主板根本不具备这个能力。 - **Flash Attention 2**：利用内存优化的Cuda内核来高效计算注意力机制。 # 0x03 启动代码 ### 硬件环境 - 2张 3090 / 4090（24GB × 2 = 48GB 显存） - 内存 128GB - CPU 32核 × 2 ### 软件环境 - Ubuntu 22.04 - CUDA 12.1 - PyTorch 2.2.0 ### 安装过程 1. Clone https://github.com/AnswerDotAI/fsdp_qlora 2. `pip install llama-recipes fastcore --extra-index-url https://download.pytorch.org/whl/test/cu118`（将118替换为你需要的Cuda版本）——这是获取大部分依赖的简便途径。 3. `pip install bitsandbytes>=0.43.0` 4. 下载模型 ### Llama-2 70B在两块24GB显卡上的微调代码

python train.py \
--model_name meta-llama/Llama-2-70b-hf \
--batch_size 2 \
--context_length 512 \
--precision bf16 \
--train_type qlora \
--use_gradient_checkpointing true \
--use_cpu_offload true \
--dataset alpaca \
--reentrant_checkpointing true

看到下面的启动信息即表示启动成功： # 0x04 实验 ### 训练的基准

llama-7B，
--batch_size 4 \
--gradient_accumulation_steps 1

### 未开启FSDP之前 - 内存占用：0 - 单张显存：22GB / 24GB ### 开启FSDP之后 - 内存占用：18GB - 单张显存：7GB / 24GB 显存占用骤降至原先的四分之一。更令人惊讶的是，训练参数甚至能够设置为：

--batch_size 8  \
--gradient_accumulation_steps =2

从原来的 batch_size 4 直接提升到 8×2，效率翻了四倍。这效果，确实让人有些难以置信。 # 0x05 总结 速度确实不快，安装过程也确实比较波折——这是实话。但不可否认的是，这项技术实实在在地让普通人也能开始微调70B大模型了。（等我的内存条到了再说吧……目前我还没真正跑起来。）对于小公司，或者只是想尝个鲜的人来说，这绝对是一个值得尝试的方向。