@@ -19,181 +19,168 @@ import numpy
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import vllm.envs as envs
|
||||
from vllm.distributed.parallel_state import get_pp_group
|
||||
from vllm.logger import logger
|
||||
from vllm.v1.utils import record_function_or_nullcontext
|
||||
|
||||
from vllm_ascend.distributed.parallel_state import get_dynamic_eplb_group
|
||||
from vllm_ascend.eplb.adaptor.vllm_adaptor import VllmEplbAdaptor
|
||||
from vllm_ascend.eplb.core.eplb_device_transfer_loader import D2DExpertWeightLoader
|
||||
from vllm_ascend.eplb.core.eplb_worker import EplbProcess
|
||||
|
||||
|
||||
class EplbUpdator:
|
||||
|
||||
def __init__(self, ascend_config, loader, eplb_process: EplbProcess,
|
||||
process):
|
||||
self.ascend_config = ascend_config
|
||||
self.init_eplb(self.ascend_config.expert_map_path, process)
|
||||
def __init__(self, eplb_config, loader: D2DExpertWeightLoader, eplb_process: EplbProcess, process):
|
||||
self.eplb_config = eplb_config
|
||||
self.multi_stage = eplb_config.eplb_policy_type == 3
|
||||
self.init_eplb(self.eplb_config.expert_map_path, process)
|
||||
self.eplb_loader = loader
|
||||
self.eplb_process = eplb_process
|
||||
self.shared_dict = self.eplb_process.shared_dict
|
||||
self.comm_group = get_dynamic_eplb_group()
|
||||
|
||||
def set_adaptor(self, adaptor):
|
||||
def set_adaptor(self, adaptor: VllmEplbAdaptor):
|
||||
self.pp_rank = get_pp_group().rank_in_group
|
||||
self.adaptor = adaptor
|
||||
self.num_moe_layers = self.adaptor.num_moe_layers
|
||||
self.global_expert_num = self.adaptor.global_expert_num
|
||||
local_load = self.adaptor.get_rank_expert_workload()
|
||||
self.world_size = dist.get_world_size()
|
||||
self.device = local_load.device
|
||||
self.eplb_loader.num_layers = self.adaptor.num_dense_layers + self.adaptor.num_moe_layers
|
||||
|
||||
def init_eplb(self, expert_map_path, process):
|
||||
self.rank_id = dist.get_rank()
|
||||
self.num_expert_load_gather = 10
|
||||
self.periodic_load_gather = True
|
||||
self.num_iterations_eplb_update: torch.int64 = self.ascend_config.num_iterations_eplb_update
|
||||
self.expert_heat_collection_interval: torch.int64 = self.eplb_config.expert_heat_collection_interval
|
||||
self.expert_map_path = expert_map_path
|
||||
self.expert_map_record_path = self.ascend_config.expert_map_record_path
|
||||
self.expert_map_record_path = self.eplb_config.expert_map_record_path
|
||||
|
||||
try:
|
||||
if not envs.VLLM_ALLOW_EXPERT_LOAD_COLLECTING:
|
||||
self.num_expert_load_gather = self.num_iterations_eplb_update
|
||||
self.num_expert_load_gather = self.expert_heat_collection_interval
|
||||
self.periodic_load_gather = False
|
||||
except Exception:
|
||||
self.num_expert_load_gather = self.num_iterations_eplb_update
|
||||
logger.debug("[eplb/updator] VLLM_ALLOW_EXPERT_LOAD_COLLECTING unavailable in current vllm version.")
|
||||
self.num_expert_load_gather = self.expert_heat_collection_interval
|
||||
self.periodic_load_gather = False
|
||||
|
||||
self.expert_map_initialized = False
|
||||
self.gate_eplb = self.ascend_config.gate_eplb
|
||||
|
||||
self.reqs = []
|
||||
self.update_info_all = []
|
||||
|
||||
self.cur_iterations: torch.int64 = 0
|
||||
|
||||
self.num_wait_worker_iterations: torch.int64 = self.ascend_config.num_wait_worker_iterations
|
||||
self.algorithm_execution_interval: torch.int64 = self.eplb_config.algorithm_execution_interval
|
||||
|
||||
self.process = process
|
||||
|
||||
logger.info(
|
||||
f"[ModelRunner] Launched EPLB process (pid={self.process.pid})")
|
||||
logger.info("[eplb/updator] Launched EPLB subprocess, pid=%s", self.process.pid)
|
||||
|
||||
def update_iteration(self):
|
||||
self.cur_iterations += 1
|
||||
if self.cur_iterations == (self.num_iterations_eplb_update + \
|
||||
self.num_wait_worker_iterations + self.num_moe_layers):
|
||||
if self.cur_iterations == (
|
||||
self.expert_heat_collection_interval + self.algorithm_execution_interval + self.num_moe_layers
|
||||
):
|
||||
logger.debug("[eplb/updator] Full EPLB cycle completed, clearing moe loads and resetting iteration counter")
|
||||
if self.expert_map_record_path is not None:
|
||||
self.adaptor._export_tensor_to_file(
|
||||
self.shared_dict["expert_maps"],
|
||||
self.expert_map_record_path)
|
||||
self.adaptor._export_tensor_to_file(self.shared_dict["expert_maps"], self.expert_map_record_path)
|
||||
|
||||
self.adaptor.model.clear_all_moe_loads()
|
||||
if not self.gate_eplb:
|
||||
self.cur_iterations = 0
|
||||
self.adaptor.clear_all_moe_loads()
|
||||
self.cur_iterations = 0
|
||||
|
||||
def get_update_info_flag(self):
|
||||
return self.cur_iterations == (self.num_iterations_eplb_update +
|
||||
self.num_wait_worker_iterations - 1)
|
||||
return self.cur_iterations == (self.expert_heat_collection_interval + self.algorithm_execution_interval - 1)
|
||||
|
||||
def wakeup_eplb_worker_flag(self):
|
||||
return self.cur_iterations == (self.num_iterations_eplb_update - 1)
|
||||
return self.cur_iterations == (self.expert_heat_collection_interval - 1)
|
||||
|
||||
def update_expert_weight_flag(self):
|
||||
weight_update_counter = self.cur_iterations - (
|
||||
self.num_iterations_eplb_update + self.num_wait_worker_iterations)
|
||||
return (weight_update_counter >= 0
|
||||
and weight_update_counter < self.num_moe_layers)
|
||||
|
||||
def get_init_expert_map(self):
|
||||
try:
|
||||
if not self.expert_map_initialized:
|
||||
self.shared_dict[
|
||||
"expert_maps"] = self.adaptor.get_init_expert_map_from_file(
|
||||
self.num_moe_layers, self.expert_map_path)
|
||||
self.expert_map_initialized = True
|
||||
except Exception as e:
|
||||
logger.warning(f"[ModelRunner] Failed to wake EPLB process: {e}",
|
||||
exc_info=True)
|
||||
self.expert_heat_collection_interval + self.algorithm_execution_interval
|
||||
)
|
||||
return weight_update_counter >= 0 and weight_update_counter < self.num_moe_layers
|
||||
|
||||
def wakeup_eplb_worker(self):
|
||||
self.eplb_process.planner_q.put(1)
|
||||
|
||||
def forward_before(self):
|
||||
if self.update_expert_weight_flag():
|
||||
(expert_send_info, expert_recv_info, updated_expert_map,
|
||||
log2phy_map, layer_id) = self.update_info_all.pop(0)
|
||||
log2phy_map_this_rank = torch.from_numpy(numpy.array(log2phy_map))
|
||||
self.eplb_loader.set_log2phy_map(log2phy_map_this_rank)
|
||||
updated_expert_map_this_rank = torch.from_numpy(
|
||||
numpy.array(updated_expert_map))
|
||||
self.eplb_loader.generate_expert_d2d_transfer_task(
|
||||
expert_send_info, expert_recv_info,
|
||||
updated_expert_map_this_rank,
|
||||
layer_id + self.adaptor.num_dense_layers)
|
||||
|
||||
# set asynchronous stream for d2d expert weight update
|
||||
self.reqs = []
|
||||
self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
|
||||
|
||||
def take_update_info_from_eplb_process(self):
|
||||
# Batch after eplb process being triggered, get update info provided by eplb process
|
||||
if self.get_update_info_flag():
|
||||
self.update_info_all = self.eplb_process.block_update_q.get()
|
||||
|
||||
def forward_end(self):
|
||||
if self.wakeup_eplb_worker_flag():
|
||||
self.compute_and_set_moe_load(is_clear=True)
|
||||
self.wakeup_eplb_worker()
|
||||
|
||||
if self.update_expert_weight_flag():
|
||||
with record_function_or_nullcontext("EPLB generate p2p task"):
|
||||
(expert_send_info, expert_recv_info, updated_expert_map, log2phy_map, layer_id) = (
|
||||
self.update_info_all.pop(0)
|
||||
)
|
||||
log2phy_map_this_rank = torch.from_numpy(numpy.array(log2phy_map))
|
||||
self.eplb_loader.set_log2phy_map(log2phy_map_this_rank)
|
||||
updated_expert_map_this_rank = torch.from_numpy(numpy.array(updated_expert_map))
|
||||
self.eplb_loader.generate_expert_d2d_transfer_task(
|
||||
expert_send_info,
|
||||
expert_recv_info,
|
||||
updated_expert_map_this_rank,
|
||||
layer_id,
|
||||
)
|
||||
|
||||
# set asynchronous stream for d2d expert weight update
|
||||
self.reqs = []
|
||||
self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
|
||||
|
||||
def forward_end(self, eplb_heat_collection_status: bool = True):
|
||||
if self.wakeup_eplb_worker_flag():
|
||||
with record_function_or_nullcontext("EPLB gather moe load"):
|
||||
self.compute_and_set_moe_load()
|
||||
self.wakeup_eplb_worker()
|
||||
|
||||
if self.update_expert_weight_flag() and self.expert_map_record_path is None:
|
||||
self.eplb_loader.update_expert_map_and_weight(self.reqs)
|
||||
|
||||
self.update_iteration()
|
||||
# One circle of eplb update includes expert_heat_collection_interval + algorithm_execution_interval
|
||||
# + num_moe_layers (for weight update). In expert_heat_collection stage, we only update the counter
|
||||
# when eplb_heat_collection_status is True. In later stages, the counter is always updated.
|
||||
# TODO(Angazenn): Decouple algorithm execution && weight update with heat collection iterations.
|
||||
if self.cur_iterations >= self.expert_heat_collection_interval - 1 or eplb_heat_collection_status:
|
||||
self.update_iteration()
|
||||
|
||||
def compute_and_set_moe_load(self, is_clear=False):
|
||||
local_load = self.adaptor.get_rank_expert_workload()
|
||||
def compute_and_set_moe_load(self):
|
||||
local_load = self.adaptor.get_rank_expert_workload().unsqueeze(1)
|
||||
moe_load = self.comm_group.all_gather(local_load, dim=1).cpu()
|
||||
|
||||
self._gather_buffer = None
|
||||
if dist.is_initialized():
|
||||
self.world_size = dist.get_world_size()
|
||||
self.device = local_load.device
|
||||
if self._gather_buffer is None:
|
||||
shape = (self.world_size, *local_load.shape)
|
||||
self._gather_buffer = torch.empty(shape,
|
||||
dtype=local_load.dtype,
|
||||
device=self.device)
|
||||
if self.multi_stage:
|
||||
moe_load = moe_load.permute(2, 0, 1, 3)
|
||||
|
||||
dist.all_gather_into_tensor(self._gather_buffer, local_load)
|
||||
self.shared_dict["moe_load"] = moe_load
|
||||
logger.debug("[eplb/updator] Updated shared_dict['moe_load'] shape=%s", moe_load.shape)
|
||||
|
||||
moe_load = self._gather_buffer.permute(1, 0, 2)
|
||||
self.shared_dict["moe_load"] = moe_load.cpu()
|
||||
logger.debug(
|
||||
f"[ModelRunner] Updated shared_dict['moe_load'] shape={moe_load.shape}"
|
||||
)
|
||||
else:
|
||||
moe_load = local_load.unsqueeze(1)
|
||||
self.shared_dict["moe_load"] = moe_load.cpu()
|
||||
logger.debug(
|
||||
f"[ModelRunner] Updated shared_dict['moe_load'] shape={moe_load.shape}"
|
||||
)
|
||||
return moe_load
|
||||
|
||||
def warm_up_eplb(self):
|
||||
|
||||
self.get_init_expert_map()
|
||||
logger.info("[eplb/updator] Starting EPLB warm-up, rank=%s, world_size=%s", self.rank_id, self.world_size)
|
||||
self.shared_dict["expert_maps"] = self.adaptor.get_global_expert_map()
|
||||
self.compute_and_set_moe_load()
|
||||
|
||||
src_tensor = torch.empty((1, ), device=self.device)
|
||||
self_rank = dist.get_rank()
|
||||
src_tensor = torch.empty((1,), device=self.device)
|
||||
|
||||
comm_op_list = []
|
||||
reqs = []
|
||||
|
||||
for dst_rank in range(self.world_size):
|
||||
if dst_rank == self_rank:
|
||||
for dst_rank in range(self.comm_group.world_size):
|
||||
if dst_rank == self.comm_group.rank_in_group:
|
||||
continue
|
||||
comm_op_list.append(dist.P2POp(dist.isend, src_tensor, dst_rank))
|
||||
global_dst = self.comm_group.ranks[dst_rank]
|
||||
comm_op_list.append(dist.P2POp(dist.isend, src_tensor, global_dst, group=self.comm_group.device_group))
|
||||
|
||||
for src_rank in range(self.world_size):
|
||||
if src_rank == self_rank:
|
||||
for src_rank in range(self.comm_group.world_size):
|
||||
if src_rank == self.comm_group.rank_in_group:
|
||||
continue
|
||||
comm_op_list.append(dist.P2POp(dist.irecv, src_tensor, src_rank))
|
||||
global_src = self.comm_group.ranks[src_rank]
|
||||
comm_op_list.append(dist.P2POp(dist.irecv, src_tensor, global_src, group=self.comm_group.device_group))
|
||||
if comm_op_list:
|
||||
reqs = dist.batch_isend_irecv(comm_op_list)
|
||||
|
||||
for req in reqs:
|
||||
req.wait()
|
||||
logger.info("[eplb/updator] EPLB warm-up completed")
|
||||
|
||||
def shutdown(self):
|
||||
"""
|
||||
@@ -202,4 +189,4 @@ class EplbUpdator:
|
||||
if self.process.is_alive():
|
||||
self.process.terminate()
|
||||
self.process.join()
|
||||
logger.info("[ModelRunner] EPLB process terminated")
|
||||
logger.info("[eplb/updator] EPLB subprocess terminated")
|
||||
|
||||
Reference in New Issue
Block a user