HiFT全参数微调新范式---逐层微调_mezo和lomo优化器

论文链接: https://arxiv.org/abs/2401.15207

HiFT 是一个端到端的层级优化策略。目前论文的结果是原始混合精度的结果，目前最新进展已将混合精度进行了分层适配，微调7B模型的内存需求约为16.87G，13B模型约为31G(batch=1,seq_length=512)

背景

在大语言模型之前，语言模型适配下游任务的首选方法就是全参数微调。随着大语言模型的出现，全参数微调对GPU的内存需求越来越高。以LORA为代表的PEFT（Parameter-Efficient Fine-tuning）方法成为微调大模型的首选。PEFT的方法以较低的GPU内存代价在多个任务上取得了媲美全参数微调的性能。但是已有的研究表明PEFT的方法整体上性能和全参数微调相比仍有差距。

最近的一些工作开始关注内存高效的全参数微调方法。这些方法的特点就是采用原始SGD优化器的思路----无动量优化。类似SGD这样的优化器没有优化器的状态，因此可以极大的降低微调过程中优化器占用的内存。其中代表的就是MeZO和LOMO优化器。MeZO优化器在prompt的微调下性能接近全参数微调，在无prompt的情况下性能和全参数微调差距较大。MEZO能做到在80G内存设备上微调30B的模型。LOMO优化器除了无动量优化器外，采用了融合梯度计算和更新的策略，这样做可以同时降低梯度的内存占用，但是LOMO需要forward两次。经过我们测试，在不使用外部优化技术的情况下，在内存节省上MeZO更有优势。这些零阶优化器虽然可以节省内存，但是类似Adamw这种优化器器训练更加稳定，更容易收敛。使用这些优化器替代被广泛验证过的优化器代价是巨大的。

内存被谁占用了

微调过程中占用内存的主要部分有：模型参数 ，梯度 ，优化器状态，剩余部分（有的地方叫激活状态，这里我们参考混合精度论文的叫法，主要包括，激活状态，图中间变量等）。模型参数由于需要进行前向传播，该部分很难进行优化，或者说我们必须要把它加载到GPU上。标准的全参数微调梯度等于模型参数；优化器状态参数量取决于使用了几阶动量，AdamW使用了二阶动量，状态是模型参数的2倍，SGD为零阶优化器，状态参数为0，SGDM加入了迭代动量，该部分参数和模型参数相等。剩余部分主要是由图的激活状态和中间变量占用，该部分占用的内存和输入呈正相关，也就是输入长句越长，batch size 越大该部分参数占用的内存就越多，也叫动态内存占用。

framework

图1：HiFT 策略的示意图。Group表示对图层的分组操作。bottom2up、top2down 和 random 代表训练策略。灰色表示对应参数处于冻结状态，棕色表示对应参数处于激活状态。K表示组数，n表示给定模型的层数，BP表示通过反向传播进行参数更新。

如何降低内存占用

