[Bugfix] Fix 2 Node and Spec Decode tests (#13341)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
2025-02-16 22:20:22 +08:00
parent a0231b7c25
commit 5d2965b7d7
2 changed files with 17 additions and 9 deletions
--- a/tests/distributed/test_pipeline_parallel.py
+++ b/tests/distributed/test_pipeline_parallel.py
@@ -275,11 +275,11 @@ def _compare_tp(
    if load_format == "dummy":
        # Avoid OOM
        text_overrides = {
-            "num_layers": 1,
-            "num_hidden_layers": 1,
-            "num_experts": 2,
-            "num_experts_per_tok": 2,
-            "num_local_experts": 2,
+            "num_hidden_layers": 4,
+            "hidden_size": 512,
+            "intermediate_size": 800,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 1,
        }

        if is_multimodal: