GaLore：推进消费级硬件上的大模型训练

GaLore 在大型语言模型 (LLM) 训练中的集成标志着深度学习领域的一项重大进步，特别是在内存效率和人工智能研究的民主化方面。通过在消费级硬件上训练数十亿参数的模型，减少优化器状态的内存占用，并利用先进的投影矩阵技术，GaLore 为计算资源有限的研究人员和从业者开辟了新的视野。

利用消费级硬件扩展LLM

GaLore 在消费级 GPU（如 NVIDIA RTX 4090）上训练多达 70 亿参数模型（如基于 Llama 架构的模型）的能力是开创性的。这通过显著降低训练过程中优化器状态和梯度所需的内存要求来实现。该方法利用深度神经网络中梯度固有的低秩结构，应用投影来降低需要存储和操作的数据的维度。

优化器状态的内存效率

优化器状态，特别是在 Adam 等自适应优化算法中，在模型训练过程中占据了内存占用的大部分。GaLore 通过在梯度被优化器处理之前将其投影到较低维的子空间来解决这个问题。这不仅减少了存储这些状态所需的内存，还保持了优化过程的有效性。

内存节省是可观的，作者报告“**在训练过程中，优化器状态的内存减少了 82.5% 以上**”，使得在相同内存限制下训练更大的模型或使用更大的批次大小成为可能。当与 8 位精度优化器结合使用时，这些节省甚至会更加显著。

子空间切换和高级投影技术

GaLore 有效性的一个关键组成部分是其动态子空间切换机制，该机制允许模型在整个训练过程中在不同的低秩子空间中导航。这确保了模型不会局限于参数空间的有限部分，从而保留了全参数学习的能力。何时以及如何切换子空间的决定至关重要，这些切换的频率是保持一致的优化轨迹和适应梯度低秩结构不断演变的平衡。

动态调整这些投影以响应梯度结构变化的能力是 GaLore 武器库中的一个强大工具，它允许更细致地控制训练大型模型固有的内存优化权衡。

GaLore 与 8 位优化器结合

GaLore 与 8 位精度优化器的结合代表了一种协同效应，它最大限度地提高了内存效率，同时保持了训练过程的完整性和性能。8 位优化器通过量化优化器状态来减少内存占用。当与 GaLore 的投影机制结合使用时，结果是一个高度内存高效的训练方案，而不会损害模型精度或收敛速度。

这种组合在内存是关键瓶颈的场景中特别有效，例如在消费级硬件上训练大型模型或在内存受限的环境中部署模型。它可以在相同的硬件限制下使用更复杂的模型和更大的数据集，从而突破有限资源所能实现的界限。

实现细节

将 8 位优化器与 GaLore 集成以训练大型语言模型 (LLM) 涉及将梯度、权重和优化器状态量化为 8 位表示。此量化过程显著减少了内存占用，从而可以在相同内存限制下训练更大的模型或使用更大的批次大小。这种集成的算法细节涉及几个关键步骤，其中一些将受益于原生 CUDA 实现以提高效率。GaLore 为将这些技术更紧密地与量化和矩阵的专门参数化相结合开辟了新的可能性，这可以进一步减少内存使用。我们目前正在 bitsandbytes 库中探索这个方向。

GaLore 8 位优化算法概述

梯度投影：GaLore 使用投影矩阵将全精度梯度投影到低秩子空间中。此步骤减少了梯度的维度，然后将其量化为 8 位格式。

量化：投影梯度以及模型权重和优化器状态（例如 Adam 中的移动平均值）从 32 位浮点数量化为 8 位整数表示。这涉及将浮点值缩放到 8 位范围并将其四舍五入到最接近的整数。

优化器更新：8 位量化梯度用于更新模型权重。此步骤涉及将梯度反量化回浮点格式，应用优化器的更新规则（例如，Adam 的动量更新和参数调整），然后将更新的优化器状态量化回 8 位以进行存储。

反量化和权重更新：8 位量化权重经过反量化为浮点表示以进行处理，尽管由于值的有限范围，它们保留了量化形式固有的 8 位精度。此步骤是必需的，因为 PyTorch 等框架中的标准操作不支持 8 位整数，并且此类整数权重无法容纳梯度。虽然此方法本质上不提高准确性，但它有助于在当前深度学习库的限制内实际应用和计算量化权重的梯度。请注意，在反量化和应用权重更新之后，GaLore 会再使用一次投影，将反量化的低秩更新投影回原始空间。

与 Hugging Face Transformers 一起使用

要将 GaLore 优化器与 Hugging Face transformers 库一起使用，您需要先将其更新到支持 GaLore 优化器的版本，可以通过安装最新更新，即 `pip install transformers>=4.39.0`，或从源代码安装 transformers。

然后使用 `pip install galore-torch` 安装 galore-torch 库。下面是 GaLore 与 transformers 配合使用的完整示例，用于在 imdb 数据集上预训练 Mistral-7B。

import torch
import datasets
from transformers import TrainingArguments, AutoConfig, AutoTokenizer, AutoModelForCausalLM
import trl

train_dataset = datasets.load_dataset('imdb', split='train')

args = TrainingArguments(
    output_dir="./test-galore",
    max_steps=100,
    per_device_train_batch_size=2,
    optim="galore_adamw",
    optim_target_modules=["attn", "mlp"]
)

model_id = "mistralai/Mistral-7B-v0.1"

config = AutoConfig.from_pretrained(model_id)

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_config(config).to(0)

trainer = trl.SFTTrainer(
    model=model, 
    args=args,
    train_dataset=train_dataset,
    dataset_text_field='text',
    max_seq_length=512,
)

trainer.train()

`TrainingArguments`：只需传递一个有效的 `optim_target_modules`（它支持单个字符串、正则表达式或字符串或正则表达式列表），以及一个有效的 GaLore 优化器（例如 `galore_adamw`、`galore_adamw_8bit`、`galore_adafactor`）作为 `optim`，即可开始使用！

逐层更新

另一个需要提及的重要点是*逐层*优化器（即一次只更新一层权重）。通常，优化器在反向传播后对所有层执行一次权重更新。这通过在内存中存储整个权重梯度来完成。通过采用逐层权重更新，我们可以进一步减少训练期间的内存占用。在底层，这是通过 PyTorch 在用户想要更新的层上使用后累积钩子实现的。

要使用此功能，只需在优化器名称后附加 `_layerwise`，例如 `galore_adamw_layerwise`。

结论

GaLore 凭借其利用梯度低秩结构的创新方法，代表了 LLM 内存高效训练方面向前迈出的重要一步。通过在消费级硬件上训练数十亿参数的模型，通过投影技术减少优化器状态的内存占用，并允许动态子空间切换，GaLore 民主化了大规模模型训练的访问。GaLore 与 8 位精度优化器的兼容性进一步增强了其实用性，为在不需要专门计算资源的情况下训练更大、更复杂的模型提供了途径。这为人工智能的研究和应用开辟了新的可能性，使其成为从业者和研究人员都感到兴奋的时刻。

资源

请参阅原始论文。Twitter 参考：1 2 3。该论文还比较了 GaLore 和 ReLoRA，这可能引起一些读者的兴趣。对于在阅读论文后仍有疑问或希望建设性地讨论结果的读者，请随时加入作者的 Slack 社区。对于对类似发布感兴趣的读者，请关注Jiawei Zhao和Titus von Koeller（获取最新 `bitsandbytes` 版本信息）以及Younes Belkada（获取 Hugging Face 生态系统内外的量化相关最新信息）。