rebase fix

Signed-off-by: Yongji Wu <[email protected]>

Author: libertyeagle

experimental/bench.sh

@@ -58,24 +58,25 @@ def batch_transfer_weights(
     state_dict = model.state_dict()
     all_params = []
 
-    for _, param in state_dict.items():
+    for name, param in state_dict.items():
+        if name.endswith("expert_map"):
+            continue
         if param.data_ptr() not in expert_weights_set:
             all_params.append(param.data)
 
-    if all_params:
-        p2p_ops = []
-        for param in all_params:
-            op = object.__new__(P2POp)
-            if is_sender:
-                op.op = torch.distributed.isend
-                op.tensor = param
-            else:
-                op.op = torch.distributed.irecv
-                op.tensor = param
-            op.group_peer = peer_rank
-            p2p_ops.append(op)
-
-        device_comm.batch_isend_irecv(p2p_ops)
+    assert len(all_params) > 0
+    p2p_ops = []
+    for param in all_params:
+        op = object.__new__(P2POp)
+        if is_sender:
+            op.op = torch.distributed.isend
+            op.tensor = param
+        else:
+            op.op = torch.distributed.irecv
+            op.tensor = param
+        op.group_peer = peer_rank
+        p2p_ops.append(op)
+    device_comm.batch_isend_irecv(p2p_ops)
 
 
 def broadcast_expert_mapping(

experimental/scale.sh

@@ -892,7 +892,7 @@ def _ensure_output_queue_task(self):
         assert output_socket is not None
 
         notification_callback_handler: Optional[
-            Callable[[AsyncMPClient, tuple[str, int]], Any]
+            Callable[[AsyncMPClient, Sequence[Any]], Any]
         ] = getattr(self.__class__, "process_worker_notification", None)
 
         async def process_outputs_socket():
@@ -913,7 +913,7 @@ async def process_outputs_socket():
                             if outputs.utility_output.result is None:
                                 continue
                             notification_data = outputs.utility_output.result.result
-                            assert isinstance(notification_data, tuple)
+                            assert isinstance(notification_data, Sequence)
                             assert len(notification_data) == 2
                             asyncio.create_task(
                                 notification_callback_handler(_self, notification_data)