MoRA: 用于参数高效微调的高阶更新

设置

我们在基于 HF peft 的 peft-mora 中实现了 MoRA，主要在 apply_mora 和 get_delta_weight 中实现。

pip install -e ./peft-mora

安装后，可以这样使用：

from peft import LoraConfig, get_peft_model
config = LoraConfig(
    # 启用 MoRA
    use_mora=True,
    # 类型 1（共享）用于大型 lora 秩，论文中的公式 6
    # 类型 6（基于 RoPE）用于小型 lora 秩，论文中的公式 9
    mora_type=6,
    # 这里是 lora 秩，我们将计算 MoRA 中相应的 $\hat{r}$
    r=lora_r,
    # MoRA 不使用 lora_alpha
    # lora_alpha=lora_alpha,
    target_modules=lora_target_modules,
    lora_dropout=lora_dropout,
    task_type="CAUSAL_LM",
    **kwargs,
)
model = get_peft_model(model, config)

# 在这里进行训练...

# 可以像 LoRA 一样通过 `merge_and_unload` 合并到模型中
model = model.merge_and_unload()

示例

使用 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 步合并一次