PEFT 中的 UniLoRA 实战指南:用共享向量库把 LoRA 参数量压到极限
2026/9/20 11:00:52 网站建设 项目流程
  • 人工智能
  • 大模型
  • 微调
  • LoRA

【免费下载链接】peft

🤗 PEFT: State-of-the-art Parameter-Efficient Fine-Tuning.

项目地址:https://gitcode.com/gh_mirrors/pe/peft
点击查看免费下载

导读

UniLoRA(Uni-LoRA: One Vector Is All You Need)是 🤗 PEFT 中一种参数高效微调(PEFT)方法,核心思想是让所有低秩适配器层共享一个紧凑的可训练向量库theta_d,并通过确定性投影把每一层的低秩矩阵映射到该共享向量上,从而把可训练参数量压缩到只剩theta_d_length个标量。本文基于仓库中的 examples/unilora_finetuning/README.md 及其配套脚本 unilora_finetuning.py,从 Quick Start、MetaMathQA 微调命令、配置参数详解到源码级实现原理(索引生成、归一化缩放、前向计算与合并),带你完整掌握在 PEFT 中使用 UniLoRA 的实战方案。

UniLoRA 是什么:一个向量搞定低秩适配

UniLoRA 的设计理念非常直接:既然 LoRA 每一层都要维护一组AB低秩矩阵(可训练参数与层数和 rank 成正比),那么能不能让所有层共享同一组可训练参数,再用廉价的确定性投影把它"展开"成每层需要的低秩矩阵?

答案就是共享向量库(vector bank)theta_d

  • 整个模型只维护一个长度为theta_d_length的可训练向量theta_d,该向量以nn.ParameterDict的形式挂在模型顶层(源码见 src/peft/tuners/unilora/model.py),所有被替换的线性层引用同一个底层张量
  • 每一层的低秩矩阵不再直接存参数,而是通过确定性生成的索引unilora_indices_A/unilora_indices_B)从theta_d中取值,再乘以归一化缩放因子(unilora_scales_A/unilora_scales_B)还原为等效的AB矩阵:
    • A = theta_d[indices_A] * scales_A
    • B = theta_d[indices_B] * scales_B
  • 前向更新等价于ΔW = B @ A,因此可以沿用熟悉的 PEFT 训练流程(get_peft_model+ Trainer 家族),无需改变任何训练习惯,即可大幅降低适配器参数量。

参数共享这一事实在测试 tests/test_unilora.py 中有明确断言:lin0lin1lin3各层的unilora_theta_d["default"]通过data_ptr()校验指向同一块内存,印证了"共享向量库"的实现细节。

UniLoraConfig 配置参数全解

UniLoRA 的配置类为UniLoraConfig,定义在 src/peft/tuners/unilora/config.py,继承自PeftConfig,并在__post_init__中完成参数校验与peft_type = UNILORA的设定。它"刻意保持精简,只包含当前实现真正用到的参数"(见类 docstring)。全部字段如下:

参数默认值说明
r4低秩适配的 rank,控制 UniLoRA 更新的表达能力;必须为正整数
proj_seed42生成固定索引分配所用的随机种子,保证跨运行可复现
theta_d_length256共享向量theta_d的长度;越大可训练参数越多,同时共享参数空间的"碰撞"越少
target_modulesNone应用 UniLoRA 的目标模块名。传字符串按正则匹配;传列表按精确名/后缀匹配;特殊值'all-linear'匹配除输出层外的所有 Linear/Conv1D 层。未指定时按模型架构推断(见下)
unilora_dropout0.0UniLoRA 层内部的 dropout 概率,取值范围[0.0, 1.0]
fan_in_fan_outFalse目标层是否以(fan_in, fan_out)存储权重,GPT-2 等使用 Conv1D 的模型需设为True
bias"none"可训练 bias 策略:'none'不更新任何 bias;'all'更新全部 bias;'unilora_only'仅更新 UniLoRA 层内部 bias。注意开启 bias 更新会在禁用适配器时也改变模型输出
modules_to_saveNone除 UniLoRA 层外需要保持可训练并随 checkpoint 保存的模块,常用于任务头(如分类头)
init_theta_d_bound0.02theta_d的初始化边界,向量从[-bound, bound]均匀采样;避免零初始化导致梯度消失,建议用较小的值保证训练稳定
init_weightsTrue是否使用 UniLoRA 默认初始化;设为Falsetheta_d保持随机(torch.randn)初始化
save_indicesFalse是否将生成的索引与缩放 buffer 一并保存进 checkpoint。开启会增加体积,但让保存的适配器不再依赖未来索引生成逻辑的变更
layers_to_transformNone仅对指定编号的 transformer 层应用 UniLoRA;仅在target_modules为列表时可用
layers_patternNone配合layers_to_transform使用的自定义层名模式,适用于层命名不遵循常规的模型;同样仅在target_modules为列表时可用

