Allow small num_heads

bench/exo_bench.py

@@ -35,14 +35,17 @@ from harness import (
     ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
+    get_all_instance_ids,
     instance_id_from_instance,
     node_ids_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
     settle_and_fetch_placements,
+    shard_split_summary,
     wait_for_instance_gone,
     wait_for_instance_ready,
+    wait_for_new_instance,
 )
 from loguru import logger
 from transformers import AutoTokenizer
@@ -506,6 +509,7 @@ def main() -> int:
         )
 
         if sharding == "Tensor" and tensor_strategy != "Naive":
+            before_ids = get_all_instance_ids(client)
             client.request_json(
                 "POST",
                 "/place_instance",
@@ -517,8 +521,13 @@ def main() -> int:
                     "tensor_strategy": tensor_strategy,
                 },
             )
+            instance_id = wait_for_new_instance(client, before_ids)
+            instance = client.get_instance(instance_id)
+            logger.info(f"place_instance created instance_id={instance_id}")
+            logger.info(f"Shard split: {shard_split_summary(instance)}")
         else:
             client.request_json("POST", "/instance", body={"instance": instance})
+            logger.info(f"Shard split: {shard_split_summary(instance)}")
         try:
             wait_for_instance_ready(client, instance_id)
         except (RuntimeError, TimeoutError) as e:

bench/harness.py

@@ -131,6 +131,24 @@ def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
     return list(inner["shardAssignments"]["nodeToRunner"].keys())
 
 
+def shard_split_summary(instance: dict[str, Any]) -> str:
+    inner = unwrap_instance(instance)
+    runner_to_shard = inner["shardAssignments"]["runnerToShard"]
+    parts: list[str] = []
+    for rid, shard_wrapper in runner_to_shard.items():
+        shard = next(iter(shard_wrapper.values()))
+        start = shard.get("startLayer", "?")
+        end = shard.get("endLayer", "?")
+        n = shard.get("nLayers", "?")
+        weights = shard.get("shardWeights")
+        rank = shard.get("deviceRank", "?")
+        part = f"rank {rank}: layers [{start},{end})/{n}"
+        if weights:
+            part += f" weights={[round(w, 3) for w in weights]}"
+        parts.append(part)
+    return " | ".join(parts)
+
+
 def runner_ready(runner: dict[str, Any]) -> bool:
     return "RunnerReady" in runner
 
@@ -151,6 +169,7 @@ def wait_for_instance_ready(
     start_time = time.time()
     instance_existed = False
     last_loaded: dict[str, int] = {}
+    last_states: dict[str, str] = {}
     while time.time() - start_time < timeout:
         instance = client.get_instance(instance_id)
 
@@ -166,21 +185,29 @@ def wait_for_instance_ready(
         rids = runner_ids_from_instance(instance)
 
         all_ready = True
+        runner_states: dict[str, str] = {}
         for rid in rids:
             runner = client.get_runner(rid) or {}
             if runner_failed(runner):
                 error_msg = get_runner_failed_message(runner) or "Unknown error"
                 raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+            state_tag = next(iter(runner), "Unknown")
             if "RunnerLoading" in runner:
                 loading = runner["RunnerLoading"]
                 loaded = loading.get("layersLoaded", 0)
                 total = loading.get("totalLayers", 0)
-                if total > 0 and last_loaded.get(rid) != loaded:
+                if total > 0:
                     last_loaded[rid] = loaded
-                    logger.debug(f"Runner {rid}: loading layers {loaded}/{total}")
+                    state_tag = f"Loading {loaded}/{total}"
+            runner_states[rid] = state_tag
             if not runner_ready(runner):
                 all_ready = False
 
+        if runner_states != last_states:
+            last_states = runner_states
+            parts = [f"{rid[:8]}: {state}" for rid, state in runner_states.items()]
+            logger.debug(f"Runners: {', '.join(parts)}")
+
         if all_ready:
             return
 
@@ -189,6 +216,24 @@ def wait_for_instance_ready(
     raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")
 
 
+def get_all_instance_ids(client: ExoClient) -> set[str]:
+    instances = client.get_state_path("instances") or {}
+    return set(instances.keys())
+
+
+def wait_for_new_instance(
+    client: ExoClient, before_ids: set[str], timeout: float = 60.0
+) -> str:
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        current_ids = get_all_instance_ids(client)
+        new_ids = current_ids - before_ids
+        if new_ids:
+            return next(iter(new_ids))
+        time.sleep(0.2)
+    raise TimeoutError(f"No new instance appeared within {timeout}s")
+
+
 def wait_for_instance_gone(
     client: ExoClient, instance_id: str, timeout: float = 3.0
 ) -> None:

src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py

@@ -23,7 +23,7 @@ REDUCED_CONFIGS = {
         "num_hidden_layers": 2,
         "hidden_size": 256,
         "num_attention_heads": 8,
-        "num_key_value_heads": 4,
+        "num_key_value_heads": 2,
         "vocab_size": 1024,
         "intermediate_size": 512,
         "rms_norm_eps": 1e-5,
@@ -132,7 +132,7 @@ REDUCED_CONFIGS = {
             "num_hidden_layers": 4,
             "hidden_size": 256,
             "num_attention_heads": 8,
-            "num_key_value_heads": 4,
+            "num_key_value_heads": 2,
             "vocab_size": 1024,
             "intermediate_size": 512,
             "num_experts": 4,