每日论文速递 | GaLore: 使用梯度低秩映射进行大模型 Memory-Efficient 全参训练

深度学习自然语言处理 分享 整理：pp

摘要：训练大型语言模型（LLMs）面临着显著的内存挑战，主要是由于权重和优化器状态的不断增大。常见的内存降低方法，如低秩适应（LoRA），在每一层中向冻结的预训练权重添加一个可训练的低秩矩阵，从而减少可训练参数和优化器状态。然而，这些方法通常在预训练和微调阶段的性能上都不如使用全秩权重训练，因为它们将参数搜索限制在低秩子空间中，改变了训练动态，并且可能需要全秩热启动。在这项工作中，我们提出了Gradient Low-Rank Projection（GaLore），一种允许全参数学习但比LoRA等常见低秩适应方法更节省内存的训练策略。我们的方法在优化器状态的内存使用上最多减少了65.5%，同时在使用C4数据集进行LLaMA 1B和7B架构的预训练以及在GLUE任务上对RoBERTa进行微调时，保持了效率和性能。我们的8位GaLore相较于BF16基准，将优化器内存进一步降低了82.5%，总训练内存降低了63.3%。值得注意的是，我们首次证明了在具有24GB内存的消费级GPU上（例如NVIDIA RTX 4090）进行7B模型的预训练是可行的，而无需模型并行、检查点策略或卸载策略。

https://arxiv.org/abs/2403.03507

Q1: 这篇论文试图解决什么问题？

A：这篇论文提出了一种名为GaLore（Gradient Low-Rank Projection）的训练策略，旨在解决大型语言模型（LLMs）训练过程中显著的内存挑战。这些挑战主要是由于权重和优化器状态的不断增长所导致的。传统的内存减少方法，如低秩适应（LoRA），通过在每层的冻结预训练权重上添加可训练的低秩矩阵来减少可训练参数和优化器状态，但这些方法通常在预训练和微调阶段的性能上不如全秩权重训练，因为它们限制了参数搜索到低秩子空间并改变了训练动态。此外，这些方法可能需要全秩的热启动。

GaLore的核心思想是利用权重矩阵梯度的低秩结构，而不是尝试将权重矩阵本身近似为低秩。这种方法在保持全参数学习的同时，比常见的低秩适应方法（如LoRA）更节省内存。论文展示了GaLore在LLaMA 1B和7B架构上的预训练以及在GLUE任务上对RoBERTa进行微调时，能够有效减少内存使用，同时保持训练的效率和性能。特别是，GaLore首次证明了在没有模型并行、检查点或卸载策略的情况下，可以在具有24GB内存的消费级GPU（如NVIDIA RTX 4090）上预训练7B模型的可行性。

Q2: 有哪些相关研究？

A：在这篇论文中，作者提到了以下几个与GaLore相关的研究领域和具体工作：

1. 低秩适应（Low-Rank Adaptation, LoRA）：LoRA是一种流行的技术，用于在不调整所有模型参数的情况下，高效地适应预训练语言模型（PLMs）到不同的下游应用。LoRA通过将权重矩阵重参数化为低秩适配器的形式，从而减少可训练参数和优化器状态的数量。
2. ReLoRA：ReLoRA是LoRA的一个变体，用于预训练，它通过定期使用之前学习的低秩适配器更新冻结的预训练权重。
3. 子空间学习（Subspace Learning）：一些研究表明，学习主要发生在一个显著低维的参数子空间内。这些发现促进了一种特殊类型的学习，称为子空间学习，其中模型权重在低秩子空间内被优化。
4. 投影梯度下降（Projected Gradient Descent, PGD）：GaLore与传统的PGD相关，但GaLore考虑了在训练多层神经网络时自然出现的特定梯度形式，并证明了其许多属性。
5. 内存高效的优化（Memory-Efficient Optimization）：一些工作尝试减少自适应优化算法的梯度统计的内存成本。例如，Adafactor通过分解二阶统计量来实现亚线性内存成本。
6. 量化（Quantization）：量化广泛用于减少优化器状态的内存成本。

7.** 融合梯度计算（Fused Gradient Computation）**：Lv等人提出了融合梯度计算以减少训练期间存储权重梯度的内存成本。

这些相关工作为GaLore提供了理论基础和实践背景，GaLore在这些方法的基础上，通过利用梯度的低秩结构来进一步减少内存使用，同时保持训练的效率和性能。

Q3: 论文如何解决这个问题？

A：论文通过提出一种名为GaLore（Gradient Low-Rank Projection）的训练策略来解决大型语言模型（LLMs）训练中的内存效率问题。GaLore的核心思想是利用梯度的低秩结构，而不是直接对权重矩阵进行低秩近似。以下是GaLore解决这个问题的关键步骤：

1. 理论基础：首先，论文理论上证明了在训练过程中，梯度矩阵G会变得低秩。这是通过分析梯度更新的数学形式得出的，特别是在可逆网络和softmax损失函数的情况下。
2. 梯度低秩投影：GaLore计算两个投影矩阵P和Q，将梯度矩阵G投影到低秩形式P'GQ。这样，依赖于梯度统计量的优化器状态的内存成本可以大幅降低。偶尔更新P和Q（例如，每200次迭代一次）会产生最小的额外计算成本。
3. 内存效率：GaLore在实践中显示出比LoRA更高的内存效率，这在预训练过程中可以减少高达30%的内存使用。
4. 预训练和微调：**GaLore在LLM的预训练和微调中都表现良好。在预训练LLaMA 7B模型时，结合8位优化器和逐层权重更新技术，GaLore在保持与全秩对应模型相当的性能的同时，优化器状态的内存成本降低了不到10%**。
5. 无需全秩热启动：与ReLoRA不同，GaLore在预训练过程中保持低内存，不需要全秩模型训练热启动。
6. 兼容性：GaLore与现有的内存高效优化技术兼容，例如8位优化器和逐层权重更新。这使得GaLore可以与现有的优化器（如AdamW、8位Adam和Adafactor）轻松集成。
7. 实验验证：通过在LLaMA模型上的预训练和在GLUE基准上的微调，论文展示了GaLore在不同模型大小和任务上的有效性。特别是在预训练7B模型时，GaLore能够在消费级GPU上实现，这在之前是不可能的。
8. 内存和吞吐量测量：论文还提供了GaLore在实际训练中的内存使用和吞吐量测量，以证明其在实际应用中的可行性。

