执行摘要
- 一句话:新增Qwen3.5模型的Megatron训练示例和代码支持。
- 推荐动作:此PR值得工程师精读,以了解如何适配新模型架构到现有训练框架。关注设计决策:chat_template的异常处理机制、tool_parser的扩展性、以及Megatron配置中use_remove_padding参数的设计权衡。对于维护者,建议审查Exception捕获逻辑,并跟踪Megatron-LM对Qwen3.5 THD格式的支持进展。
功能与动机
PR body中说明,感谢@LiuXTao在mbridge上的工作(https://github.com/ISEEKYAN/mbridge/pull/83)支持了Qwen3.5,此PR基于此在verl上成功运行Qwen3.5 SFT。Issue评论中用户询问是否支持RL,作者更新了RL脚本,但提到vllm依赖问题,体现了添加Qwen3.5训练支持的需求。
实现拆解
实现方案分为多个模块:
- 示例脚本:新增examples/grpo_trainer/run_qwen3_5-35b-megatron.sh和examples/sft/gsm8k/run_qwen3_5_megatron.sh,配置Megatron并行参数和Qwen3.5特定设置(如use_remove_padding=False)。
- chat模板:修改verl/utils/chat_template.py中的apply_chat_template函数,处理Qwen3.5至少需要一个用户消息的模板需求。
- 工具解析器:在verl/experimental/agent_loop/tool_parser.py中新增Qwen3XMLToolParser类,支持Qwen3.5模型的XML格式工具调用解析。
- agent循环:调整verl/experimental/agent_loop/agent_loop.py中的_compute_position_ids函数,以兼容transformers>=5.3.0引入的mm_token_type_ids参数。
- 数据处理器:修改verl/utils/dataset/multiturn_sft_dataset.py,增强对视觉内容的处理和错误处理。
- Megatron工具:更新verl/utils/megatron_utils.py中的load_megatron_model_to_gpu函数,修复Gated DeltaNet等非标准层的grad buffer处理。
- 其他修改:包括更新CI工作流以支持transformers>=5.0.0,修复NestedTensor处理(verl/utils/tensordict_utils.py)等。
关键文件:
examples/grpo_trainer/run_qwen3_5-35b-megatron.sh(模块 examples): 新增GRPO训练脚本,演示Qwen3.5-35B-A3B MoE模型的Megatron配置、并行参数设置和Qwen3.5特定注意事项(如use_remove_padding=False)。
examples/sft/gsm8k/run_qwen3_5_megatron.sh(模块 examples): 新增SFT训练脚本,用于Qwen3.5-397B-A17B模型的Megatron后端训练,包含分布式设置和模型架构注释。
verl/utils/chat_template.py(模块 utils): 修改apply_chat_template函数以处理Qwen3.5的chat模板需求,特别是至少需要一个用户消息的特殊情况,增强兼容性。
verl/experimental/agent_loop/tool_parser.py(模块 agent_loop): 新增Qwen3XMLToolParser类以支持Qwen3.5模型的工具调用解析,基于外部实现适配,扩展工具解析能力。
verl/utils/tensordict_utils.py(模块 utils): 修复chunk_tensordict函数以正确处理NestedTensor,特别是3D jagged tensors(如MRoPE position_ids),解决PyTorch bug导致的unbind失败。
关键符号:apply_chat_template, extract_tool_calls, _compute_position_ids, chunk_tensordict, load_megatron_model_to_gpu
评论区精华
Review评论中的核心讨论包括:
风险与影响
- 风险:技术风险具体包括:
- 兼容性风险:依赖特定版本的vllm、transformers和mbridge,Issue评论中显示版本冲突(如transformers==5.2.0与vllm nightly的兼容性问题),可能导致用户环境配置复杂。
- 性能风险:Qwen3.5使用Gated Delta Net线性注意力,不支持packed sequences(THD格式),必须设置use_remove_padding=False,可能影响训练性能,直到Megatron-LM添加支持(PR讨论中提及)。
- 代码健壮性风险:verl/utils/dataset/multiturn_sft_dataset.py中捕获宽泛Exception可能掩盖非预期错误,增加调试难度。
- 回归风险:修改核心函数如verl/utils/tensordict_utils.py中的chunk_tensordict可能影响其他使用NestedTensor的模块,需确保向后兼容。
- 依赖外部项目风险:mbridge支持是关键依赖,但需从源码安装,增加部署复杂度。
- 影响:影响范围:
- 对用户:现在可以在verl中运行Qwen3.5的SFT和GRPO训练,通过提供的示例脚本快速上手,但需注意硬件要求(如128+ GPUs for 397B模型)和依赖版本管理。
- 对系统:引入新模型架构支持,需升级transformers到>=5.0.0(CI工作流已更新),可能影响现有模型训练流程的兼容性。
- 对团队:需维护Qwen3.5相关代码,处理后续bug和性能优化,Issue评论中显示用户反馈了多个运行问题(如vllm初始化、transformers版本冲突),增加了支持负担。
- 风险标记:版本兼容性问题, 宽泛Exception捕获, NestedTensor处理复杂性
关联脉络
- PR #5695 [fully_async] fix: Patch vllm013 weight loader for qwen3-moe series: 类似地修复了vllm中的模型权重加载问题,与本PR的模型兼容性工作相关,都涉及新模型架构的支持。
- PR #5575 [megatron] feat: checkpoint save as HF PEFT format: 涉及Megatron模块的改进,与本PR的Megatron训练示例相关,展示了Megatron功能的持续扩展。
参与讨论