跳转至

megatron_peft.py — Megatron PEFT 配置

文件概述

提供 Megatron 后端的 PEFT(Parameter-Efficient Fine-Tuning)配置工具,支持 LoRA、DoRA、CanonicalLoRA、VLMLoRA 等高效微调方式。

核心函数

get_peft_cls

def get_peft_cls(model_config, bridge, provider, dtype):
    """根据配置返回对应的 PEFT 类

    支持的 PEFT 类型:
    - LoRA: 标准低秩适应
    - DoRA: 带方向的 LoRA
    - CanonicalLoRA: 规范化 LoRA
    - VLMLoRA: 视觉语言模型专用 LoRA
    """
    lora_config = model_config.get("lora", None)
    if lora_config is None or lora_config.get("rank", 0) <= 0:
        return None

    # 根据 peft_type 选择对应的类
    peft_type = lora_config.get("peft_type", "lora")
    ...

与其他模块的关系

  • 被 engine/megatron/transformer_impl.py 使用
  • 提供 Megatron 后端的 LoRA 支持

小结

为 Megatron 后端的大模型提供内存高效的微调能力,通过 PEFT 技术大幅减少可训练参数。