Chapter 07

PEFT in NeMo:LoRA / Adapter / IA³

📌 commit 2ea3e0fNeMo 自家的 PEFT 实现

NeMo 不直接用 HF PEFT,而是自己实现了一套兼容 Megatron-Core ParallelLinear 的 PEFT 子系统。本章讲它的设计、用法、与 HF PEFT 的差别。

7.1为什么 NeMo 不用 HF PEFT

NeMo 模型的 Linear 用的是 Megatron-Core 的 ColumnParallelLinear / RowParallelLinear(带 TP 切分)。HF PEFT 假设 nn.Linear,套上来会出错——LoRA 的 BA 矩阵也需要按 TP 切分。NeMo 自己写了一套 ParallelLinear-aware 的 PEFT。

7.2支持的 PEFT 方法

方法NeMo 类类比 HF PEFT
LoRALoraModuleLoRA
DoRA同上,开 variant='dora'DoRA
IA³IA3ModuleIA3
P-TuningPTuningModuleP-Tuning v1
AdapterAdapterModule—(HF PEFT 无)
Canonical LoRA同 LoRA

7.3用 NeMo LoRA(2.x 风格)

from nemo.collections.llm.peft import LoRA
from nemo.collections.llm.recipes import llama3_8b

# 1) 选 base recipe
recipe = llama3_8b.finetune_recipe(dir="./out", name="llama3-lora")

# 2) 注入 PEFT
recipe.peft = LoRA(
    target_modules=["linear_qkv", "linear_proj"],
    dim=32,
    alpha=64,
    dropout=0.05,
)

# 3) 训
import nemo_run as run
run.run(recipe, executor=run.LocalExecutor(ntasks_per_node=8))

7.4用 NeMo LoRA(1.x YAML 风格)

# 在 finetune YAML 里
model:
  peft:
    peft_scheme: lora       # lora / dora / ia3 / ptuning / adapter
    lora_tuning:
      target_modules: [linear_qkv, linear_proj]
      adapter_dim: 32
      adapter_dropout: 0.05
      alpha: 64
      a_init_method: xavier
      column_init_method: xavier
      row_init_method: zero

7.5target_modules 命名

NeMo 的层名跟 HF 不同,因为是 Megatron-Core 风格:

HF 名NeMo 名
q_proj / k_proj / v_proj(分开)linear_qkv(fused)
o_projlinear_proj
gate_proj + up_projlinear_fc1(fused)
down_projlinear_fc2

NeMo 默认把 QKV 和 gate+up 各自合并成一个 fused linear,因此 LoRA 也只对 fused 后整体加。

7.6SFT + LoRA 完整脚本

from nemo.collections.llm.recipes import llama3_8b
from nemo.collections.llm.peft import LoRA
import nemo_run as run

recipe = llama3_8b.finetune_recipe(
    dir="./sft_lora_out",
    name="llama3-sft",
    num_nodes=1, num_gpus_per_node=8,
)
recipe.peft = LoRA(
    target_modules=["linear_qkv", "linear_proj", "linear_fc1", "linear_fc2"],
    dim=16, alpha=32,
)
recipe.data = run.Config(
    "nemo.collections.llm.SquadDataModule",
    seq_length=2048, global_batch_size=128,
)
recipe.trainer.max_steps = 1000
recipe.optim.config.lr = 1e-4

run.run(recipe, executor=run.LocalExecutor(ntasks_per_node=8))

7.7训完保存 / 加载

# 保存:NeMo .nemo 格式
recipe.resume.path = "./sft_lora_out/checkpoints/last.ckpt"
# 训完目录会有 .nemo 包,包含 adapter weights

# 加载推理
from nemo.collections.llm.peft import LoraModelMixin
model = ...  # 加载 base
model.add_adapter("my_lora", "./sft_lora_out/.../*.nemo")
model.set_active_adapter("my_lora")

7.8转换为 HF 格式(推理常需)

python scripts/checkpoint_converters/convert_llama_nemo_to_hf.py \
    --input_name_or_path ./sft_lora_out/llama3-8b-lora.nemo \
    --output_path ./hf_lora_out \
    --hf_input_path meta-llama/Meta-Llama-3-8B

转完可以用 HF + PEFT 直接加载,详见 ch12。

7.9性能(NeMo LoRA vs HF PEFT)

方案速度显存备注
HF PEFT + transformersbaselinebaseline简单
NeMo LoRA + MCore1.5-2×稍省支持 TP/PP/SP
NeMo + DoRA~1.3×(vs HF)+5%变体支持

NeMo 的 LoRA 在大模型 + 多机场景明显更快(HF PEFT 在 TP > 1 时不工作)。

7.10这章你需要带走的