执行摘要

• 一句话：统一 VeOmni MoE 参数处理为默认 fallback
• 推荐动作：建议合并。该 PR 修复了一个 MoE 参数处理的缺陷，使 VeOmni 引擎能自动处理所有模型的 gate_up_proj 拆分，消除了对显式注册的依赖。虽然改动小（仅 2 个文件 5 行增加 7 行删除），但修复逻辑清晰，风险低。

功能与动机

VeOmni 引擎在专有专家并行（EP）优化中，将所有模型的 gate_proj 和 up_proj 融合为单个 gate_up_proj 以减少通信量。但 checkpoint 加载和参数获取时需要将融合后的张量重新拆分为 gate_proj 和 up_proj。之前仅在 MOE_PARAM_HANDERS 字典中显式注册了部分模型（如 qwen3_5_moe）的处理器，其他模型使用默认的恒等映射 lambda n, t, ep_rank: iter([(n, t)])，导致融合层未被拆分，引发 State Dict Key Mismatch 错误。

实现拆解

1. 重命名函数并清空字典 (verl/workers/engine/veomni/utils.py): 将 _map_moe_params_qwen3_moe 重命名为 default_moe_param_handler，移除原有的 MOE_PARAM_HANDERS 字典（包含 qwen3_moe deepseek_v3 qwen3_5_moe 三个条目），改为空字典 {}，并添加注释说明该字典用于覆盖默认映射。
2. 修改 fallback 逻辑 (verl/workers/engine/veomni/transformer_impl.py): 在 get_per_tensor_param 方法中，将 MOE_PARAM_HANDERS.get(model_type, lambda n, t, ep_rank: iter([(n, t)])) 改为 MOE_PARAM_HANDERS.get(model_type, default_moe_param_handler)，同时导入 default_moe_param_handler。
3. 效果：现在任何未在 MOE_PARAM_HANDERS 中注册的模型都会自动使用 default_moe_param_handler，该处理器会检查参数名是否包含 gate_up_proj，若是则沿最后一维拆分为 gate_proj 和 up_proj。

关键文件：

• verl/workers/engine/veomni/utils.py（模块 VeOmni引擎；类别 source；类型 core-logic；符号 _map_moe_params_qwen3_moe, default_moe_param_handler）: 核心逻辑文件。重命名 _map_moe_params_qwen3_moe 为 default_moe_param_handler，并清空 MOE_PARAM_HANDERS 字典，实现默认 fallback。
• verl/workers/engine/veomni/transformer_impl.py（模块 VeOmni引擎；类别 source；类型 core-logic）: 调用方文件。修改 fallback 逻辑，将默认处理从恒等映射改为 default_moe_param_handler，并新增导入。

关键符号：default_moe_param_handler, get_per_tensor_param

关键源码片段

verl/workers/engine/veomni/utils.py

核心逻辑文件。重命名 _map_moe_params_qwen3_moe 为 default_moe_param_handler，并清空 MOE_PARAM_HANDERS 字典，实现默认 fallback。

# verl/workers/engine/veomni/utils.py
​
def _map_moe_params_common(name, tensor, ep_rank):
    num_experts_per_rank = tensor.size(0)
    for i in range(num_experts_per_rank):
        idx = ep_rank * num_experts_per_rank + i
        new_key = name.replace("mlp.experts.", f"mlp.experts.{idx}.") + ".weight"
        yield new_key, tensor[i].to(get_device_id(), non_blocking=True)
​
def default_moe_param_handler(name, tensor, ep_rank):
    # Handles the case where VeOmni fused gate_proj and up_proj into gate_up_proj
    if "gate_up_proj" in name:
        gate, up = tensor.chunk(2, dim=1)
        params = {
            name.replace("gate_up_proj", "gate_proj"): gate,
            name.replace("gate_up_proj", "up_proj"): up,
        }
    else:
        params = {name: tensor}
​
    for key, value in params.items():
        yield from _map_moe_params_common(key, value, ep_rank)
​
# Empty dict by default; users can override for model-specific handling
MOE_PARAM_HANDERS = {}

verl/workers/engine/veomni/transformer_impl.py

调用方文件。修改 fallback 逻辑，将默认处理从恒等映射改为 default_moe_param_handler，并新增导入。

# verl/workers/engine/veomni/transformer_impl.py
from .utils import (
    MOE_PARAM_HANDERS,
    VL_TYPE2INDEX,
    default_moe_param_handler, # newly imported
    load_veomni_model_to_gpu,
    ...
)
​
class VeOmniEngine(FSDPEngine):
    def get_per_tensor_param(self, **kwargs):
        # ...
        model_type = getattr(self.module.config, "model_type", "default")
        # Use default_moe_param_handler when model_type is not found in MOE_PARAM_HANDERS
        process_func = MOE_PARAM_HANDERS.get(model_type, default_moe_param_handler)
        # ...

评论区精华

本 PR 没有人工 review 评论，只有 Gemini Code Assist 机器人自动生成的总结性评论，未提出任何修改意见或争议。

• 暂无高价值评论线程

风险与影响

• 风险：
  • 回归风险低：改动仅涉及 MoE 参数的获取流程，且新逻辑与原 _map_moe_params_qwen3_moe 完全相同，只是移至默认路径。已有注册的三种模型（qwen3_moe, deepseek_v3, qwen3_5_moe）行为完全不变，因为它们的处理器在 MOE_PARAM_HANDERS 为空后会 fallback 到相同逻辑（注：原注册的 _map_moe_params_common 不处理 gate_up_proj 拆分，而新的 fallback 会拆分，这可能构成行为变化——但 PR 说明需要拆分，因此这是修复而非回归）。
  • 潜在兼容性风险：如果某模型原本不需要拆分（例如其 checkpoint 已经将 gate_proj 与 up_proj 分开存储），则使用 default_moe_param_handler 会产生错误拆分。但 PR 描述表明所有 VeOmni 模型均已融合，所以风险可控。
  • 无性能影响：仅在 checkpoint 加载和参数提取时调用一次，不在训练热路径上。
• 影响：
  • 开发者：不再需要为每个新模型手动添加 MoE 参数处理器，降低了集成成本。
  • 系统：修复了非注册模型 checkpoint 加载失败的问题，提升 VeOmni 引擎的兼容性。
  • 影响范围：仅限使用 VeOmni 引擎且模型使用 MoE（专家混合）结构的场景。
  • 风险标记：暂无

关联脉络

• 暂无明显关联 PR