model_forward.py — mcore 模型前向函数

文件路径

verl/models/mcore/model_forward.py

文件概述

提供 Megatron-Core 模型的前向函数生成器。核心函数 model_forward_gen() 返回一个闭包，处理输入预处理（THD/BSHD 打包）、模型前向、输出后处理（解包）的完整流程。

关键代码讲解

model_forward_gen() -- 前向函数工厂

def model_forward_gen(vision_model: bool = False):
    def model_forward(model, input_ids, attention_mask, position_ids,
                      multi_modal_inputs, logits_processor=None,
                      data_format="thd", mtp_config=None):
        pre_process = unwrap_model(model).pre_process if not vision_model else False
        post_process = unwrap_model(model).post_process

        # 1. 预处理：打包序列
        if data_format == "thd":
            input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(
                input_ids, attention_mask, pre_process=pre_process)
        else:  # bshd
            input_ids, attention_mask, position_ids = preprocess_bshd(
                input_ids, attention_mask, position_ids)

        # 2. 模型前向
        output = model(
            input_ids=input_ids_rmpad,
            position_ids=position_ids,
            attention_mask=None,  # THD 使用 packed_seq_params
            packed_seq_params=packed_seq_params,
            **model_kwargs,
        )

        # 3. 后处理：解包 + logits_processor
        if post_process:
            output = postprocess_packed_seqs(output, packed_seq_params, attention_mask, ...)
            if logits_processor:
                output = logits_processor(output, **logits_processor_args)
        return output
    return model_forward

VLM 处理

当 vision_model=True 时： - pre_process 设为 False（VLM 内部处理 embedding） - 多模态输入通过 model_kwargs 传递

if "pixel_values" in multi_modal_inputs:
    model_kwargs["pixel_values"] = multi_modal_inputs["pixel_values"]
if "image_grid_thw" in multi_modal_inputs:
    model_kwargs["image_grid_thw"] = multi_modal_inputs["image_grid_thw"]

无 padding 版本 -- gptmodel_forward_no_padding()

def gptmodel_forward_no_padding(model, input_ids, ...):
    """输入为 nested tensor 的前向函数"""
    input_ids_rmpad, packed_seq_params = preprocess_thd_no_padding(input_ids)
    output = model(input_ids=input_ids_rmpad, packed_seq_params=packed_seq_params, ...)
    output = postprocess_thd_no_padding(output, packed_seq_params, input_ids, ...)

核心函数列表

函数名	作用
model_forward_gen()	前向函数工厂（THD/BSHD 格式）
gptmodel_forward_no_padding()	无 padding 前向（nested tensor 输入）

与其他模块的关系

• 被 registry.py 注册为前向函数
• 使用 util.py 的打包/解包函数
• 被训练流程中的 actor/critic 模型调用

小结

这个文件封装了 mcore 模型前向的完整流程：打包输入 -> 模型计算 -> 解包输出。通过工厂模式生成前向函数，支持 THD/BSHD 两种格式和 VLM 多模态输入。