utils.py — VeOmni 工具函数

文件概述

提供 VeOmni 引擎的模型/优化器 CPU 卸载工具和 MoE 参数处理函数。

核心内容

视觉语言模型 token 索引

VL_TYPE2INDEX = {
    "qwen2_5_vl": {"IMAGE_INPUT_INDEX": 151655, "VIDEO_INPUT_INDEX": 151656},
    "qwen3_vl": {"IMAGE_INPUT_INDEX": 151655, "VIDEO_INPUT_INDEX": 151656},
}

CPU 卸载函数

def offload_veomni_model_to_cpu(model):
    """将 VeOmni 模型卸载到 CPU（FSDP2 兼容）"""
    model.reshard()
    model.cpu()

def load_veomni_model_to_gpu(model):
    """将模型加载到 GPU"""
    model.to(get_device_id())

def offload_veomni_optimizer(optimizer):
    """卸载优化器状态到 CPU，支持 MultiOptimizer"""

def load_veomni_optimizer(optimizer, device_id):
    """加载优化器状态到 GPU"""

MoE 参数处理

MOE_PARAM_HANDERS = {
    "qwen3_moe": _map_moe_params_qwen3_moe,
}

def _map_moe_params_qwen3_moe(name, tensor):
    """将堆叠的专家权重拆分为独立的专家参数"""
    for i in range(tensor.size(0)):
        new_key = name.replace("mlp.experts.", f"mlp.experts.{i}.") + ".weight"
        yield new_key, tensor[i]

与其他模块的关系

• 被 engine/veomni/transformer_impl.py 使用
• 提供 VeOmni 特有的 FSDP2 兼容 offload 逻辑

小结

这些工具函数处理了 VeOmni 框架与 verl 之间的差异，特别是 FSDP2 的 CPU 卸载和 MoE 权重格式转换。