diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 7f8895cd2c..d9e6a739b0 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import asyncio import random import pytest @@ -361,16 +360,14 @@ def _test_async_transfer_layer_without_mtp_worker( communicator.set_stream(cuda_stream) for layer_idx in range(num_layers): - transfer_metadata = asyncio.run( - transfer_layer( - old_layer_indices=old_indices_cpu[layer_idx], - new_layer_indices=new_indices_cpu[layer_idx], - expert_weights=expert_weights[layer_idx], - expert_weights_buffer=expert_buffer, - ep_group=ep_group, - communicator=communicator, - cuda_stream=cuda_stream, - ) + transfer_metadata = transfer_layer( + old_layer_indices=old_indices_cpu[layer_idx], + new_layer_indices=new_indices_cpu[layer_idx], + expert_weights=expert_weights[layer_idx], + expert_weights_buffer=expert_buffer, + ep_group=ep_group, + communicator=communicator, + cuda_stream=cuda_stream, ) cuda_stream.synchronize() move_from_buffer( diff --git a/vllm/distributed/eplb/async_worker.py b/vllm/distributed/eplb/async_worker.py index a47b5ce29c..542606fe74 100644 --- a/vllm/distributed/eplb/async_worker.py +++ b/vllm/distributed/eplb/async_worker.py @@ -4,7 +4,6 @@ The async worker that transfers experts in the background. """ -import asyncio import threading from typing import TYPE_CHECKING @@ -36,21 +35,15 @@ def start_async_worker( assert device_index is not None torch.accelerator.set_device_index(device_index) cuda_stream = torch.cuda.Stream(device=device_index) - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) try: - loop.run_until_complete( - transfer_run_periodically( - state=state, - eplb_group=eplb_group, - cuda_stream=cuda_stream, - is_profile=is_profile, - ) + transfer_run_periodically( + state=state, + eplb_group=eplb_group, + cuda_stream=cuda_stream, + is_profile=is_profile, ) except Exception as exc: # pragma: no cover - diagnostic path logger.exception("async loop error (Rank %d): %s", rank, str(exc)) - finally: - loop.close() thread = threading.Thread(target=thread_target, daemon=True) thread.start() @@ -83,7 +76,7 @@ def run_rebalance_experts( return new_physical_to_logical_map -async def transfer_run_periodically( +def transfer_run_periodically( state: "EplbState", eplb_group: ProcessGroup, cuda_stream: torch.cuda.Stream, @@ -118,7 +111,7 @@ async def transfer_run_periodically( # model_state.expert_buffer, which will be consumed by the main thread in # move_to_workspace while model_state.rebalanced and layer_idx < num_layers: - transfer_metadata = await transfer_layer( + transfer_metadata = transfer_layer( old_layer_indices=physical_to_logical_map_cpu[layer_idx], new_layer_indices=new_physical_to_logical_map[layer_idx], expert_weights=model_state.model.expert_weights[layer_idx], diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index a68fbda86c..f348521c00 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -418,7 +418,7 @@ def move_from_buffer( w[dst].copy_(w[src], non_blocking=True) -async def transfer_layer( +def transfer_layer( old_layer_indices: torch.Tensor, new_layer_indices: torch.Tensor, expert_weights: Sequence[torch.Tensor],