__post_init__中的校验规则(config.py)包括:

  • rtheta_d_length必须为正整数;
  • unilora_dropout必须在[0.0, 1.0]
  • init_theta_d_bound必须为正浮点数;
  • bias只能是'none''all''unilora_only'三者之一;
  • target_modules为字符串时不允许同时指定layers_to_transformlayers_pattern

target_modules未显式给出时,会从TRANSFORMERS_MODELS_TO_UNILORA_TARGET_MODULES_MAPPING(定义于 src/peft/utils/constants.py,即 LoRA 目标模块映射的副本)按模型架构推断默认目标模块;若架构不受支持则会报错。

快速开始:用 SFTTrainer 微调 Llama-3.2-3B

原文档给出的 Quick Start 与经典 PEFT 流程几乎一致,只需把LoraConfig换成UniLoraConfig。完整示例(原文未删减):

import torch from datasets import load_dataset from peft import UniLoraConfig, get_peft_model from transformers import AutoModelForCausalLM, AutoTokenizer from trl import SFTConfig, SFTTrainer model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B", dtype=torch.bfloat16, device_map="auto") tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B") tokenizer.pad_token_id = tokenizer.eos_token_id config = UniLoraConfig( r=32, theta_d_length=256, proj_seed=42, target_modules=["q_proj", "v_proj"], unilora_dropout=0.0, task_type="CAUSAL_LM", ) peft_model = get_peft_model(model, config) peft_model.print_trainable_parameters() dataset = load_dataset("imdb", split="train[:1%]") training_args = SFTConfig(dataset_text_field="text", max_length=128) trainer = SFTTrainer( model=peft_model, args=training_args, train_dataset=dataset, processing_class=tokenizer, ) trainer.train() peft_model.save_pretrained("unilora-llama-3.2-3b")

这里需要注意两点:

  1. theta_d_length就是参数量:上例中即使r=32且模型有几十个目标层,可训练参数也只有theta_d_length = 256个标量(外加 bias 策略为"none"时不产生额外参数)。测试 tests/test_unilora.py 印证了theta_d.shape == (theta_d_length,),并在注释中指出"theta_d_length实际上充当共享参数池(码本)的大小"。
  2. 索引取值范围indices中的值域为[0, theta_d_length)(见 tests/test_unilora.py 注释),这意味着索引所需存储可以非常紧凑,同时缩放因子用于抵消不同索引被使用次数不均带来的数值尺度差异。

微调完成后,用标准的PeftModel加载适配器(保持与训练时一致的 dtype 与device_map):

import torch from peft import PeftModel from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3.2-3B", dtype=torch.bfloat16, device_map="auto" ) peft_model = PeftModel.from_pretrained(model, "unilora-llama-3.2-3b")

在 MetaMathQA 上微调:完整命令行方案

原文档提供了针对数学推理数据集 MetaMathQA 的完整命令行微调脚本,直接使用仓库中的 unilora_finetuning.py:

