GaLore

此仓库包含由 GaLore: 通过梯度低秩投影实现高效内存的大型语言模型训练 提出的 GaLore 算法预发布版本。

梯度低秩投影（GaLore）是一种高效的低秩训练策略，它支持全参数学习，同时比常见的低秩微调方法（如 LoRA）更加节省显存。作为一种梯度投影方法，GaLore 不依赖于优化器的选择，只需两行代码即可轻松集成到现有优化器中，具体如下面的算法 1 所示。

最新动态

• 2024-09-01: 我们正在开发 GaLore 2，这是一个更高效、更易用的 GaLore 版本。敬请期待！

• 2024-07-11: 我们发布了 Q-GaLore：采用 INT4 投影的量化 GaLore。[论文] [代码]

• 2024-07-01: GaLore 已被 ICML 2024 接受为口头报告！

• 2024-04-20: 欢迎加入我们的 Slack 工作区 GaLore-Social，与我们及社区成员一起交流讨论。

安装

安装 GaLore 优化器

通过 pip 安装：

pip install galore-torch

或者从源码安装：

git clone git@github.com:jiaweizzhao/GaLore.git
cd GaLore
pip install -e .

安装实验依赖

pip install -r exp_requirements.txt

我们的实验脚本已在 Python 3.8 和 PyTorch 2.1 环境下测试通过。

使用方法

使用 GaLore 优化器节省优化器内存

from galore_torch import GaLoreAdamW, GaLoreAdamW8bit, GaLoreAdafactor
# 定义参数组 galore_params 和 non_galore_params
param_groups = [{'params': non_galore_params}, 
                {'params': galore_params, 'rank': 128, 'update_proj_gap': 200, 'scale': 0.25, 'proj_type': 'std'}]
optimizer = GaLoreAdamW(param_groups, lr=0.01)

使用逐层权重更新节省权重梯度内存

我们使用 PyTorch（torch>=2.1.0）提供的 register_post_accumulate_grad_hook 来启用逐层权重更新。示例如下：

# 为每个参数 p 定义一个优化器，并存储在 optimizer_dict 中
for p in model.parameters():
    if p.requires_grad:
        optimizer_dict[p] = GaLoreAdamW([{'params': p, 'rank': 128, 'update_proj_gap': 200, 'scale': 0.25, 'proj_type': 'std'}], lr=0.01)

# 定义一个钩子函数，在反向传播过程中更新参数 p
def optimizer_hook(p):
    if p.grad is None: 
        return
    optimizer_dict[p].step()
    optimizer_dict[p].zero_grad()

# 将钩子注册到每一个参数上
for p in model.parameters():
    if p.requires_grad:
        p.register_post_accumulate_grad_hook(optimizer_hook)

更多细节请参阅 torchrun_main.py。

基准测试 1：在 C4 数据集上预训练 LLaMA 模型

torchrun_main.py 是使用 GaLore 在 C4 数据集上训练 LLaMA 模型的主要脚本。我们针对不同规模模型的基准测试脚本位于 scripts/benchmark_c4 文件夹中。 例如，要训练一个 6000 万参数的模型在 C4 数据集上，可以执行以下命令：

# LLaMA-60M，GaLore-Adam，1 张 A100 显卡，1 个节点
torchrun --standalone --nproc_per_node 1 torchrun_main.py \
    --model_config configs/llama_60m.json \
    --lr 0.01 \
    --galore_scale 0.25 \
    --rank 128 \
    --update_proj_gap 200 \
    --batch_size 256 \
    --total_batch_size 512 \
    --num_training_steps 10000 \
    --warmup_steps 1000 \
    --weight_decay 0 \
    --dtype bfloat16 \
    --eval_every 1000 \
    --optimizer galore_adamw 

使用单张 24GB 显存的 GPU 训练 70 亿参数模型

要使用单张 NVIDIA RTX 4090 等显卡训练 70 亿参数的模型，只需指定 --optimizer=galore_adamw8bit_per_layer，即可启用带有逐层权重更新的 GaLoreAdamW8bit。 配合激活检查点技术，您可以在 NVIDIA RTX 4090 上保持 16 的批量大小进行训练。

# LLaMA-7B，8 位 GaLore-Adam，单 GPU，激活检查点
# bsz=16，22.8G，
torchrun --standalone --nproc_per_node 1 torchrun_main.py \
    --model_config configs/llama_7b.json \
    --lr 0.005 \
    --galore_scale 0.25 \
    --rank 1024 \
    --update_proj_gap 500 \
    --batch_size 16 \
    --total_batch_size 512 \
    --activation_checkpointing \
    --num_training_steps 150000 \
    --warmup_steps 15000 \
    --weight_decay 0 \
    --grad_clipping 1.0 \
    --dtype bfloat16 \
    --eval_every 1000 \
    --single_gpu \
    --optimizer galore_adamw8bit_per_layer

