73 lines
2.6 KiB
Diff
73 lines
2.6 KiB
Diff
diff --git a/frontier/profiling/linear_op/profiling_plan.py b/frontier/profiling/linear_op/profiling_plan.py
|
|
--- a/frontier/profiling/linear_op/profiling_plan.py
|
|
+++ b/frontier/profiling/linear_op/profiling_plan.py
|
|
@@ -112,5 +112,11 @@ def build_profiling_plan(
|
|
|
|
ffn_sharded_enabled = tp_size in ffn_tp_set
|
|
+ if is_moe and not _supports_share_expert(model_config):
|
|
+ # Routed experts are profiled by the MoE profiler. Keeping the
|
|
+ # dense FFN surrogate enabled executes an unprofiled fake weight
|
|
+ # path and can reject otherwise-valid attention TP layouts (for
|
|
+ # example Qwen3-235B TP8 with 128x128 block FP8 weights).
|
|
+ ffn_sharded_enabled = False
|
|
|
|
padded_n_embd = model_config.embedding_dim
|
|
padded_n_expanded_embd = model_config.mlp_hidden_dim
|
|
diff --git a/tests/unit/test_moe_linear_profiling_plan.py b/tests/unit/test_moe_linear_profiling_plan.py
|
|
new file mode 100644
|
|
--- /dev/null
|
|
+++ b/tests/unit/test_moe_linear_profiling_plan.py
|
|
@@ -0,0 +1,52 @@
|
|
+from types import SimpleNamespace
|
|
+
|
|
+from frontier.profiling.linear_op.profiling_plan import build_profiling_plan
|
|
+
|
|
+
|
|
+def test_routed_moe_plan_does_not_execute_dense_ffn_surrogate() -> None:
|
|
+ model_config = SimpleNamespace(
|
|
+ no_tensor_parallel=False,
|
|
+ embedding_dim=4096,
|
|
+ mlp_hidden_dim=1536,
|
|
+ num_q_heads=64,
|
|
+ num_kv_heads=4,
|
|
+ model_type="qwen3_moe",
|
|
+ post_attn_norm=True,
|
|
+ is_moe=True,
|
|
+ is_step2_mini=False,
|
|
+ )
|
|
+
|
|
+ plan = build_profiling_plan(
|
|
+ model_config=model_config,
|
|
+ tp_size=8,
|
|
+ attn_tp=[8],
|
|
+ ffn_tp=[8],
|
|
+ disable_replicated=False,
|
|
+ is_moe=True,
|
|
+ )
|
|
+
|
|
+ assert plan["attn_sharded_enabled"] is True
|
|
+ assert plan["ffn_sharded_enabled"] is False
|
|
+ assert plan["ffn_enabled"] is True
|
|
+ assert plan["padded_n_expanded_embd"] == 1536
|
|
+ assert "attn_pre_proj" in plan["enabled_ops"]
|
|
+ assert "post_attention_layernorm" in plan["enabled_ops"]
|
|
+ assert "mlp_up_proj" not in plan["enabled_ops"]
|
|
+ assert "mlp_down_proj" not in plan["enabled_ops"]
|
|
+
|
|
+
|
|
+def test_dense_plan_still_executes_ffn_surrogate() -> None:
|
|
+ model_config = SimpleNamespace(
|
|
+ no_tensor_parallel=False,
|
|
+ embedding_dim=4096,
|
|
+ mlp_hidden_dim=1536,
|
|
+ num_q_heads=64,
|
|
+ num_kv_heads=4,
|
|
+ model_type="qwen2",
|
|
+ post_attn_norm=True,
|
|
+ is_moe=False,
|
|
+ is_step2_mini=False,
|
|
+ )
|
|
+ plan = build_profiling_plan(model_config, 8, [8], [8], is_moe=False)
|
|
+ assert plan["ffn_sharded_enabled"] is True
|
|
+ assert "mlp_up_proj" in plan["enabled_ops"]
|