python unilora_finetuning.py \ --base_model_name_or_path meta-llama/Llama-3.2-3B \ --output_dir output/unilora-llama-3.2-3b-metamath \ --unilora_r 32 \ --theta_d_length 256 \ --proj_seed 42 \ --unilora_dropout 0.0 \ --bits bf16 \ --data_path meta-math/MetaMathQA \ --dataset_split train[:100000] \ --dataset_field query response \ --bf16 True \ --num_train_epochs 1 \ --per_device_train_batch_size 2 \ --gradient_accumulation_steps 8 \ --save_strategy steps \ --save_steps 1000 \ --save_total_limit 1 \ --logging_steps 1 \ --learning_rate 1e-4 \ --weight_decay 0. \ --warmup_steps 0.03 \ --tf32 True \ --report_to none

该脚本继承SFTConfig构建ScriptArguments(unilora_finetuning.py),因此上述参数中--num_train_epochs--per_device_train_batch_size--gradient_accumulation_steps--save_strategy--learning_rate等均为 TRL SFT 的标准训练超参数。此外还有一组 UniLoRA 专属与数据相关的参数值得说明:

参数默认值说明
--base_model_name_or_pathNone(必填)基础模型名或本地路径,支持 fp32/fp16/bf16 权重
--bitsbf16模型加载精度:bf16/fp16/fp32,由get_dtype()映射(unilora_finetuning.py)
--unilora_r32UniLoRA 适配器 rank
--theta_d_length256共享theta_d向量库长度
--proj_seed42确定性投影索引的种子
--unilora_dropout0.0UniLoRA 层 dropout 概率
--init_weightsTrue是否启用 UniLoRA 专属初始化
--target_modules7 个默认模块q_proj, o_proj, k_proj, v_proj, gate_proj, up_proj, down_proj
--merge_and_saveFalse训练结束后把适配器合并进基础模型并另存一份
--data_pathimdb训练数据集路径或 Hub id
--dataset_splittrain[:1%]数据集划分
--dataset_fieldNone指令数据的输入/输出字段名;提供后脚本会格式化为### USER: {query}\n### ASSISTANT: {response}模板(unilora_finetuning.py)

脚本的数据流(unilora_finetuning.py)为:加载模型 → 设置 pad token → 组装UniLoraConfigbias="none"task_type="CAUSAL_LM")→get_peft_model并打印可训练参数 → 加载并格式化数据集 →SFTTrainer训练 →save_state→ 保存适配器到output_dir/unilora_ft。若开启--merge_and_save,还会调用merge_and_unload()生成合并模型并保存到output_dir/unilora_merged

源码级原理:确定性投影与归一化缩放

1. 平衡的确定性索引分配

每一层最初通过torch.randint(0, theta_d_length, ...)生成随机索引(indices_A形状为(r, in_features)indices_B形状为(out_features, r),见 layer.py),随后UniLoraModel._assign_unilora_indices_and_scales会在全模型范围内重新分配索引(model.py)。

关键函数是静态方法generate_index(model.py)。其 docstring 明确指出:如果简单地用np.random.choice独立采样每个位置,对于较小的适配器会有一部分theta_d条目永远不被用到;UniLoRA 改为平衡的确定性分配——把展平后的 LoRA 参数总数D平均分到d = theta_d_length个索引上,每个索引恰好出现floor(D/d)ceil(D/d)次,再用proj_seed打乱顺序。这样既保证所有共享向量都被使用,也让后续归一化保持稳定,同时完全可复现。

2. 基于使用频率的逆平方根缩放

索引分配完毕后,_assign_unilora_indices_and_scales会统计每个索引被引用的次数counts = torch.bincount(indices, minlength=theta_d_length),然后计算inv_sqrt_counts = counts.rsqrt()(未使用的索引缩放为 0),并据此为每一层的scales_A/scales_B赋值(model.py)。这一归一化机制保证:即使不同索引被不同数量的低秩位置共享,聚合后的更新幅度也不会失衡——这正是共享码本方法能稳定训练的关键。

3. 前向计算与权重合并

在 layer.py 的_get_lora_matrices中,先按索引取出theta_d的切片并乘以对应缩放,得到等效的AB矩阵;forward(layer.py)随后执行:

