执行摘要
修复bridge模式中因moe_token_dispatcher_type参数传播缺失导致的训练失败错误,通过添加兼容性传播提升系统稳定性,影响所有bridge模式模型。
功能与动机
此PR旨在解决Issue #1725描述的bug:当使用bridge模式时,由于模型提供者未传播moe_token_dispatcher_type参数,而slime始终设置variable_seq_lengths=True,导致Megatron-Core验证失败,finalize()抛出ValueError: Token dispatcher type: allgather does not support variable sequence length。这影响了所有bridge模式训练,需传播参数以兼容Megatron版本并纠正错误。
实现拆解
在slime/backends/megatron_utils/model_provider.py的wrapped_model_provider函数中,添加了以下代码:
if hasattr(args, "moe_token_dispatcher_type"):
provider.moe_token_dispatcher_type = args.moe_token_dispatcher_type
这确保了args中的moe_token_dispatcher_type值(在arguments.py中已修正为alltoall)传播到provider,并使用hasattr检查处理属性可能不存在的情况,以兼容较新的Megatron版本。变更仅涉及两行添加,逻辑简单直接。
评论区精华
Review评论为空,无讨论记录。
风险与影响
- 风险:变更风险低,
hasattr检查确保了向后兼容性,但需依赖args.moe_token_dispatcher_type已正确设置;无测试覆盖,但逻辑简单,回归可能性小。
- 影响:修复了所有bridge模式训练中的bug,提升系统稳定性;对用户透明,无需额外操作;对团队,减少了错误排查成本并示范了兼容性处理模式。
关联脉络
此PR与Issue #1725直接关联,修复了其中描述的bug。从Issue可追溯到PR #1715,后者添加了variable_seq_lengths传播但遗漏了moe_token_dispatcher_type,导致此问题。整体来看,这是bridge模式参数传播机制的完善,反映了在集成Megatron时对参数兼容性的持续关注。
参与讨论