MoRA: 高效参数微调的高阶更新方法

MoRA

MoRA: 开启参数高效微调的新纪元

在深度学习领域,如何高效地微调大型预训练模型一直是一个热门话题。近日,来自康奈尔大学和微软研究院的研究人员提出了一种新的参数高效微调方法 - MoRA(High-Rank Updating for Parameter-Efficient Fine-Tuning)。这种方法通过高阶矩阵分解实现了高效的参数更新,在各种下游任务中都表现出色。让我们一起深入了解这项创新技术。

MoRA的核心思想

MoRA的核心思想是利用高阶矩阵分解来实现参数的高效更新。传统的低秩适应(LoRA)方法虽然可以减少可训练参数的数量,但其表达能力有限。MoRA通过引入高阶分解,大大提高了参数更新的灵活性和表达能力,同时保持了计算效率。

具体来说,MoRA使用了两种主要的更新策略:

共享型更新(Sharing): 适用于较大的LoRA秩,对应论文中的公式6。
基于RoPE的更新: 适用于较小的LoRA秩,对应论文中的公式9。

这两种策略都能在保持参数效率的同时,提供更强的表达能力。

MoRA的实现与使用

MoRA的实现基于Hugging Face的PEFT库。研究人员在GitHub上提供了完整的代码实现(https://github.com/kongds/MoRA)。使用MoRA非常简单,只需要在LoRA配置中启用相关选项即可:

from peft import LoraConfig, get_peft_model

config = LoraConfig(
    use_mora=True,  # 启用MoRA
    mora_type=6,    # 选择MoRA类型
    r=lora_r,       # 设置LoRA秩
    target_modules=lora_target_modules,
    lora_dropout=lora_dropout,
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, config)

值得注意的是,MoRA不使用lora_alpha参数,这是它与传统LoRA的一个区别。

MoRA在实践中的应用

研究人员在多个任务上对MoRA进行了测试,包括语言模型微调和预训练。以下是两个典型的应用场景:

微调MetaMath模型:

RANK=8
deepspeed --num_gpus=8 --num_nodes=2 train.py \
    --base_model <LLAMA-2> --micro_batch_size 4 \
    --wandb_run_name mora_math_r8 \
    --lora_target_modules q_proj,k_proj,v_proj,o_proj,gate_proj,down_proj,up_proj \
    --num_epochs 3 --deepspeed ds.config --wandb_project lora-math --lora_r $RANK \
    --batch_size 128 --data_path meta-math/MetaMath --save_steps 3000 \
    --learning_rate 3e-4 --mora_type 6 --logging_steps 5 --use_bf16 --use_16bit --use_mora

预训练场景:

deepspeed --num_gpus=8 --num_nodes=4 train.py \
    --micro_batch_size 16 --wandb_run_name mora-pretrain250m-r128 \
    --num_epochs 1 --wandb_project lora-pretrain --batch_size 1024 \
    --data_path <processed C4> --logging_steps 1 \
    --lora_target_modules q_proj,k_proj,v_proj,o_proj,gate_proj,down_proj,up_proj \
    --lora_r 128 --lora_alpha 64 --warmup_steps 1000 \
    --force_tqdm_update --lr_scheduler_type cosine \
    --max_steps 10000 --pretrain 250m \
    --train_embhead --learning_rate 5e-4 \
    --use_mora --use_relora --use_relora_step 2000  # ReMoRA每2000步合并一次

这些示例展示了MoRA在不同规模和类型的任务中的灵活应用。

MoRA architecture