result = result + F.linear(F.linear(dropout(x), A), B)

即先降维(A)再升维(B),与 LoRA 的低秩通路完全等价。与标准 LoRA 一样,UniLoRA 也支持merge()/unmerge()get_delta_weight计算ΔW = transpose(B @ A)(在 CPU 上且权重为 fp16/bf16 时会先转 fp32 计算再转回,避免精度问题),merge将其累加到基础层权重,safe_merge模式下还会在合并前后检查 NaN(layer.py)。层实现兼容torch.nn.Lineartransformers.pytorch_utils.Conv1D(GPT-2 等模型),后者会自动强制fan_in_fan_out=True(model.py)。

4. 共享参数的设备与梯度管理

由于所有层共享同一个theta_d参数,UniLoraModel通过nn.ParameterDict统一持有它(prefix = "unilora_"uses_shared_state = True,见 model.py),state_dict 中对应的键为base_model.unilora_theta_d.<adapter_name>(model.py)。set_adapterenable_adapters会同步管理theta_drequires_grad,确保切换适配器或禁用适配器时行为正确(layer.py)。peft_model.py中的特殊处理(src/peft/peft_model.py)也说明unilora_theta_d是"跨层共享"参数,在按层遍历时不会被重复归属。

进阶用法与注意事项

save_indices:让 checkpoint 更稳健

默认save_indices=False时,索引与缩放 buffer 为非持久化 buffer,不写入 state_dict,加载时按proj_seed重新生成(layer.py),因此 checkpoint 极小(只有theta_d本体)。测试中的_get_unilora_index_state也专门注释了"索引是非持久化 buffer,需从模块收集而非 state_dict"(tests/test_unilora.py)。若未来索引生成算法变更可能影响复现,可将save_indices=True,此时索引与缩放会随适配器一起保存,代价是 checkpoint 体积增大。

初始化选择

默认init_weights=True时,theta_d[-init_theta_d_bound, init_theta_d_bound]均匀采样(默认边界0.02),目的是避免零初始化——注意get_delta_weighttheta_d二次依赖(B @ A中两个因子都来自theta_d),零初始化会直接导致梯度消失(测试 tests/test_custom_models.py 注释也说明 UniLoRA 默认不能用零初始化)。设为False则改用torch.randn初始化。

模块定位与层级裁剪

target_modules支持列表、正则字符串与'all-linear';配合layers_to_transform/layers_pattern可以只改造部分 transformer 层。回归测试 tests/regression/test_state_dict.py 使用theta_d_length=257验证了非默认长度的 state_dict 往返一致性。

低层 API 支持

UniLoRA 同样可通过低层 API 手动创建(见 tests/test_low_level_api.py 中UniLoraConfig, {"r": 2, "theta_d_length": 16}的用法),并支持convert_to_lora转换(tests/test_unilora.py),便于与其他 PEFT 方法互操作。

总结

UniLoRA 用一个长度为theta_d_length的共享向量库替代了所有层的 LoRA 低秩参数:通过proj_seed控制的确定性平衡索引分配、基于使用频率的逆平方根缩放、以及完全等价的ΔW = B @ A前向通路,在保持熟悉 PEFT 训练流程的同时把可训练参数量削减到一个极小的常量级别。在 PEFT 中使用它只需三步:构造UniLoraConfigget_peft_model包装模型、用任意 Trainer 训练;进阶场景可结合save_indicesmerge_and_savelayers_to_transform等选项灵活控制 checkpoint 体积、推理形态与改造范围。相关实现与测试分别位于 src/peft/tuners/unilora/、tests/test_unilora.py 与 docs/source/package_reference/unilora.md,可供进一步深入阅读。

  • 人工智能
  • 大模型
  • 微调
  • LoRA

【免费下载链接】peft

🤗 PEFT: State-of-the-art Parameter-Efficient Fine-Tuning.

项目地址:https://gitcode.com/gh_mirrors/pe/peft
点击查看免费下载

相关推荐

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询