Support dynamically rebalancing experts using EPLB (#6469)

2025-05-21 23:13:21 -07:00
parent 9484eba4ad
commit 7a80f56513
6 changed files with 226 additions and 3 deletions
@@ -210,8 +210,8 @@ def _execute_test(info: _TestInfo, rank: int, num_gpus: int, device: str):
            temp_buffers=expert_location_updater.create_temp_buffers(
                routed_experts_weights
            ),
-            old_physical_to_logical_map=physical_to_logical_map,
-            new_physical_to_logical_map=new_physical_to_logical_map,
+            old_physical_to_logical_map=physical_to_logical_map.tolist(),
+            new_physical_to_logical_map=new_physical_to_logical_map.tolist(),
            num_local_physical_experts=num_local_physical_experts,
            num_gpu_per_node=num_gpu_per_node,
            rank=rank,