十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

verl 使用 Megatron-LM 后端添加新模型:基于 mcore GPTModel 的四步扩展指南

verl 使用 Megatron-LM 后端添加新模型:基于 mcore GPTModel 的四步扩展指南 verl 使用 Megatron-LM 后端添加新模型基于 mcore GPTModel 的四步扩展指南【免费下载链接】verlverl/HybridFlow: A Flexible and Efficient RL Post-Training Framework项目地址: https://gitcode.com/GitHub_Trending/ve/verl导读verlHybridFlow在 Megatron 训练后端上通过 Megatron-Core 的统一GPTModel抽象支持各类语言模型无需为每个模型重写整套训练逻辑。本文基于 docs/advance/megatron_extension.rst 展开结合 verl/models/mcore 目录下的真实实现系统讲解在 verl 中为 Megatron 后端添加新模型的完整四步流程、LayerSpec/TransformerConfig/ HF Config 三者之间的关系以及 config 转换、模型初始化、注册表与权重加载等配套机制。读完本文你将具备为 verl 的 Megatron 后端接入新架构模型含 MoE、MLA 类模型的实战能力。一、背景为什么 verl 的 Megatron 后端统一使用 GPTModelverl 早期版本基于 Megatron-LM 0.4 并对 HuggingFace 模型类做 workaround 适配维护成本高、无法及时享受新特性。随后 verl 迁移到 Megatron-Coremcore并采用官方推荐的GPTModel类承载所有语言模型。这样做的直接收益是免费获得 mcore 持续演进的底层能力如context parallel上下文并行、expert parallel专家并行、dist_checkpointing分布式检查点等训练引擎与 rollout 引擎vLLM / SGLang / TensorRT-LLM 等使用 HuggingFace 格式之间的权重同步只需处理 HF 格式与 mcore 格式的映射新增模型时绝大多数情况下不必触碰并行策略、调度等上层逻辑只需要回答如何把这个模型建模成 mcore GPTModel这一个问题。从 verl/models/mcore/readme.md 可以看到verl 集成 megatron-core 先后有三种方式目录内逐模型转换代码现已弃用、mbridge弃用以及 Megatron-Bridge当前默认的官方演进方向。配置项megatron.use_mbridge与megatron.vanilla_mbridge用于切换目前use_mbridge被断言为 truevanilla_mbridgefalse时走 Megatron-Bridgevanilla_mbridgetrue时兼容旧的 mbridge。无论走哪条 bridge 路线模型层统一建模为 mcoreGPTModel的原则不变这也是本文后续所有内容的基础。二、添加新模型的核心流程四步法原文档 docs/advance/megatron_extension.rst 给出的操作步骤非常精炼本质是复用 mcore 标准建模能力的四步定位模型初始化入口找到 verl/models/mcore/model_initializer.py。这是 verl 中所有 mcore 模型初始化的核心文件其内部定义了统一的BaseModelInitializer抽象基类及一系列具体模型的初始化器。判断模型能否由TransformerLayerSpec直接配置mcore 的GPTModel把每层长什么样抽象为TransformerLayerSpec在 Megatron-LM 的megatron/core/models/gpt/gpt_layer_specs.py中定义。如果新模型的标准层attention MLP 归一化恰好能被现有 spec 描述就直接复用GPTModel无需任何自定义层代码否则需要在这里实现新的ModelLayerSpec和对应的ModelLayer把新架构的算子组织进 mcore 的层框架中。用正确的三元组初始化 GPTModel将LayerSpec、TransformerConfig和HuggingfaceConfig作为参数传入GPTModel构造函数。其中TransformerConfig描述并行规模、数据类型、层数与各类开关HuggingfaceConfig提供词表大小、最大序列长度、RoPE 参数等模型本身的结构信息。返回模型初始化完成后把GPTModel实例对 MoE / 价值模型等场景可能还需要额外处理见下文返回给上层训练流程。这四步的代码落点就是BaseModelInitializer的initialize方法下面深入剖析其源码实现。三、源码级剖析BaseModelInitializer 与 initialize 的关键参数BaseModelInitializer定义在 verl/models/mcore/model_initializer.py它是一个抽象基类构造时接收 mcore 的TransformerConfig与 HF 的PretrainedConfigclass BaseModelInitializer(ABC): def __init__(self, tfconfig: TransformerConfig, hf_config: PretrainedConfig): self.tfconfig tfconfig self.hf_config hf_config self.has_vp_stage inspect.signature( get_gpt_decoder_block_spec).parameters.get(vp_stage, None) is not None abstractmethod def get_transformer_layer_spec(self, vp_stageNone): Get the transformer layer specification. pass子类唯一必须实现的方法是get_transformer_layer_spec——这正是原文档步骤 2 的代码形态通过覆写该方法决定层的规格。has_vp_stage用inspect.signature探测当前 Megatron 版本的get_gpt_decoder_block_spec是否支持vp_stage参数从而兼容不同 mcore 版本。核心的initialize方法model_initializer.py签名与参数语义如下参数默认值含义pre_processTrue是否包含 embedding 层流水线并行首段为 Truepost_processTrue是否包含输出层流水线并行末段为 Trueshare_embeddings_and_output_weightsFalse输入 embedding 与输出 logit 权重是否共享valueFalse是否为 Critic/价值模型额外挂一个输出维度为 1 的线性层**extra_kwargs—透传vp_stage、mtp_block_spec等方法内部先取得transformer_layer_spec再通过get_rope_scaling_args()从 HF config 中提取rope_scaling的factor映射为 mcore 的seq_len_interpolation_factor当前仅支持 linear 类缩放随后构造GPTModelmodel GPTModel( configself.tfconfig, transformer_layer_spectransformer_layer_spec, vocab_sizeself.hf_config.vocab_size, max_sequence_lengthself.hf_config.max_position_embeddings, pre_processpre_process, post_processpost_process, share_embeddings_and_output_weightsshare_embeddings_and_output_weights, position_embedding_typerope, rotary_baseget_hf_rope_theta(self.hf_config), **rope_scaling_args, mtp_block_specmtp_block_spec, )注意rotary_base取自get_hf_rope_theta(self.hf_config)RoPE 基数直接来自 HF config保证与原始权重语义一致。若post_process and value为真还会用 verl/models/mcore/bridge.py 中的LinearForLastLayer把model.output_layer替换为hidden_size - 1的线性层RL 训练中 Critic 输出标量价值并遵循sequence_parallel配置。四、内置模型初始化器从 Dense 到 MoE 与 MLAverl 的 model_initializer.py 内置了覆盖 Dense、MoE、MLA 三类架构的初始化器可直接作为新模型的参照模板DenseModel面向 Llama、Qwen2、Qwen3 等稠密模型。get_transformer_layer_spec断言normalization RMSNorm直接调用get_gpt_decoder_block_spec(tfconfig, use_transformer_engineTrue)返回标准解码层 spec。Qwen2MoEModel / Qwen3MoEModel面向 Qwen 系列 MoE。在标准 spec 基础上把shared_experts的gate参数置为TrueQwen MoE 有共享专家initialize中默认freeze_moe_routerTrue将每一层layer.mlp.router.weight.requires_grad置为False这与 Qwen 官方默认冻结 router 的策略一致。MixtralModelMixtral 的 MoE 无共享专家spec 无需 patch仅按需冻结 router默认freeze_moe_routerFalse。DeepseekV3ModelMLAMulti-head Latent Attention MoE 的典型代表。get_rope_scaling_args返回空MLA 的 RoPE 参数由MLATransformerConfig管理initialize中默认freeze_moe_routerTrue并把moe_router_load_balancing_type设为none当mtp_num_layers 0时通过get_gpt_mtp_block_spec构造 MTPMulti-Token Prediction块并注入mtp_block_spec。从这些内置实现可以看出 verl 的一个通用模式初始化器不只是建模型还负责把模型家族在 RL 训练中的既有约定如冻结 router、关闭 aux loss、挂价值头落到模型结构上。自定义新模型时覆写get_transformer_layer_spec与initialize即可复用这套约定。五、配套组件 1HF Config 到 mcore TransformerConfig 的转换GPTModel需要 mcore 风格的TransformerConfig而用户手里通常是 HuggingFace 的PretrainedConfig。这一转换集中在 verl/models/mcore/config_converter.py其函数命名与模型架构一一对应hf_to_mcore_config_denseLlama / Qwen2 / Qwen3 等稠密模型。自动推导qkv_biasQwen2 为 True其余取attention_bias与qk_layernormQwen3 为 True。hf_to_mcore_config_qwen2moe/hf_to_mcore_config_qwen3moeQwen MoE 系列写入moe_ffn_hidden_size、num_moe_experts、moe_router_topk、moe_shared_expert_intermediate_size、moe_aux_loss_coeff等RL 场景下默认moe_router_load_balancing_typenone注释明确说明turn off aux_loss as it hurts perf in RL。qwen3moe版本额外支持 MTP从hf_config或hf_config.text_config读取mtp_num_hidden_layers并设置mtp_num_layers与mtp_loss_scaling_factor默认 0.1。hf_to_mcore_config_mixtralMixtral 无共享专家moe_shared_expert_intermediate_sizeNone、moe_shared_expert_overlapFalse。hf_to_mcore_config_dpskv3DeepSeek-V3 的 MLA 版本返回MLATransformerConfig。内部设置beta_fast32, beta_slow1等 YaRN 相关参数attention_backendAttnBackend.fused、qk_layernormTrue、moe_router_score_functionsigmoid、moe_router_dtypefp64、disable_bf16_reduced_precision_matmulTrue等 DeepSeek 特有配置并通过first_k_dense_replace构造moe_layer_freq前 K 层为稠密层。同时对当前不支持的特性显式断言num_nextn_predict_layers必须为 0MTP 暂不支持、不允许存在quantization_config。这些转换函数的公共骨架是_get_base_transformer_config它从 HF config 抽取num_layers、hidden_size、num_attention_heads、num_query_groups取num_key_value_heads、ffn_hidden_size取intermediate_size、kv_channels取head_dim等并把当前进程组的并行规模TP / PP / EP / VP / CP通过 mpu 读入设置sequence_parallel tp 1、variable_seq_lengthsTrue、masked_softmax_fusionTrue等训练友好默认值override_transformer_config_kwargs允许上层如 YAML 中actor_rollout_ref.actor.megatron.override_transformer_config.*按需覆盖任意字段。最后check_and_construct_configs会剔除当前 Megatron 版本不认识的 key 并给出告警避免版本不匹配导致崩溃。六、配套组件 2注册表机制与模型 forward 函数初始化器与配置转换器并不会被上层直接调用而是通过 verl/models/mcore/registry.py 中的注册表按 HFarchitectures[0]字符串如LlamaForCausalLM、Qwen2MoeForCausalLM、DeepseekV3ForCausalLM自动分发。registry 维护了四张表MODEL_CONFIG_CONVERTER_REGISTRY架构名 → 配置转换函数MODEL_INITIALIZER_REGISTRY架构名 → 初始化器类MODEL_FORWARD_REGISTRY/MODEL_FORWARD_FUSED_REGISTRY架构名 → forward 函数含 VLM 版本MODEL_WEIGHT_CONVERTER_REGISTRY架构名 → mcore→HF 权重转换器。对外入口为hf_to_mcore_config()、init_mcore_model()与get_mcore_weight_converter()。添加新模型时除了实现初始化器与配置转换器还必须把新架构登记进这些注册表否则会抛出Model Type: xxx not supported的NotImplementedError。模型的前向调用由 verl/models/mcore/model_forward.py 提供model_forward_gen(vision_modelFalse)生成的 forward 支持thdpacked / 去 paddingRL 长序列默认与bshdGPT-OSS 等暂不支持 thd 的模型使用需 padding、建议 batch size1两种数据格式并处理value_model输出压缩、MTP 的 label/loss_mask 对齐_build_mtp_loss_mask_nested与 VLM 的pixel_values/image_grid_thw等多模态输入。对于 mcore 0.12 的 MLA 等已知缺陷verl/models/mcore/patch.py 提供apply_patch()系列运行时补丁例如修正get_query_key_value_tensors在packed_seq_params下的行为、为 ROCm 环境注入纯 PyTorch 的fast_hadamard_transform兜底DeepSeek 稀疏注意力 DSA 需要等。七、配套组件 3权重加载与转换初始化出GPTModel后还要把 HF 权重灌入并行切分的 mcore 模型。verl/models/mcore/loader.py 的load_state_dict_to_megatron_gptmodel是早期逐模型转换方式的核心实现它按_megatron_calc_layer_map计算全局层号到(pp_rank, vp_rank, local_layer_idx)的映射通过torch.distributed.broadcast在 TP 组内分发 embedding、QKV把 HF 的q_proj/k_proj/v_proj合并为 mcore 的 fused QKV见_broadcast_tp_shard_tensor_qkv、gate/up合并为 SwiGLU 的 fused 形式见_broadcast_tp_shard_tensor_gate_up等权重最后在 DP 组内broadcast_params完成复制。不过按 readme.md 的说明这种方式已被弃用当前主流路径是走Megatron-Bridge或兼容的 mbridge以 HuggingFace 格式直接在线加载/保存 mcore 权重并支持megatron 0.13的新优化特性——这正是 readme 中支持新模型三步的核心推理引擎支持该模型 → bridge 支持该模型新架构向 megatron-bridge 提 issue 或贡献实现→ 修改脚本中的模型路径直接运行。离线转换脚本可参考 scripts/converter_hf_to_mcore.py在线 mcore→HF 转换逻辑在 verl/utils/megatron_utils.py 的convert_megatron_model_to_transformers_model。八、实战清单把新模型接入 verl Megatron 后端的检查项综合原文档四步与仓库源码接入一个全新架构模型如某私有 MoE的完整检查清单如下确认推理引擎vLLM / SGLang / TensorRT-LLM与 bridge首选 Megatron-Bridge已支持该架构注意 Megatron 与 TransformerEngine 版本需匹配若该模型架构不在 registry.py 的SupportedModel枚举中新增枚举值并在四张注册表中登记在 config_converter.py 新增hf_to_mcore_config_arch复用_get_base_transformer_configDense或_get_mla_transformer_configMLA按模型特性补充 MoE / MTP / qk_layernorm 等字段在 model_initializer.py 新增继承BaseModelInitializer的初始化器能被现有 spec 描述则仅覆写get_transformer_layer_spec否则按 mcore 规范实现新的ModelLayerSpec/ModelLayer按需在initialize中处理冻结 router、关闭 aux loss、挂价值头、拼接 MTP 块等 RL 训练约定以init_mcore_model与hf_to_mcore_config为入口自测模型可初始化、可前向并与 vLLM 用相同权重对比生成结果一致性在训练脚本中通过actor_rollout_ref.actor.megatron.override_transformer_config.*覆盖默认TransformerConfig并切换数据格式默认thd模型不支持 packed 时改用bshd。现有示例可参考 examples/grpo_trainer/run_qwen3_8b_megatron.sh、examples/grpo_trainer/run_deepseek_v3_671b_megatron.sh 等 Megatron 后端脚本以及 tests/utils/megatron 与 tests/special_distributed 下的相关测试用例。结语verl 把为 Megatron 后端添加新模型收敛为一次围绕 mcoreGPTModel的建模工作层规格LayerSpec决定模型结构TransformerConfig决定并行与数值行为HF Config 决定参数与词表语义三者通过BaseModelInitializer汇聚再经注册表接入训练管线。对于绝大多数主流架构你只需参照内置的 Dense / MoE / MLA 三类初始化器完成配置转换与 spec 选择即可只有遇到全新算子结构时才需要深入实现自定义ModelLayer。这套设计让 verl 能在保持 mcore 高性能训练能力的同时以极小的成本跟上社区模型迭代的速度。【免费下载链接】verlverl/HybridFlow: A Flexible and Efficient RL Post-Training Framework项目地址: https://gitcode.com/GitHub_Trending/ve/verl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表