clean debug logs

exo/inference/mlx/models/sharded_llama.py

@@ -90,6 +90,8 @@ class Attention(nn.Module):
     ) -> mx.array:
         B, L, D = x.shape
 
+        print("q_proj: ", self.q_proj)
+        print("x: ", x.shape)
         queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
 
         # Prepare the queries, keys and values for the attention computation
@@ -190,7 +192,8 @@ class LlamaModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        for layer, c in zip(self.layers, cache):
+        for i, (layer, c) in enumerate(zip(self.layers, cache)):
+            print(f"layer: {i}")
             h = layer(h, mask, cache=c)
 
         if self.args.shard.is_last_layer():

exo/inference/mlx/test_sharded_llama.py

@@ -3,6 +3,7 @@ from exo.inference.mlx.sharded_model import StatefulShardedModel
 from exo.inference.mlx.sharded_utils import load_shard
 from exo.inference.shard import Shard
 
+# 79, 80 for Llama-3-70B
 shard_full = Shard("llama", 0, 31, 32)
 shard1 = Shard("llama", 0, 12, 32)
 shard2 = Shard("llama", 13, 31, 32)
@@ -16,7 +17,7 @@ m1 = StatefulShardedModel(shard1, model_shard1)
 m2 = StatefulShardedModel(shard2, model_shard2)
 
 prompt = "write a beautiful haiku about a utopia where people own their AI with edge intelligence:"
-prompt_tokens = mx.array(tokenizer1.encode(prompt))
+prompt_tokens = mx.array(full_tokenizer.encode(prompt))
 max_tokens = 50
 
 resp = prompt_tokens
@@ -25,7 +26,7 @@ for _ in range(max_tokens):
     resp = full.step(resp)
     full_generated_tokens.append(resp.item())
 
-print("full response: ", tokenizer1.decode(full_generated_tokens))
+print("full response: ", full_tokenizer.decode(full_generated_tokens))
 
 
 sharded_generated_tokens = []

exo/orchestration/standard_node.py

@@ -7,6 +7,7 @@ from exo.topology.topology import Topology
 from exo.topology.device_capabilities import device_capabilities
 from exo.topology.partitioning_strategy import PartitioningStrategy
 from exo.topology.partitioning_strategy import Partition
+from exo import DEBUG
 import asyncio
 import uuid
 
@@ -29,7 +30,7 @@ class StandardNode(Node):
         await self.discovery.start()
         await self.update_peers(wait_for_peers)
         await self.collect_topology()
-        print(f"Collected topology: {self.topology}")
+        if DEBUG >= 2: print(f"Collected topology: {self.topology}")
         asyncio.create_task(self.periodic_topology_collection(5))
 
     async def stop(self) -> None:
@@ -42,7 +43,7 @@ class StandardNode(Node):
         if request_id not in self.buffered_token_output:
             self.buffered_token_output[request_id] = ([], False)
 
-        print(f"[{request_id}] process prompt: {shard}, {prompt}")
+        if DEBUG >= 2: print(f"[{request_id}] process prompt: {shard}, {prompt}")
         result, is_finished = await self.inference_engine.infer_prompt(self.get_current_shard(shard), prompt)
         is_finished = is_finished or len(self.buffered_token_output[request_id]) >= self.max_generate_tokens
         if is_finished:
@@ -52,7 +53,7 @@ class StandardNode(Node):
             self.buffered_token_output[request_id][0].append(result.item())
             self.on_token(self.buffered_token_output[request_id][0])
 
-        print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id])}")
+        if DEBUG >= 2: print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id])}")
 
         if not is_finished:
             asyncio.create_task(self.forward_tensor_to_next_shard(shard, result, request_id))
@@ -66,7 +67,7 @@ class StandardNode(Node):
             self.buffered_token_output[request_id] = ([], False)
 
         try:
-            print(f"[{request_id}] process_tensor: {shard}, {tensor}")
+            if DEBUG >= 2: print(f"[{request_id}] process_tensor: {shard}, {tensor}")
             result, is_finished = await self.inference_engine.infer_tensor(self.get_current_shard(shard), tensor)
             is_finished = is_finished or len(self.buffered_token_output[request_id]) >= self.max_generate_tokens
             if is_finished:
@@ -75,7 +76,7 @@ class StandardNode(Node):
             if result.size == 1:  # we got a new token out
                 self.buffered_token_output[request_id][0].append(result.item())
                 self.on_token(self.buffered_token_output[request_id][0])
