GaLore

本仓库包含GaLore算法的预发布版本，该算法由GaLore: 通过梯度低秩投影实现内存高效的LLM训练提出。

梯度低秩投影（GaLore）是一种内存高效的低秩训练策略，它允许全参数学习，但比常见的低秩适应方法（如LoRA）更加内存高效。作为一种梯度投影方法，GaLore独立于优化器的选择，可以轻松地插入到现有的优化器中，只需两行代码，如下面的算法1所示。

讨论 (GaLore-Social)

我们欢迎任何关于GaLore的讨论、问题和反馈。请加入我们的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文件夹中。例如，要在C4上训练一个60m的模型，请执行以下操作：

# 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训练7B模型

要使用单个GPU（如NVIDIA RTX 4090）训练7B模型，你只需要指定--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: 通过梯度低秩投影实现高效大规模语言模型训练}, 
      author={赵嘉伟 and 张振宇 and 陈北笛 and 王张扬 and Anima Anandkumar and 田园东},
      year={2024},
      eprint={2403.03507},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}