目前，逐层权重更新技术仅支持单 GPU 训练（--single_gpu），不适用于 nn.parallel.DistributedDataParallel。我们正在努力实现多 GPU 训练时的逐层权重更新功能。

基准测试 2：在 GLUE 任务上微调 RoBERTa 模型

run_glue.py 是使用 GaLore 在 GLUE 任务上微调 RoBERTa 模型的主要脚本。示例脚本如下：

python run_glue.py \
    --model_name_or_path roberta-base \
    --task_name mrpc \
    --enable_galore \
    --lora_all_modules \
    --max_length 512 \
    --seed=1234 \
    --lora_r 4 \
    --galore_scale 4 \
    --per_device_train_batch_size 16 \
    --update_proj_gap 500 \
    --learning_rate 3e-5 \
    --num_train_epochs 30 \
    --output_dir results/ft/roberta_base/mrpc

引用

@misc{zhao2024galore,
      title={GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection}, 
      author={Jiawei Zhao and Zhenyu Zhang and Beidi Chen and Zhangyang Wang and Anima Anandkumar and Yuandong Tian},
      year={2024},
      eprint={2403.03507},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

GaLore 快速上手指南

GaLore (Gradient Low-Rank Projection) 是一种内存高效的低秩训练策略。它支持全参数学习，但比 LoRA 等常见低秩适配方法更节省显存。作为梯度投影方法，GaLore 独立于优化器选择，仅需少量代码即可集成到现有训练中。

环境准备

• 操作系统: Linux / macOS
• Python 版本: 推荐 Python 3.8+
• PyTorch 版本: 推荐 PyTorch 2.1+ (使用逐层权重更新功能必须 >= 2.1.0)
• 硬件要求:
  • 基础训练：任意支持 CUDA 的 GPU
  • 单卡训练 7B 模型：推荐 24GB 显存 (如 RTX 4090)，需配合激活重计算 (Activation Checkpointing)

安装步骤

方式一：通过 pip 安装（推荐）

pip install galore-torch

国内加速建议：如果下载缓慢，可使用清华或阿里镜像源：

pip install galore-torch -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

git clone git@github.com:jiaweizzhao/GaLore.git
cd GaLore
pip install -e .

安装实验依赖（可选）

如果您需要运行官方提供的基准测试脚本，请安装额外依赖：

pip install -r exp_requirements.txt

基本使用

1. 基础用法：替换优化器

最简单的使用方式是将原有的优化器（如 AdamW）替换为 GaLoreAdamW。您需要将参数分为两组：普通参数和需要应用 GaLore 的参数。

from galore_torch import GaLoreAdamW, GaLoreAdamW8bit, GaLoreAdafactor

# 定义参数组：non_galore_params 为普通参数，galore_params 为应用低秩投影的参数
param_groups = [
    {'params': non_galore_params}, 
    {'params': galore_params, 'rank': 128, 'update_proj_gap': 200, 'scale': 0.25, 'proj_type': 'std'}
]

# 初始化 GaLore 优化器
optimizer = GaLoreAdamW(param_groups, lr=0.01)

关键参数说明：

• rank: 低秩投影的秩（例如 128）。
• update_proj_gap: 更新投影矩阵的步数间隔（例如 200）。
• scale: 缩放因子（例如 0.25）。
• proj_type: 投影类型（通常为 'std'）。

2. 进阶用法：逐层权重更新（极致省显存）

为了进一步节省显存（例如在单张 24GB 显卡上训练 7B 模型），可以使用 PyTorch 的钩子机制实现逐层权重更新。这需要 PyTorch >= 2.1.0。

# 1. 为每个需要梯度的参数单独创建优化器
optimizer_dict = {}
for p in model.parameters():
    if p.requires_grad:
        optimizer_dict[p] = GaLoreAdamW(
            [{'params': p, 'rank': 128, 'update_proj_gap': 200, 'scale': 0.25, 'proj_type': 'std'}], 
            lr=0.01
        )

# 2. 定义反向传播时的钩子函数
def optimizer_hook(p):
    if p.grad is None: 
        return
    optimizer_dict[p].step()
    optimizer_dict[p].zero_grad()

# 3. 将钩子注册到每个参数上
for p in model.parameters():
    if p.requires_grad:
        p.register_post_accumulate_grad_hook(optimizer_hook)

注意：使用此模式时，通常不需要再调用全局的 optimizer.step()，因为更新已在反向传播过程中通过钩子完成。目前该技巧主要支持单卡训练模式。