-            print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id])}")
+            if DEBUG >= 2: print(f"[{request_id}] result size: {result.size}, is finished: {is_finished}, buffered tokens: {len(self.buffered_token_output[request_id])}")
 
             if not is_finished:
                 asyncio.create_task(self.forward_tensor_to_next_shard(shard, result, request_id))
@@ -89,16 +90,16 @@ class StandardNode(Node):
 
     async def forward_tensor_to_next_shard(self, shard: Shard, tensor: np.ndarray, request_id: str) -> None:
         if not self.partitioning_strategy:
-            print("No partitioning strategy found. Skipping forward.")
+            if DEBUG >= 1: print("No partitioning strategy found. Skipping forward.")
             return
 
         partitions = self.partitioning_strategy.partition(self.topology)
         current_partition_index = next((i for i, p in enumerate(partitions) if p.node_id == self.id), None)
-        print(f"Current partition index: {current_partition_index}")
+        if DEBUG >= 2: print(f"Current partition index: {current_partition_index}")
         if current_partition_index is not None:
             next_partition_index = (current_partition_index + 1) % len(partitions)
             next_partition: Partition = partitions[next_partition_index]
-            print(f"Computed next from: {shard}, {self.topology}. Next partition: {next_partition}")
+            if DEBUG >= 2: print(f"Computed next from: {shard}, {self.topology}. Next partition: {next_partition}")
 
             if next_partition:
                 if next_partition.node_id == self.id:
@@ -113,7 +114,7 @@ class StandardNode(Node):
                 end_layer = int(next_partition.end * shard.n_layers) - 1
                 next_shard = Shard(shard.model_id, start_layer, end_layer, shard.n_layers)
 
-                print(f"Sending tensor to {target_peer.id()} for shard: {next_shard}: {tensor}")
+                if DEBUG >= 2: print(f"Sending tensor to {target_peer.id()} for shard: {next_shard}: {tensor}")
 
                 await target_peer.send_tensor(next_shard, tensor, request_id)
 
@@ -131,20 +132,20 @@ class StandardNode(Node):
 
     async def reset_shard(self, shard: Shard) -> None:
         # Implement shard reset logic
-        print(f"Resetting shard: {shard}")
+        if DEBUG >= 2: print(f"Resetting shard: {shard}")
         self.buffered_token_output = {}
         await self.inference_engine.reset_shard(self.get_current_shard(shard))
 
     async def update_peers(self, wait_for_peers: int = 0) -> None:
         self.peers = await self.discovery.discover_peers(wait_for_peers)
-        print(f"Starting with the following peers: {self.peers}")
-        print("Connecting to new peers...")
+        if DEBUG >= 2: print(f"Starting with the following peers: {self.peers}")
+        if DEBUG >= 2: print("Connecting to new peers...")
         for peer in self.peers:
             is_connected = await peer.is_connected()
-            print(f"Connected to {peer.id()}: {is_connected}")
+            if DEBUG >= 2: print(f"Connected to {peer.id()}: {is_connected}")
             if not is_connected:
                 await peer.connect()
-                print(f"Connected to peer {peer.id()}")
+                if DEBUG >= 2: print(f"Connected to peer {peer.id()}")
 
     async def collect_topology(self, max_depth: int = 4) -> Topology:
         self.topology.update_node(self.id, self.device_capabilities)
@@ -156,7 +157,7 @@ class StandardNode(Node):
             if max_depth > 0:
                 try:
                     other_topology = await peer.collect_topology(max_depth = max_depth - 1)
-                    print(f"Collected topology from: {peer.id()}: {other_topology}")
+                    if DEBUG >= 2: print(f"Collected topology from: {peer.id()}: {other_topology}")
                     self.topology.merge(other_topology)
                 except Exception as e:
                     print(f"Error collecting topology from {peer.id()}: {e}")
@@ -172,8 +173,8 @@ class StandardNode(Node):
             except Exception as e:
                 print(f"Error collecting topology: {e}")
 
-            print("Topology collection task executed.")
-            print(f"Current topology: {self.topology}")
+            if DEBUG >= 2: print("Topology collection task executed.")
+            if DEBUG >= 2: print(f"Current topology: {self.topology}")
 
     async def get_inference_result(self, request_id: str) -> Tuple[Optional[np.ndarray], bool]:
         if request_id not in self.buffered_token_output: