vllm.distributed.eplb.async_worker

The async worker that transfers experts in the background.

_transfer_cycle(state, stream, is_profile)

Run one rebalance pass over all model states and layers.

Source code in vllm/distributed/eplb/async_worker.py

def _transfer_cycle(
    state: "EplbState",
    stream: torch.Stream,
    is_profile: bool,
) -> None:
    """Run one rebalance pass over all model states and layers."""
    # Snapshot the FT reset epoch. A change in this value means a fault
    # recovery happened, and the in-flight rebalance must be abandoned.
    reset_epoch = state.ft_reset_epoch

    eplb_group = get_eplb_group().device_group
    eplb_cpu_group = get_eplb_group().cpu_group
    ep_rank = eplb_group.rank()

    assert state.is_async
    for model_state in state.model_states.values():
        layer_idx = 0
        # Set the async worker's CUDA stream on the communicator
        model_state.communicator.set_stream(stream)
        num_layers = model_state.model.num_moe_layers

        # Snapshot the physical_to_logical_map (synchronized with
        # rearrange_event) and copy it to CPU
        with device_stream(stream):
            physical_to_logical_map_cpu = model_state.physical_to_logical_map.cpu()

        t_start = time.perf_counter()
        new_physical_to_logical_map = run_rebalance_experts(
            model_state, state, physical_to_logical_map_cpu, stream
        )

        migration_batches_by_layer = None
        if state.parallel_config.eplb_config.migration_batching_enabled:
            profile_name = "eplb: schedule migration batches"
            num_local_experts = model_state.model.expert_weights[0][0].shape[0]
            # The placement snapshots stay fixed for this rebalance cycle.
            # Precompute all layer schedules once, but execute transfers one
            # layer at a time below to preserve the migration lifecycle. This
            # profiling range therefore covers scheduling for every MoE layer.
            with torch.profiler.record_function(profile_name):
                migration_batches_by_layer = schedule_migration_batches_for_layers(
                    num_local_experts,
                    physical_to_logical_map_cpu.numpy(),
                    new_physical_to_logical_map.numpy(),
                )

        # Execute one EPLB layer transfer per model forward pass. Each iteration
        # of this loop will copy the new set of expert weights into
        # model_state.expert_buffer, which will be consumed by the main thread in
        # move_to_workspace.
        # We sync the rebalanced flag across ranks before each iteration so
        # all ranks make a coordinated decision to continue or stop.
        while layer_idx < num_layers:
            flag = torch.tensor(
                [int(model_state.rebalanced)],
                dtype=torch.int32,
                device="cpu",
            )
            torch.distributed.all_reduce(flag, group=eplb_cpu_group)
            if int(flag.item()) != eplb_cpu_group.size():
                logger.warning(
                    "async worker (rank=%d): layer %d coordinated stop "
                    "(flag_sum=%d, group_size=%d)",
                    ep_rank,
                    layer_idx,
                    int(flag.item()),
                    eplb_cpu_group.size(),
                )
                model_state.rebalanced = False
                break

            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],
                expert_weights_buffer=model_state.expert_buffer,
                communicator=model_state.communicator,
                ep_group=eplb_group,
                is_profile=is_profile,
                stream=stream,
                layer_idx=layer_idx,
                enable_migration_batching=(
                    state.parallel_config.eplb_config.migration_batching_enabled
                ),
                migration_batches=(
                    # This is only the current layer's schedule; it does not
                    # combine transfers from different layers.
                    migration_batches_by_layer[layer_idx]
                    if migration_batches_by_layer is not None
                    else None
                ),
            )

            # Wait until all writes to expert_buffer have finished before making the
            # AsyncEplbLayerResult visible to the main thread.
            stream.synchronize()

            # This event guarantees that expert_buffer will not be overwritten by
            # subsequent iterations of this loop until the main thread has consumed
            # it. Record is called by the main thread after move_from_buffer().
            consumed_event = CpuGpuEvent()

            model_state.pending_result = AsyncEplbLayerResult(
                layer_idx=layer_idx,
                new_physical_to_logical_map=new_physical_to_logical_map[layer_idx],
                transfer_metadata=transfer_metadata,
                consumed_event=consumed_event,
            )
            # A fault recovery happened during the transfer, so the main
            # thread will never consume this result. Drop it and abandon
            # the rebalance.
            if state.ft_reset_epoch != reset_epoch:
                model_state.pending_result = None
                return

            # Block this thread until the main thread and main stream
            # finish copying model_state.expert_buffer into
            # model_state.model.expert_weights[layer_idx]
            consumed_event.wait(stream=stream)
            # Host-issued writes into expert_buffer (e.g. NIXL RDMA reads) are not
            # stream-ordered, so wait until move_from_buffer has run on the GPU.
            stream.synchronize()
            assert model_state.pending_result is None
            # Woken by an FT reset's record, not a normal consume.
            if state.ft_reset_epoch != reset_epoch:
                return
            layer_idx += 1
        if ep_rank == 0:
            logger.info(
                "Async EPLB rearranged %d/%d layers of %s in %.2f s.",
                layer_idx,
                num_layers,
                model_state.model_name,
                time.perf_counter() - t_start,
            )