揭秘两张RTX 3090卡微调70B 大模型的背后

导读

近日Answer AI 联合bitsandbytes 作者Tim Dettmers、HuggingFace和Mobius Labs整合了FDSP和QLoRA 技术搞出了一个新项目，在个人电脑上用消费级游戏显卡就可以完成70B 规模的大模型微调训练，再次把性价比推向了一个新高度。

QLoRA & FSDP

主要思路：

1. 把量化技术和LoRA 相结合，降低模型可训练参数和单个参数内存大小，整体降低模型内存开销，但是传统的量化处理完后无法在之上进行fine-tun，因此引入结合后的QLoRA 技术把一个70B 模型的内存占用降低到35GB。

2. 用Pytoch 提供的FSDP 分布式训练技术把模型切成2 个17.5GB 分片到2 张24G GPU卡上，结合CPU offloading和梯度累计，剩下的6.5GB 足够应对训练模型的额外开销了。

LoRA：

通常大模型中使用原始精度16或者32位时内存开销很大，特别是参数规模较大的模型，考虑到成本，微软提出了一种优化算法LoRA（Low-Rank Adaptation of Large Language Models），主要思路就是把训练过程中模型权重参数更新拆成一个固定部分+一个小的可训练部分，并且可训练部分的维度可以通过rank 参数控制，rank 越大越接近原始模型，如果等于预训练权重矩阵的rank，那相当于就是原始模型了。

QLoRA：

QLoRA 是充分利用了量化和LoRA，用一个量化的base 模型，再加上一个可训练、非量化的LoRA 适配器，充分结合双方的优点，同时又规避掉缺点，实现了1+1 > 2的效果。

FSDP：

对于一个70B 的模型用QLoRA 技术处理完之后仍然需要高达35GB 内存，如果是单卡训练的话，至少需要一张48GB 内存的卡，但通常成本很高。因此，如何把训练扩展到多张GPU 卡上就成了一个新的需要解决的问题。一个常见的做法就是把模型不同层放到不同GPU 上训练，在transformers 库只需要把device_map 参数值设置成auto 就行，缺点就是同一个时刻只有一张卡在运行，其他都等待，效率很低。幸好PyTorch 中提供了FSDP，把参数切分到不同GPU 上，支持GPU 并行训练，当网络中层在某个GPU 上训练时会从其他GPU 拷贝需要的参数，计算完后再删掉，通过巧妙的实现可以做到同等DDP 性能。

如何使用：

如果想用上了这个新技术，需要在Python 环境中安装transformers、peft和bitsandbytes 。准备完基础环境后，克隆项目（https://github.com/AnswerDotAI/fsdp_qlora/tree/main）到本地，然后运行train.py 脚本就可以完成模型训练，具体例子如下，

python train.py \--train_type qlora \--dataset alpaca \--batch_size 8 \--gradient_accumulation_steps 2 \--output_dir qlora_output --log_to wandb

具体模型的训练逻辑可以参考脚本里的详细内容，在这里面就不再展开了。

总结

通过组合量化、LoRA和FSDP 等一系列基础的优化技术，实现了在个人消费级GPU 卡上微调70B 大模型，也算得上是科技普惠的好成果了，特别是对于资金没那么雄厚的小Lab 来说，可以用更低的成本去做更多探索，进一步推动大模型技术的发展和落地。