1块3090就能训7B大模型，山东大学低带宽低显存训练法，解决显卡限购卡脖子

2024-06-19 06:06

夕小瑶科技说原创
作者 | 谢年年

大模型的全参数微调对资源要求非常高，当前业界更倾向于采用LoRA，Parallel Adapter等参数高效微调（PEFT）方法，通过添加只占用LLMs全部参数很小部分（例如，0.1%）的可训练模块，损失小部分精度以换取低资源高效率的微调。

但对于问答（QA）等知识密集型任务来说，当可训练参数受限时，性能下降较为显著。如下图所示，相比全参数微调，其他PEFT方法下降10%左右。

但我们也从中发现，在Parallel Adapter中随着适配器参数数量的增加，答案准确率呈现出明显的上升趋势。大约需要更新10%的参数，可以达到全量微调的性能。但这一方案需要远超24G的GPU内存支持，这在实际应用中仍然面临较高的资源成本。

今天我们介绍一篇来自山东大学的研究，在可训练参数增加的同时显著降低了GPU内存使用，可实现仅需1块3090（24G）训练7B大模型。并且在保持相近性能的同时，相比其他PEFT方法，内存占用率下降了50%。

论文标题:
MEFT: Memory-Efficient Fine-Tuning through Sparse Adapter

论文链接：
https://arxiv.org/pdf/2406.04984

前置知识：Parallel Adapter

本文是在Parallel Adapter基础上开展的研究，因此需要先简单了解一下该方法。

在Transformer中，FFN（前馈神经网络）起着关键值记忆的作用，其中每个键对应一种文本模式，每个值则指导输出词汇的分布。基于这一发现，Parallel Adapter通过为下游任务定制特定的知识记忆来扩展原始FFN。Parallel Adapter将adapter与FFN并行放置，adapter由两个线性变换矩阵和以及ReLU非线性激活函数组成。其计算过程可以表述为：

Parallel Adapter中并不是每一个神经元都有作用，即存在稀疏性。作者在Natural Questions数据集上训练了一个瓶颈大小为4096的并行适配器模型。然后在测试集（包含4000个tokens）上提取了适配器的FFNs层的激活，并计算了平均激活值与累积激活值。如下图所示：

可以看到，适配器中的激活非常稀疏，即只有部分神经元对模型预测有重大贡献，而大部分神经元未被激活。根据这一观察，本文考虑在训练时仅复制这些重要参数，从而减少CPU-GPU通信量和内存使用。

接下来，我们来看看作者是如何激活重要参数的。

方法

总的来说，本文提出的MEFT方法如上图所示。虚线将参数分成两部分，左侧为GPU，右侧为CPU。大多数可训练的参数将分配给CPU。在前向传播阶段，注意力块的输出将被高效地传递给CPU，使用类似于moe的结构来检索与当前环境高度相关的神经元，将激活神经元将传输到GPU。在反向传播期间，将梯度传输到CPU并更新设置CPU参数。

1. 稀疏激活

FFN块内因ReLU或GELU激活函数存在上下文稀疏性，进一步导致了稀疏梯度。因此本文探索了稀疏Adapter训练，仅更新高激活神经元。在前向计算中，FFN层基于与中个最相似键进行激活:其中，表示所选键的索引,然后在CPU上使用相关的索引构建W_B^K

W[·]WW_AW_BW_A^KW_B^K$ 移动到GPU作为复制适配器，然后作为加宽的FFNs进行计算。

在反向传播中，仅更新激活神经元的梯度，因为未激活神经元不参与FFNPA的计算。通过保留大部分Parallel Adapter参数在CPU内存中，并在每次FFN计算前仅将激活神经元临时复制到GPU内存。由于K远小于总神经元数r，且激活比例通常低于5%，这一策略非常高效,可以显著节省GPU资源。

2. Key-Experts机制

在稀疏激活中，检索最相似权重的TopK操作在CPU上。考虑到当r较大时，给CPU较低的TFLOPs可能成为计算速度的瓶颈。作者进一步提出Key-Experts机制，提高计算效率。

该机制基于MoE的思想，权重和被划分为个专家，并使用一个路由器将输入导向特定的专家。每个专家是一个包含和的FFN。对于输入的token，路由器计算每个专家被选中的分数：然后选择得分最高的K位专家,将这些选定专家的权重连接到和上：选取前k个键值对，得到,,按照下列算法所示计算FFNPA(h)。

3. 效率分析

虽然该方法通过仅将激活的神经元部分放置在GPU上，可以减少GPU内存使用，但CPU与GPU之间的通信以及CPU计算可能会导致GPU等待。作者分析了该方法的通信量和计算复杂度。

通信量

CPU与GPU之间的参数通信分为前向传播和反向传播两部分。

前向传播。对于每一层，隐藏状态h需要从GPU传输到CPU，这导致了B×l×d的通信开销。在参数选择后，大小为2×d×β×K的激活参数会从CPU传输到GPU。这里，B表示批次大小，l表示批次中序列的长度，β是一个与l相关的稀疏因子。
反向传播。对于每一层，GPU上计算得到的激活参数的梯度被移动到CPU，用于更新CPU端的对应参数。因此，通信量大小等于激活参数的大小，即 2 × d × β ×K。

因此，模型训练的总通信开销为：

计算复杂度

在CPU上的额外计算包括路由器上的计算和TopK操作。根据提出的Key-Experts机制，CPU上的复杂度为。因此，当 N 接近时，能达到最优计算复杂度 , 显著降低了计算量，非常适合在CPU上执行。

实证结果

当使用LLaMA-7B作为基础模型，配置为6144的键值对大小、批量大小为2、序列长度为256时，每批的双向通信量约为0.56M次（M代表LLaMA-7B的总可训练参数数）。相比之下，采用deepspeed-offload的Parallel Adapter在相同规模下，每次迭代需要2M的通信量。因此，本文方法在GPU-CPU通信上减少了3.57倍。

在训练效率方面，本文方法相比排除额外通信和CPU计算时间的基线，实证结果显示至少提高了63%的效率。

主要实验结果

下表列出了主要结果。在知识密集型任务（如NQ、SQuAD和Tool）中，MEFT方法在24GB GPU内存限制下显著优于其他PEFT方法。这一提升得益于在有限GPU容量内有效利用了更高的可训练参数比例（即10%）

此外，MEFT在性能上与其他同样包含10%可训练参数的PEFT方法相当，但仅消耗50%的GPU内存，甚至可与全参数模型微调媲美，显著提高了资源利用效率。对于非知识密集型任务如GSM8k，MEFT的稀疏训练策略也展现出不减损性能的稳健性。

▲VRAM代表训练所需的GPU内存。Param显示模型中可训练参数的百分比。Base Model表示原始模型在任务上的zero-shot性能。

稀疏激活机制是否真的有用？

为了减轻CPU的计算负担，本文中使用类似MoE机制。但实验结果表明，这种机制并不是性能提升的主要原因。如下图所示：

参数对知识密集型任务影响甚微，甚至在SQuAD和ToolBench上提高了性能。brutal offload指的是CPU和GPU之间参数的直接交换。然而，在逻辑成分较强的GSM8k任务上，性能略有下降。这表明逻辑任务可能并不需要大量的参数。

效率分析

下图展示了在RTX 3090 GPU和32核CPU（支持AVX）的服务器上每批次训练的延迟。

图中通过消融研究比较了数据传输、CPU计算和GPU计算的训练时间。其中，“"MEFT w/o both”代表将所有可训练参数移至CPU计算，导致了最高延迟；“MEFT w/o Sparse”移除了稀疏激活,但通过PCIe优化传输必要神经元，降低了数据传输时间，提高了GPU效率；“MEFT w/o KE”采用MoE方法管理参数，减少计算负载但涉及完整参数传输；而“Parallel Adapter”则在GPU上执行所有操作，实现最低延迟。