通过这些方法，GaLore不仅减少了内存使用，而且保持了训练的效率和模型性能，使得在资源有限的硬件上训练大型语言模型成为可能。

Q4: 论文做了哪些实验？

A：论文中进行了一系列的实验来验证GaLore方法的有效性。以下是主要的实验内容：

1. 预训练性能比较：在C4数据集上，对不同大小的LLaMA模型进行预训练，比较GaLore与其他低秩方法（如LoRA和ReLoRA）以及全秩训练方法的性能。实验结果包括验证困惑度（perplexity）和内存估计。
2. 与内存高效优化器结合：将GaLore应用于不同的优化器，如AdamW、8位Adam和Adafactor，以进一步减少内存占用，并评估其对训练性能的影响。
3. 扩展到LLaMA 7B架构：在LLaMA 7B架构上进行预训练，使用8位GaLore与8位Adam进行比较，以展示GaLore在大规模模型上的应用能力。
4. 内存效率微调：在GLUE基准测试上对预训练的RoBERTa模型进行微调，比较GaLore与全秩微调和LoRA的性能和内存占用。
5. 内存和吞吐量测量：在NVIDIA A100 GPU上，对LLaMA 1B模型进行预训练，测量不同方法的实际内存占用和训练吞吐量。
6. 消融研究：进行消融研究以理解GaLore中不同超参数的影响，例如子空间更新频率和子空间的秩。
7. 理论分析：提供了理论证明，包括梯度变得低秩的证明，以及GaLore收敛性的分析。

这些实验旨在全面评估GaLore在不同规模的模型、不同的优化器以及不同的训练阶段（预训练和微调）中的性能和内存效率。实验结果表明，GaLore能够在显著减少内存使用的同时，保持或甚至提高训练效率和模型性能。

Q5: 有什么可以进一步探索的点？

A：论文中提出了几个可以进一步探索的开放性问题和潜在的研究方向：

其他类型的模型：GaLore目前主要针对大型语言模型。探索将其应用于其他类型的模型，如视觉Transformer和扩散模型，可能会揭示新的应用场景和挑战。

内存效率的进一步改进：研究如何通过量化或特殊的参数化技术进一步减少投影矩阵的内存占用，以提高内存效率。

弹性数据分布式训练：探索在低带宽的消费级硬件上进行弹性数据分布式训练的可能性，这可能需要对GaLore的实现进行优化，以适应资源受限的环境。

子空间更新策略：虽然论文中提到了子空间更新频率对训练性能的影响，但最佳的更新策略可能依赖于模型的大小、复杂性和训练任务。进一步研究如何动态调整子空间更新策略可能会提高训练效率。

理论分析的深入：对GaLore的收敛性和稳定性进行更深入的理论分析，特别是在不同的网络架构和损失函数下。

实际部署和应用：将GaLore应用于实际的NLP任务和生产环境中，评估其在实际应用中的性能和可扩展性。

环境影响：研究GaLore在减少大型模型训练环境影响方面的潜力，例如通过降低能源消耗和碳足迹。

与其他优化技术的结合：探索GaLore与其他内存优化技术（如梯度累积、模型并行化等）的结合，以进一步提高训练效率。

这些研究方向不仅有助于提高GaLore的性能和适用性，还可能推动整个领域在内存高效训练策略方面的进步。

Q6: 总结一下论文的主要内容？

A：这篇论文的主要内容可以总结如下：

1. 问题定义：论文针对大型语言模型（LLMs）训练过程中的内存挑战，特别是权重和优化器状态的不断增长导致的内存需求问题。
2. 现有方法的局限性：传统的低秩适应方法（如LoRA）在预训练和微调阶段通常性能不佳，因为它们限制了参数搜索空间并改变了训练动态。
3. GaLore方法：提出了一种新的训练策略GaLore，它通过利用梯度的低秩结构来减少内存使用，同时允许全参数学习，提高了内存效率。
4. 理论基础：论文提供了理论分析，证明了在训练过程中梯度矩阵会变得低秩，并提出了GaLore的数学定义和收敛性分析。
5. 实验验证：通过在LLaMA模型上的预训练和在GLUE基准上的微调实验，展示了GaLore在不同模型大小和任务上的有效性。实验结果表明，GaLore在减少内存使用的同时，保持了训练的效率和模型性能。
6. 内存和吞吐量测量：论文还提供了GaLore在实际训练中的内存使用和吞吐量测量，证明了其在资源受限的硬件上的可行性。
7. 开放性问题和未来工作：论文提出了一些开放性问题，如将GaLore应用于其他类型的模型、进一步改进内存效率、在低带宽硬件上进行弹性数据分布式训练等，为未来的研究提供了方向。

总的来说，这篇论文提出了一种创新的内存高效训练策略GaLore，通过在保持全参数学习的同时显著减少内存使用，为在资源受限的环境中训练大型语言模型提供了一种有效的解决方案。

以上内容均由KimiChat生成，深入了解论文内容仍需精读论文