梯度，优化器状态和神经元的激活都和模型的可训练参数量相关。既然一次更新所有参数有困难，那是否可以一次只更新一层参数？这好比爬山，当路比较宽的时候，大家可以同时走，当路比较窄的时候只能排着队走，只要保证后面的人不掉队保持步调一致就行。实际上模型更新也可以这么做。具体的做法如下：第一步，先将模型进行分组，如上图所示，假设分了K组，K最大为模型的层数。分完组后，在微调过程中，每个step只更新其中一组参数，假设我们选择了自低向上 (bottom2up）的更新方式, 当更新低2组的参数时候，冻结其它层的参数，下一个step更新第三组的参数同时冻结其它层，按照自低向上的顺序依次更新每一层参数，直至模型收敛。这样更新方式有一个问题，如果每个step都对学习率更新，那会导致第K组参数更新幅度过大导致整个模型参数更细幅度不一致出现loss震荡的情况，这就好比排队爬山，如果第一个人永远迈大步，而最后一个人永远迈小步，最后就会出现人员掉队的问题，模型的更新也是同样的道理。因此我们采用了学习率延迟更新的策略，只有当所有层都更细一遍的时候，才对学习率进行更新。

这种分层更新的方式每个step只有一组参数是需要计算梯度的，剩下的参数都是冻结状态。当K等于模型层数时，训练过程中的峰值训练参数量等于模型中参数量最大的那一层的参数量。可训练参数量的减少，将直接会降低梯度参数的内存使用。由于每个step只更新部分参数，因此没有必要将所有的优化器状态都同时保存在GPU上，所以每个step只有需要进行梯度更新参数的状态会在GPU上，而其它参数的优化器状态都会被保留在CPU上。

图2：在 E2E  数据集上微调 LLaMA2-7B 的 GPU 内存使用情况。序列长度和批量大小分别设置为 512 和 6。#Dtype表示训练时使用的数据类型，其中FP32表示以32位精度对模型进行全参数微调，mixed表示以混合精度进行微调。#Trainable 参数表示微调过程中单个步骤中出现的最大可训练参数数量。#Para表示模型参数占用的内存，#Gra表示梯度占用的内存，#Sta表示优化器状态占用的内存。#PGS 表示模型参数（即#Para）、梯度（即#Gra）和优化器状态（即#Sta）占用的内存总和。残留状态主要包括激活状态、临时缓冲区和不可用的碎片内存。Total 表示微调期间使用的总内存。HiFT的参数K设置为最大，既每组一层参数。 需要说明的是，混合精度的结果是混合精度未适配分层微调下的结果。 最新的结果混合精度适配分层策略后，支持24G设备全参数微调7B模型

以adamw 优化器为例，标准全参数微调下需要GPU的参数量为=。 分层微调下 =.   两者的差距为 = 。 以使用AdamW优化器对LLaMA-7B单精度微调为例，模型参数 约为26.08G，理论上，7B 模型的, 和微调所需的 GPU 内存约为 104.32 GB。LLaMA-7B 有34 层(包括embedding层和head层）。可推断, 和所需GPU显存为约 31.13G。与标准的全参数微调相比节省GPU显存约73.19G。

大模型下混合精度的问题

图2为 使用不同优化器微调LLaMA2-7B模型是的GPU使用情况，需要说明的是，混合精度的结果是混合精度未适配分层微调下的结果。可以看出混合精度下的内存需求比不使用混合精度的更高。解释这一结果我们需要了解混合精度的原理。混合精度虽然会以半精度进行前向传播，但是在更新模型的时候，会备份出32位的权重进行梯度更新，同时梯度也会在32位环境下进行更新，原因是位解决数值下溢的问题，在训练后期，激活函数的梯度会非常小， 甚至在梯度乘以学习率后，值会更加小。如果利用 fp16 来进行参数更新的话，会出现舍入误差（fp16的表示范围有限，超过这个范围的值会被置为0）问题，导致更新无效。

问题来了，为什么混合精度能降低内存使用？结论是混合精度能大幅降低动态内存的使用，既主要是激活神经元的内存占用。以往的模型比如RoBERTa，GPT-2这些模型，模型的参数量较小（相比于大模型说），这个时候，模型占用的参数量十分有限，微调过程中使用内存较多的是和输入相关的动态内存部分。比如RoBERTa-base，以adamw 微调过程中（batch=8，max length=512)，模型参数，梯度，优化器状态占据的固定内存约为1.9G.  而动态内存部分约为5G.  小模型时代，微调过程中动态内存占主导，当使用混合精度微调时，由于前向传播半精度该部分内存会大幅降低，batch越大，句子长度越大，混合精度的降低的内存越多，可以抵消多出来的16为模型权重占据的内存。

大模型背景下，因为设备限制，很多情况下无法使用大的batch发挥出混合精度的优势。以65B模型为例，单精度仅模型参数需求内存约为242G，半精度的内存需求约为121G，混合精度下仅模型参数需求的内存为242+121=363G，只有当混合精度降低的动态内存超过121G时候，混合精度的优势才能体现出来，但是大多数情况下，设备的限制，无法使用大的batch size. 根据我们的实验结果看，当微调3B（GPT-Neo）左右的模型时候，在小的batch下（我们设置的batch size 为8，句子长度为512），混合精度已经没有内存优势。当然微调GPT-large时 混合精度仍然有优势。

混合精度适配分层策略

目前论文的结果是原始混合精度的结果，目前最新进展已将混合精度进行了分层适配，batch=1，句子长度为512下，实测结果7B模型大约混合精度下内存需求为16.87G，13B模型为31G。请持续关注我们的工作，代码基于基于hugging face，代码即将开源，可移植性友好，可与LORA等其PEFT的方式融合。

图3:（a）显示了RoBERTa-base在不同策略下的HiFT性能。B2U、T2D和RAN分别代表使用策略bottom2up、top2bottom和随机微调策略。(b)表示RoBERTa-base在不同分组设置下的HiFT性能。m表示分组设置中每组的层数。

图3显示了采用不同的微调策略和设置不同组队模型性能的影响。结论影响可以忽略不计。 一个有意思的现象是微调的顺序对模型性能几乎没有影响。

VS MeZO和LOMO

1. 我们的方法和MeZO的在下游任务的比较参请考原始论文，我们的方法比MEZO具有明显的性能优势。请持续关注我们的最新论文获取最新的结果。

2. 和LOMO相比，在不使用其它内存节省技术下，使用e2e 数据，对llama2-7B 测试，batch为1，最大句子长度为512。混合精度下，LOMO的峰值内存为21.57GB，HiFT为16.87GB；单精度下：LOMO 60.06GB，HiFT 29.73GB