已开启
[Bug]: plan_overrides 的 glob "*.mlp" 会把 EP 的 local_compute_fn 注入到 dense MLP 边界上 #350
shawnylee233创建于  25 天前
shawnylee233
25 天前 创建

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.pyDeepseekV3DecoderLayer.__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):

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=3num_hidden_layers=61,即 3 个 dense + 58 个 MoE,全部叫 .mlp

代码路径(基于 origin/trainer_dev @ 2f14ce85

hyper_parallel/auto_models/components/distributed/sharding_planner.py

  • L1150(glob 分支):self._merge_into(plan.modules[fqn], user_spec)
  • L1156(精确键分支):self._merge_into(plan.modules[key], user_spec)

两个分支调用完全相同、没有任何区分 glob / exact 的参数

  • L1280-1295 _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_EXPERT + 容器叶名 → "moe_mlp"(≈L699)
  • SHARED_EXPERT → "mlp"(≈L708)

但这个结果只是个局部变量,既没有存到 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"]

而 dense 的 DeepseekV3MLP 的属性是:

['config', 'hidden_size', 'intermediate_size', 'gate_proj', 'up_proj', 'down_proj', 'act_fn']

一个 gate/experts/shared_experts 都没有。所以计划一旦这样产出,下游只有两种结局:要么在 apply 期断言炸掉,要么(若某条路径绕过了断言)静默走错分支。无论哪种,计划本身在产出时就已经是错的。


Expected behavior

plan_overrides 里带 EP 语义 injection(local_compute_fn / region_dispatch)的 glob 规则,不应落到非 MoE 边界上。期望以下之一:

  1. planner 用它已有的 boundary_type 把关 —— glob 规则携带 local_compute_fn 时,只并进 _boundary_type == "moe_mlp" 的边界;或者
  2. 提供显式的意图声明,让 YAML 能把范围写出来,例如 only_boundary_type: moe_mlp,由用户声明而非框架猜测;或者
  3. 如果认为「glob 命中 dense」是使用者的责任,那么至少 在 planner 层给出明确报错,指出某条 glob 规则把 EP compute fn 挂到了 DeepseekV3MLP 上 —— 而不是产出一份要到 apply 期才炸的计划。

另外无论选哪种,ep_compute.py:25-28 的示例 YAML 都建议同步更新:目前文档推荐的 *.mlp 写法,在 dense/MoE 混排的模型上是不安全的。

我个人倾向 2(用户显式声明意图、不静默丢弃、可推广到别的 archetype),但这是维护者的设计取舍。


Additional context

  • 触发条件是 模型里 dense 与 MoE 共享叶名。DeepSeek-V3 是这样;Qwen2-MoE(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.mlpDeepseekV3MLP)也会带上 compute fn。

plan_overrides 的构造形式请以仓库现有用法为准(见 hyper_parallel/auto_models/_transformers/infrastructure.pytests/components/distributed/test_dist_s2_apply.py)—— 上面这段是按公开签名 plan(model, mesh, *, tp_size=1, cp_size=1, ep_size=1, ...) 写的示意,若与内部约定不符请以仓库为准,结论不受影响。


Environment info

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 🎉!

likedislike