bridge.py — Megatron-Bridge 集成

文件路径

verl/models/mcore/bridge.py

文件概述

集成 Megatron-Bridge 库，提供 HF 到 mcore 的自动权重转换、LoRA 适配等功能。Megatron-Bridge 是一个更高层的工具库，简化了 mcore 模型的使用。

关键代码讲解

导入和依赖

from megatron.bridge import AutoBridge
from megatron.bridge.models.conversion.param_mapping import AutoMapping
from megatron.bridge.peft.canonical_lora import CanonicalLoRA
from megatron.bridge.peft.dora import DoRA
from megatron.bridge.peft.lora import LoRA, VLMLoRA

LinearForLastLayer -- Value 模型输出层

class LinearForLastLayer(torch.nn.Linear):
    """支持序列并行的最后一层线性层（用于 critic 模型）"""
    def __init__(self, input_size, output_size, *, sequence_parallel):
        super().__init__(in_features=input_size, out_features=output_size, bias=False)
        self.sequence_parallel = sequence_parallel

    def forward(self, input_, weight=None, runtime_gather_output=None):
        if self.sequence_parallel:
            # 先 gather 序列并行的分片
            input_ = tensor_parallel.gather_from_sequence_parallel_region(input_)
        output = torch.nn.functional.linear(input_, self.weight)
        return output, None

导出的功能

名称	作用
AutoBridge	自动的 HF <-> mcore 权重转换
LoRA	标准 LoRA 适配
VLMLoRA	VLM 专用 LoRA
DoRA	权重分解 LoRA
CanonicalLoRA	标准化 LoRA
LinearForLastLayer	Value 模型输出层

与其他模块的关系

• 被 mcore 训练流程中的 LoRA 和权重转换代码调用
• LinearForLastLayer 被 model_initializer.py 使用

小结

Bridge 模块封装了 Megatron-Bridge 库的功能，提供了更便捷的权重转换和 LoRA 适配接口。LinearForLastLayer 解决了 critic 模型在序列并行下的输出层问题。