ShardingPlanner._merge_into() 在合并用户 plan_overrides 时,对 glob 匹配和精确匹配一视同仁地无条件覆盖 local_compute_fn 等 injection 字段。当规则用 glob 书写、而模型里多个语义不同的模块共享同一个叶名时,EP 的 compute fn 会被挂到不该挂的边界上。
ShardingPlanner._merge_into()
plan_overrides
local_compute_fn
DeepSeek-V3 正好构成这个场景:dense 层和 MoE 层都叫 .mlp。
.mlp
transformers/models/deepseek_v3/modeling_deepseek_v3.py,DeepseekV3DecoderLayer.__init__:
transformers/models/deepseek_v3/modeling_deepseek_v3.py
DeepseekV3DecoderLayer.__init__
if layer_idx >= config.first_k_dense_replace: self.mlp = DeepseekV3MoE(config) else: self.mlp = DeepseekV3MLP(config)
而 hyper_parallel/auto_models/components/distributed/ep_compute.py:25-28 的文档里,推荐的写法正是 glob:
hyper_parallel/auto_models/components/distributed/ep_compute.py:25-28
- match: "*.mlp" when: ep local_compute_fn: _target_: hyper_parallel.auto_models.components.distributed.ep_compute.qwen2moe_ep_compute_fn
两者相乘的结果:*.mlp 同时命中 dense 和 MoE。已在本地用 transformers 5.14.1 实测(hidden_size=64, num_hidden_layers=6, first_k_dense_replace=3):
*.mlp
hidden_size=64, num_hidden_layers=6, first_k_dense_replace=3
glob "*.mlp" 命中 6 个模块: layers.0.mlp DeepseekV3MLP <-- dense,无 gate/experts layers.1.mlp DeepseekV3MLP <-- dense,无 gate/experts layers.2.mlp DeepseekV3MLP <-- dense,无 gate/experts layers.3.mlp DeepseekV3MoE layers.4.mlp DeepseekV3MoE layers.5.mlp DeepseekV3MoE
DeepSeek-V3 官方 config 的默认值是 first_k_dense_replace=3、num_hidden_layers=61,即 3 个 dense + 58 个 MoE,全部叫 .mlp。
first_k_dense_replace=3
num_hidden_layers=61
origin/trainer_dev
2f14ce85
hyper_parallel/auto_models/components/distributed/sharding_planner.py:
hyper_parallel/auto_models/components/distributed/sharding_planner.py
self._merge_into(plan.modules[fqn], user_spec)
self._merge_into(plan.modules[key], user_spec)
两个分支调用完全相同、没有任何区分 glob / exact 的参数。
_merge_into
def _merge_into(derived: ModuleShardingSpec, user_spec: ModuleShardingSpec) -> None: for attr in ShardingPlanner._CONTRACT_FIELDS: ShardingPlanner._merge_contract_field(derived, user_spec, attr) for attr in ("local_compute_fn", "inner_target", "inner_wrapper", "inner_out_src", "region_dispatch", "tp_divide_attrs"): value = getattr(user_spec, attr) if value is not None: setattr(derived, attr, value) # <-- 无条件覆盖 _normalize_out_fields(derived)
值得注意的是:planner 已经知道每个边界是不是 MoE。_infer_boundary_type()(L658)能正确区分:
_infer_boundary_type()
"moe_mlp"
"mlp"
但这个结果只是个局部变量,既没有存到 spec 上,_merge_into 也从不查询它。也就是说,判断所需的信息在同一个类里已经算出来了,只是没被用上。
ep_compute.py 自己在 apply 期会做接口断言 —— L129 _require_moe_interface(module, expected_attrs, archetype_key),DeepSeek-V3 的 fn(L308 deepseekv3_ep_compute_fn)在 L341 声明 expected_attrs=["gate", "experts", "shared_experts"]。
ep_compute.py
_require_moe_interface(module, expected_attrs, archetype_key)
deepseekv3_ep_compute_fn
expected_attrs=["gate", "experts", "shared_experts"]
而 dense 的 DeepseekV3MLP 的属性是:
DeepseekV3MLP
['config', 'hidden_size', 'intermediate_size', 'gate_proj', 'up_proj', 'down_proj', 'act_fn']
一个 gate/experts/shared_experts 都没有。所以计划一旦这样产出,下游只有两种结局:要么在 apply 期断言炸掉,要么(若某条路径绕过了断言)静默走错分支。无论哪种,计划本身在产出时就已经是错的。
gate/experts/shared_experts
plan_overrides 里带 EP 语义 injection(local_compute_fn / region_dispatch)的 glob 规则,不应落到非 MoE 边界上。期望以下之一:
region_dispatch
boundary_type
_boundary_type == "moe_mlp"
only_boundary_type: moe_mlp
另外无论选哪种,ep_compute.py:25-28 的示例 YAML 都建议同步更新:目前文档推荐的 *.mlp 写法,在 dense/MoE 混排的模型上是不安全的。
ep_compute.py:25-28
我个人倾向 2(用户显式声明意图、不静默丢弃、可推广到别的 archetype),但这是维护者的设计取舍。
sharding_applier.py
getattr(spec, "local_compute_fn", None)
has_injection
ShardingPlanner.plan()
plan_overrides: - match: "*.mlp" when: ep region_dispatch: false local_compute_fn: _target_: hyper_parallel.auto_models.components.distributed.ep_compute.deepseekv3_ep_compute_fn
配 tp_size: 2, ep_size: 2 的 DeepSeek-V3。
tp_size: 2, ep_size: 2
# repro_glob_ep.py -> torchrun --nproc_per_node=2 repro_glob_ep.py import torch, torch.distributed as dist from torch.distributed.device_mesh import init_device_mesh from transformers import DeepseekV3Config from transformers.models.deepseek_v3 import modeling_deepseek_v3 as m from hyper_parallel.auto_models.components.distributed.sharding_planner import ShardingPlanner from hyper_parallel.auto_models.components.distributed.ep_compute import deepseekv3_ep_compute_fn dist.init_process_group("gloo") mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("ep",)) cfg = DeepseekV3Config( hidden_size=64, intermediate_size=128, moe_intermediate_size=32, num_hidden_layers=6, first_k_dense_replace=3, # 0..2 dense, 3..5 MoE n_routed_experts=4, n_shared_experts=1, num_experts_per_tok=2, num_attention_heads=4, num_key_value_heads=4, vocab_size=128, q_lora_rank=None, kv_lora_rank=16, qk_nope_head_dim=16, qk_rope_head_dim=8, v_head_dim=16, ) model = m.DeepseekV3Model(cfg) planner = ShardingPlanner(plan_overrides={ "*.mlp": {"when": "ep", "region_dispatch": False, "local_compute_fn": deepseekv3_ep_compute_fn}, }) plan = planner.plan(model, mesh, ep_size=2) for fqn, spec in plan.modules.items(): if getattr(spec, "local_compute_fn", None) is not None: cls = type(dict(model.named_modules())[fqn]).__name__ flag = " <-- BUG: dense" if cls == "DeepseekV3MLP" else "" print(f"{fqn:24s} {cls}{flag}")
期望输出里只有 layers.3/4/5.mlp;实际预计 layers.0/1/2.mlp(DeepseekV3MLP)也会带上 compute fn。
layers.3/4/5.mlp
layers.0/1/2.mlp
plan_overrides 的构造形式请以仓库现有用法为准(见 hyper_parallel/auto_models/_transformers/infrastructure.py 与 tests/components/distributed/test_dist_s2_apply.py)—— 上面这段是按公开签名 plan(model, mesh, *, tp_size=1, cp_size=1, ep_size=1, ...) 写的示意,若与内部约定不符请以仓库为准,结论不受影响。
hyper_parallel/auto_models/_transformers/infrastructure.py
tests/components/distributed/test_dist_s2_apply.py
plan(model, mesh, *, tp_size=1, cp_size=1, ep_size=1, ...)
Branch / commit : origin/trainer_dev @ 2f14ce85 transformers : 5.14.1(DeepSeek-V3 为原生支持,无需 trust_remote_code) torch : 2.12.0 Python : 3.11 Model : deepseek-ai/DeepSeek-V3(默认 first_k_dense_replace=3, num_hidden_layers=61) Parallel config : tp_size=2, ep_size=2, sequence_parallel=true 后端 : 与本问题无关 —— plan() 在建立 process group 之前完成
上面「命中 6 个模块」的模块清单是在 transformers 5.14.1 + torch 2.12.0 上实际跑出来的。
Thanks for contributing 🎉!
Checklist
🐛 Describe the bug
ShardingPlanner._merge_into()在合并用户plan_overrides时,对 glob 匹配和精确匹配一视同仁地无条件覆盖local_compute_fn等 injection 字段。当规则用 glob 书写、而模型里多个语义不同的模块共享同一个叶名时,EP 的 compute fn 会被挂到不该挂的边界上。DeepSeek-V3 正好构成这个场景:dense 层和 MoE 层都叫
.mlp。transformers/models/deepseek_v3/modeling_deepseek_v3.py,DeepseekV3DecoderLayer.__init__:if layer_idx >= config.first_k_dense_replace: self.mlp = DeepseekV3MoE(config) else: self.mlp = DeepseekV3MLP(config)而
hyper_parallel/auto_models/components/distributed/ep_compute.py:25-28的文档里,推荐的写法正是 glob:- match: "*.mlp" when: ep local_compute_fn: _target_: hyper_parallel.auto_models.components.distributed.ep_compute.qwen2moe_ep_compute_fn两者相乘的结果:
*.mlp同时命中 dense 和 MoE。已在本地用 transformers 5.14.1 实测(hidden_size=64, num_hidden_layers=6, first_k_dense_replace=3):DeepSeek-V3 官方 config 的默认值是
first_k_dense_replace=3、num_hidden_layers=61,即 3 个 dense + 58 个 MoE,全部叫.mlp。代码路径(基于
origin/trainer_dev@2f14ce85)hyper_parallel/auto_models/components/distributed/sharding_planner.py:self._merge_into(plan.modules[fqn], user_spec)self._merge_into(plan.modules[key], user_spec)两个分支调用完全相同、没有任何区分 glob / exact 的参数。
_merge_into:def _merge_into(derived: ModuleShardingSpec, user_spec: ModuleShardingSpec) -> None: for attr in ShardingPlanner._CONTRACT_FIELDS: ShardingPlanner._merge_contract_field(derived, user_spec, attr) for attr in ("local_compute_fn", "inner_target", "inner_wrapper", "inner_out_src", "region_dispatch", "tp_divide_attrs"): value = getattr(user_spec, attr) if value is not None: setattr(derived, attr, value) # <-- 无条件覆盖 _normalize_out_fields(derived)值得注意的是:planner 已经知道每个边界是不是 MoE。
_infer_boundary_type()(L658)能正确区分:"moe_mlp"(≈L699)"mlp"(≈L708)但这个结果只是个局部变量,既没有存到 spec 上,
_merge_into也从不查询它。也就是说,判断所需的信息在同一个类里已经算出来了,只是没被用上。为什么这是错误而不只是「配置写法问题」
ep_compute.py自己在 apply 期会做接口断言 —— L129_require_moe_interface(module, expected_attrs, archetype_key),DeepSeek-V3 的 fn(L308deepseekv3_ep_compute_fn)在 L341 声明expected_attrs=["gate", "experts", "shared_experts"]。而 dense 的
DeepseekV3MLP的属性是:一个
gate/experts/shared_experts都没有。所以计划一旦这样产出,下游只有两种结局:要么在 apply 期断言炸掉,要么(若某条路径绕过了断言)静默走错分支。无论哪种,计划本身在产出时就已经是错的。Expected behavior
plan_overrides里带 EP 语义 injection(local_compute_fn/region_dispatch)的 glob 规则,不应落到非 MoE 边界上。期望以下之一:boundary_type把关 —— glob 规则携带local_compute_fn时,只并进_boundary_type == "moe_mlp"的边界;或者only_boundary_type: moe_mlp,由用户声明而非框架猜测;或者DeepseekV3MLP上 —— 而不是产出一份要到 apply 期才炸的计划。另外无论选哪种,
ep_compute.py:25-28的示例 YAML 都建议同步更新:目前文档推荐的*.mlp写法,在 dense/MoE 混排的模型上是不安全的。我个人倾向 2(用户显式声明意图、不静默丢弃、可推广到别的 archetype),但这是维护者的设计取舍。
Additional context
ep_compute.py文档里举的例子)不是 —— 它每层都是 MoE,所以*.mlp在 Qwen2-MoE 上恰好是安全的。这可能正是该问题一直没暴露的原因。sharding_applier.py消费(该文件里getattr(spec, "local_compute_fn", None)参与has_injection判定),所以这不是某一条下游链路独有的问题,而是计划层的问题。ShardingPlanner.plan()在任何 process group 建立之前就跑完了,hccl / gloo / 单卡都一样。触发用的 YAML(最小片段)
plan_overrides: - match: "*.mlp" when: ep region_dispatch: false local_compute_fn: _target_: hyper_parallel.auto_models.components.distributed.ep_compute.deepseekv3_ep_compute_fn配
tp_size: 2, ep_size: 2的 DeepSeek-V3。复现脚本(2 卡,gloo/CPU 即可,无需 NPU)
# repro_glob_ep.py -> torchrun --nproc_per_node=2 repro_glob_ep.py import torch, torch.distributed as dist from torch.distributed.device_mesh import init_device_mesh from transformers import DeepseekV3Config from transformers.models.deepseek_v3 import modeling_deepseek_v3 as m from hyper_parallel.auto_models.components.distributed.sharding_planner import ShardingPlanner from hyper_parallel.auto_models.components.distributed.ep_compute import deepseekv3_ep_compute_fn dist.init_process_group("gloo") mesh = init_device_mesh("cpu", (2,), mesh_dim_names=("ep",)) cfg = DeepseekV3Config( hidden_size=64, intermediate_size=128, moe_intermediate_size=32, num_hidden_layers=6, first_k_dense_replace=3, # 0..2 dense, 3..5 MoE n_routed_experts=4, n_shared_experts=1, num_experts_per_tok=2, num_attention_heads=4, num_key_value_heads=4, vocab_size=128, q_lora_rank=None, kv_lora_rank=16, qk_nope_head_dim=16, qk_rope_head_dim=8, v_head_dim=16, ) model = m.DeepseekV3Model(cfg) planner = ShardingPlanner(plan_overrides={ "*.mlp": {"when": "ep", "region_dispatch": False, "local_compute_fn": deepseekv3_ep_compute_fn}, }) plan = planner.plan(model, mesh, ep_size=2) for fqn, spec in plan.modules.items(): if getattr(spec, "local_compute_fn", None) is not None: cls = type(dict(model.named_modules())[fqn]).__name__ flag = " <-- BUG: dense" if cls == "DeepseekV3MLP" else "" print(f"{fqn:24s} {cls}{flag}")期望输出里只有
layers.3/4/5.mlp;实际预计layers.0/1/2.mlp(DeepseekV3MLP)也会带上 compute fn。Environment info
上面「命中 6 个模块」的模块清单是在 transformers 5.14.1 + torch 2.12.0 上实际跑出来的。
Thanks for contributing 🎉!