EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Lint]Style: Convert vllm-ascend/ to ruff format(Batch #6) (#6001)

Browse Source 
### What this PR does / why we need it?
| File Path |
| :--- |
| ` vllm_ascend/eplb/adaptor/abstract_adaptor.py` |
| ` vllm_ascend/eplb/adaptor/vllm_adaptor.py` |
| ` vllm_ascend/eplb/core/eplb_device_transfer_loader.py` |
| ` vllm_ascend/eplb/core/eplb_utils.py` |
| ` vllm_ascend/eplb/core/eplb_worker.py` |
| ` vllm_ascend/eplb/core/policy/policy_abstract.py` |
| ` vllm_ascend/eplb/core/policy/policy_default_eplb.py` |
| ` vllm_ascend/eplb/core/policy/policy_factory.py` |
| ` vllm_ascend/eplb/core/policy/policy_flashlb.py` |
| ` vllm_ascend/eplb/core/policy/policy_random.py` |
| ` vllm_ascend/eplb/core/policy/policy_swift_balancer.py` |
| ` vllm_ascend/eplb/eplb_updator.py` |
| ` vllm_ascend/eplb/utils.py` |
| ` vllm_ascend/model_loader/netloader/executor/elastic_load.py` |
| ` vllm_ascend/model_loader/netloader/executor/netloader_pg.py` |
| ` vllm_ascend/model_loader/netloader/interaction/elastic.py` |
| ` vllm_ascend/model_loader/netloader/load.py` |
| ` vllm_ascend/model_loader/netloader/netloader.py` |
| ` vllm_ascend/model_loader/netloader/utils.py` |
| ` vllm_ascend/patch/platform/__init__.py` |
| ` vllm_ascend/patch/platform/patch_balance_schedule.py` |
| ` vllm_ascend/patch/platform/patch_ec_connector.py` |
| ` vllm_ascend/patch/platform/patch_mamba_config.py` |
| ` vllm_ascend/patch/platform/patch_multiproc_executor.py` |
| ` vllm_ascend/patch/platform/patch_sched_yield.py` |


- vLLM version: v0.13.0
- vLLM main:
2c24bc6996

---------

Signed-off-by: MrZ20 <2609716663@qq.com>
This commit is contained in:
SILONG ZENG
2026-01-24 22:08:33 +08:00
committed by GitHub
parent 153da1a669
commit 4e53c1d900
26 changed files with 894 additions and 1148 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
vllm_ascend/eplb/core/policy/policy_abstract.py
View File

@@ -12,7 +12,6 @@ class DynamicConfig:
class EplbPolicy:
def __init__(self, config: DynamicConfig):
self.config = config

178
vllm_ascend/eplb/core/policy/policy_default_eplb.py
View File

@@ -25,13 +25,11 @@ class DynamicTable:
class DefaultEplb(EplbPolicy):
def __init__(self, config: DynamicConfig):
super().__init__(config)
@staticmethod
def add_redundant(current_expert_table, expert_workload,
num_original_expert):
def add_redundant(current_expert_table, expert_workload, num_original_expert):
layer_num, npu_num, experts_per_npu = expert_workload.shape
workload_new = np.zeros((layer_num, num_original_expert))
for layer_idx in range(layer_num):
@@ -40,31 +38,24 @@ class DefaultEplb(EplbPolicy):
workload_layer = expert_workload[layer_idx].copy()
for npu_idx in range(npu_num):
for expert_idx in range(experts_per_npu):
workload_dict[placement_layer[npu_idx][
expert_idx]] += workload_layer[npu_idx][expert_idx]
workload_dict[placement_layer[npu_idx][expert_idx]] += workload_layer[npu_idx][expert_idx]
for expert_idx in range(num_original_expert):
workload_new[layer_idx][expert_idx] = workload_dict[expert_idx]
return workload_new
@staticmethod
# Split hot (high-load) experts into redundant experts
def original_compute_balanced_pack_redundancy(origin_weights, card_num,
num_redundancy_expert):
def original_compute_balanced_pack_redundancy(origin_weights, card_num, num_redundancy_expert):
# Step 1: Sort the items by weight in descending order (we are sorting by weight now)
# Sort based on the second element (the second value of each tuple)
route_expert_num = len(origin_weights)
route_expert_redundancy: list[list[int]] = [
[] for _ in range(route_expert_num)
]
route_expert_redundancy: list[list[int]] = [[] for _ in range(route_expert_num)]
for i in range(num_redundancy_expert):
sorted_indices = np.argsort([t[1] for t in origin_weights],
kind='stable')[::-1]
sorted_indices = np.argsort([t[1] for t in origin_weights], kind="stable")[::-1]
weights = [origin_weights[idx] for idx in sorted_indices]
tmp_raw_weight = weights[0][1] * (
len(route_expert_redundancy[weights[0][0]]) + 1)
tmp_raw_weight = weights[0][1] * (len(route_expert_redundancy[weights[0][0]]) + 1)
route_expert_redundancy[weights[0][0]].append(route_expert_num + i)
avg_weight = tmp_raw_weight / (
len(route_expert_redundancy[weights[0][0]]) + 1)
avg_weight = tmp_raw_weight / (len(route_expert_redundancy[weights[0][0]]) + 1)
weights[0] = (weights[0][0], avg_weight)
origin_weights = weights
@@ -93,8 +84,7 @@ class DefaultEplb(EplbPolicy):
box_counts[index] += 1
index += 1
sorted_indices = np.argsort([t[1] for t in origin_weights],
kind='stable')[::-1]
sorted_indices = np.argsort([t[1] for t in origin_weights], kind="stable")[::-1]
origin_weights = [origin_weights[idx] for idx in sorted_indices]
# Step 4: Distribute items into boxes based on weight
for item_id, weight in origin_weights:
@@ -104,11 +94,8 @@ class DefaultEplb(EplbPolicy):
if item_id in boxes[i]:
continue
# Only choose boxes that still have space (box_counts[i] < items_per_box)
if box_counts[i] < items_per_box or (box_counts[i]
== items_per_box
and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[
min_box_index]:
if box_counts[i] < items_per_box or (box_counts[i] == items_per_box and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[min_box_index]:
min_box_index = i
# Place the item (id) into the selected box
@@ -118,40 +105,35 @@ class DefaultEplb(EplbPolicy):
box_counts[min_box_index] += 1
# If there's an imbalance in the remaining items, reduce the "remaining_items" counter
if box_counts[min_box_index] == (items_per_box +
1) and remaining_items > 0:
if box_counts[min_box_index] == (items_per_box + 1) and remaining_items > 0:
remaining_items -= 1
# Step 5: Output each box's contents and total weight
result = []
for i in range(card_num):
result.append({
"box_index": i + 1,
"items": boxes[i], # List of item IDs in the box
"weight": boxes_weights[i],
"total_weight": box_weights[i], # Total weight in this box
"item_count": box_counts[i] # Number of items in the box
})
result.append(
{
"box_index": i + 1,
"items": boxes[i], # List of item IDs in the box
"weight": boxes_weights[i],
"total_weight": box_weights[i], # Total weight in this box
"item_count": box_counts[i], # Number of items in the box
}
)
return result, boxes
# Split hot (high-load) experts into redundant experts
@staticmethod
def compute_balanced_pack_redundancy(origin_weights, card_num,
num_redundancy_expert):
def compute_balanced_pack_redundancy(origin_weights, card_num, num_redundancy_expert):
route_expert_num = len(origin_weights)
route_expert_redundancy: list[list[int]] = [
[] for _ in range(route_expert_num)
]
route_expert_redundancy: list[list[int]] = [[] for _ in range(route_expert_num)]
for i in range(num_redundancy_expert):
sorted_indices = np.argsort([t[1] for t in origin_weights],
kind='stable')[::-1]
sorted_indices = np.argsort([t[1] for t in origin_weights], kind="stable")[::-1]
weights = [origin_weights[idx] for idx in sorted_indices]
tmp_raw_weight = weights[0][1] * (
len(route_expert_redundancy[weights[0][0]]) + 1)
tmp_raw_weight = weights[0][1] * (len(route_expert_redundancy[weights[0][0]]) + 1)
route_expert_redundancy[weights[0][0]].append(route_expert_num + i)
avg_weight = tmp_raw_weight / (
len(route_expert_redundancy[weights[0][0]]) + 1)
avg_weight = tmp_raw_weight / (len(route_expert_redundancy[weights[0][0]]) + 1)
weights[0] = (weights[0][0], avg_weight)
origin_weights = weights
@@ -166,7 +148,7 @@ class DefaultEplb(EplbPolicy):
box_weights = [0] * card_num
box_counts = [0] * card_num
all_weights = np.zeros((expert_num, ), dtype='object')
all_weights = np.zeros((expert_num,), dtype="object")
all_weights[:route_expert_num] = origin_weights
index = route_expert_num
@@ -178,17 +160,13 @@ class DefaultEplb(EplbPolicy):
all_weights[index] = (item, weight)
index += 1
sorted_indices = np.argsort([t[1] for t in all_weights],
kind='stable')[::-1]
sorted_indices = np.argsort([t[1] for t in all_weights], kind="stable")[::-1]
all_weights = [all_weights[idx] for idx in sorted_indices]
for item_id, weight in all_weights:
min_box_index = -1
for i in range(card_num):
if box_counts[i] < items_per_box or (box_counts[i]
== items_per_box
and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[
min_box_index]:
if box_counts[i] < items_per_box or (box_counts[i] == items_per_box and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[min_box_index]:
if item_id not in boxes[i]:
min_box_index = i
@@ -197,19 +175,20 @@ class DefaultEplb(EplbPolicy):
box_weights[min_box_index] += weight
box_counts[min_box_index] += 1
if box_counts[min_box_index] == (items_per_box +
1) and remaining_items > 0:
if box_counts[min_box_index] == (items_per_box + 1) and remaining_items > 0:
remaining_items -= 1
result = []
for i in range(card_num):
result.append({
"box_index": i + 1,
"items": boxes[i],
"weight": boxes_weights[i],
"total_weight": box_weights[i],
"item_count": box_counts[i]
})
result.append(
{
"box_index": i + 1,
"items": boxes[i],
"weight": boxes_weights[i],
"total_weight": box_weights[i],
"item_count": box_counts[i],
}
)
return result, boxes
@@ -232,11 +211,8 @@ class DefaultEplb(EplbPolicy):
for item_id, weight in weights:
min_box_index = -1
for i in range(card_num):
if box_counts[i] < items_per_box or (box_counts[i]
== items_per_box
and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[
min_box_index]:
if box_counts[i] < items_per_box or (box_counts[i] == items_per_box and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[min_box_index]:
min_box_index = i
boxes[min_box_index].append(item_id)
@@ -244,19 +220,20 @@ class DefaultEplb(EplbPolicy):
box_weights[min_box_index] += weight
box_counts[min_box_index] += 1
if box_counts[min_box_index] == (items_per_box +
1) and remaining_items > 0:
if box_counts[min_box_index] == (items_per_box + 1) and remaining_items > 0:
remaining_items -= 1
result = []
for i in range(card_num):
result.append({
"box_index": i + 1,
"items": boxes[i],
"weight": boxes_weights[i],
"total_weight": box_weights[i],
"item_count": box_counts[i]
})
result.append(
{
"box_index": i + 1,
"items": boxes[i],
"weight": boxes_weights[i],
"total_weight": box_weights[i],
"item_count": box_counts[i],
}
)
return result, boxes
@@ -274,16 +251,11 @@ class DefaultEplb(EplbPolicy):
return max_heat_per_layer
@staticmethod
def constraint_expert_local_exchange(current_expert_table,
global_deployment):
def constraint_expert_local_exchange(current_expert_table, global_deployment):
for layer_id in range(len(global_deployment)):
for card_id in range(len(global_deployment[layer_id])):
current_list = [
int(x) for x in current_expert_table[layer_id][card_id]
]
new_list = [
int(x) for x in global_deployment[layer_id][card_id]
]
current_list = [int(x) for x in current_expert_table[layer_id][card_id]]
new_list = [int(x) for x in global_deployment[layer_id][card_id]]
num = len(new_list)
new_index = [-1] * num
@@ -293,8 +265,7 @@ class DefaultEplb(EplbPolicy):
for i in range(num):
flag = True
for j in range(num):
if new_list[i] == current_list[j] and new_index[
j] == -1:
if new_list[i] == current_list[j] and new_index[j] == -1:
new_index[j] = 0
new_result[j] = current_list[j]
flag = False
@@ -313,7 +284,6 @@ class DefaultEplb(EplbPolicy):
return global_deployment
def rebalance_experts(self, current_expert_table, expert_workload):
info = DynamicTable()
info.workload_table = np.array(expert_workload)
info.placement_table = np.array(current_expert_table)
@@ -324,17 +294,15 @@ class DefaultEplb(EplbPolicy):
expert_ids, counts = np.unique(row, return_counts=True)
num_redundancy_expert = self.get_redundant_num(num_npus, counts)
num_original_expert = len(expert_ids)
layer_workloads = self.add_redundant(info.placement_table,
info.workload_table,
num_original_expert)
max_heat_per_layer_before = self.calculate_max_heat_per_layer(
info.workload_table, layer_num)
layer_workloads = self.add_redundant(info.placement_table, info.workload_table, num_original_expert)
max_heat_per_layer_before = self.calculate_max_heat_per_layer(info.workload_table, layer_num)
npu_heat_all_origin = sum(max_heat_per_layer_before)
# Perform load balancing and deploy redundant experts
layer_num = layer_workloads.shape[0]
expert_num = layer_workloads.shape[1]
# Validate that the number of experts, number of cards, and number of redundant experts do not exceed the number of cards
# Validate that the number of experts, number of cards, and number of redundant experts
# do not exceed the number of cards.
if num_original_expert != expert_num:
raise ValueError(
f"the number of original experts {num_original_expert} must be equal to expert_num {expert_num}"
@@ -345,38 +313,35 @@ class DefaultEplb(EplbPolicy):
if num_npus < num_redundancy_expert:
raise ValueError(
f"the number of NPUs {num_npus} must be greater than or equal to the number of redundant experts {num_redundancy_expert}"
"the number of NPUs "
f"{num_npus} must be greater than or equal to the number of redundant experts "
f"{num_redundancy_expert}"
)
# Number of experts deployed on each card includes one redundant expert
global_deployment: list[list[list[int]]] = [[[]
for _ in range(num_npus)]
for _ in range(layer_num)]
global_deployment: list[list[list[int]]] = [[[] for _ in range(num_npus)] for _ in range(layer_num)]
# Iterate to obtain the placement strategy for each layer, taking computational balance into account
max_heat_per_layer_after = np.zeros([layer_num])
for layer in range(layer_num):
# Get the expert IDs and their corresponding workloads for the current layer;
# workloads need to be normalized, and one redundant expert is added per card
weights = np.zeros((expert_num, ), dtype='object')
for expert_id, workload_weight in enumerate(
layer_workloads[layer]):
weights = np.zeros((expert_num,), dtype="object")
for expert_id, workload_weight in enumerate(layer_workloads[layer]):
weights[expert_id] = (expert_id, workload_weight)
# Obtain the globally balanced placement strategy for each layer
result, layer_deployment = self.original_compute_balanced_pack_redundancy(
weights, num_npus, num_redundancy_expert)
weights, num_npus, num_redundancy_expert
)
global_deployment[layer] = layer_deployment
max_heat_per_layer_after[layer] = max(
result, key=lambda x: x['total_weight'])['total_weight']
max_heat_per_layer_after[layer] = max(result, key=lambda x: x["total_weight"])["total_weight"]
new_global_deployment = self.constraint_expert_local_exchange(
current_expert_table, global_deployment)
new_global_deployment = self.constraint_expert_local_exchange(current_expert_table, global_deployment)
# Obtain the priority of each layer
layer_changed_ratio = []
for layer_idx in range(layer_num):
layer_changed_ratio.append(max_heat_per_layer_after[layer_idx] /
max_heat_per_layer_before[layer_idx])
layer_changed_ratio.append(max_heat_per_layer_after[layer_idx] / max_heat_per_layer_before[layer_idx])
per_layer_priority = np.argsort(layer_changed_ratio)
npu_heat_all_after = sum(max_heat_per_layer_after)
@@ -385,5 +350,4 @@ class DefaultEplb(EplbPolicy):
if npu_heat_all_after < 0.95 * npu_heat_all_origin:
change = 1
return change, per_layer_priority, np.array(
new_global_deployment).tolist()
return change, per_layer_priority, np.array(new_global_deployment).tolist()

19
vllm_ascend/eplb/core/policy/policy_factory.py
View File

@@ -2,29 +2,26 @@
# Todo: Once https://github.com/vllm-project/vllm/pull/24069 is merged in vllm. Remove this factory.
from .policy_abstract import DynamicConfig, EplbPolicy
from .policy_default_eplb import DefaultEplb
from .policy_swift_balancer import SwiftBalanceEplb
from .policy_flashlb import FlashLB, warm_up
from .policy_random import RandomLoadBalance
from .policy_swift_balancer import SwiftBalanceEplb
class PolicyFactory:
@staticmethod
def generate_policy(policy_type: int, config: DynamicConfig) -> EplbPolicy:
policy = {
# Constraint applying Dynamic EPLB policy V2:
# If there exists redundant expert:
# only one redundant expert can be placed in one NPU and its physical expert index must be 0
# Applying greedy d2d expert weight update composing
0:
RandomLoadBalance, # RandomLoadBalance: shuffle last physical expert on NPU 1 and 3
1:
DefaultEplb, # Dynamic EPLB policy: overall expert replacement based on current moe load
2:
SwiftBalanceEplb, # Dynamic EPLB policy V2: expert replacement with constrained number of expert shuffle
3:
FlashLB, # FlashLB EPLB policy: expert replacement based on Joint Optimization, Multi-Shot Enhancement and Incremental Adjustment
0: RandomLoadBalance, # RandomLoadBalance: shuffle last physical expert on NPU 1 and 3
1: DefaultEplb, # Dynamic EPLB policy: overall expert replacement based on current moe load
# Dynamic EPLB policy V2: expert replacement with constrained number of expert shuffle
2: SwiftBalanceEplb,
# FlashLB EPLB policy: expert replacement based on Joint Optimization,
# Multi-Shot Enhancement and Incremental Adjustment
3: FlashLB,
}
policy_class = policy.get(policy_type, RandomLoadBalance)
policy_instance = policy_class(config)

127
vllm_ascend/eplb/core/policy/policy_flashlb.py
View File

@@ -3,7 +3,6 @@
import logging
from collections import deque
from typing import Dict
import numpy as np
import torch
@@ -45,8 +44,7 @@ def compute_piece_counts(X, P, stage_weights):
secv = unit[i, idx2]
alt = X[i, idx1] / (pieces[idx1] + 1)
delta = origin - (alt if alt > secv else secv)
deltas[idx1] += delta * stage_weights[i] if np.any(
delta) != 0 else stage_weights[i]
deltas[idx1] += delta * stage_weights[i] if np.any(delta) != 0 else stage_weights[i]
max_idx = np.argmax(deltas)
pieces[max_idx] += 1
@@ -157,9 +155,7 @@ def jsq_placement(X, pieces, M, stage_weights):
# Get elements already in current column
current_rank_elements = set(deployment[rank, :])
# Filter elements from range(N) not in current column
available = [
x for x in range(N) if x not in current_rank_elements
]
available = [x for x in range(N) if x not in current_rank_elements]
# Randomly select an available element to fill
if len(available) > 0:
rand_idx = np.random.randint(0, len(available))
@@ -187,8 +183,7 @@ def slice_values(X, pieces):
@njit
def group_based_adaptive_bloating_kernel(X, P, M, simulated_pieces,
simulated_deployment, stage_weights):
def group_based_adaptive_bloating_kernel(X, P, M, simulated_pieces, simulated_deployment, stage_weights):
n_stage, N = X.shape
num_group = P // M
@@ -207,12 +202,11 @@ def group_based_adaptive_bloating_kernel(X, P, M, simulated_pieces,
flat_deployment = simulated_deployment.reshape(-1)
simulated_load = np.zeros(M, dtype=np.float32)
for i in range(flat_deployment.shape[0]):
simulated_load[i // (flat_deployment.shape[0] //
M)] += unit_load[flat_deployment[i]]
simulated_load[i // (flat_deployment.shape[0] // M)] += unit_load[flat_deployment[i]]
slice_vals = slice_values(X_all, simulated_pieces)
sorted_slices = np.sort(slice_vals)[::-1]
simulated_slopes = (sorted_slices[:-M + 1] - sorted_slices[M - 1:]) / M
simulated_slopes = (sorted_slices[: -M + 1] - sorted_slices[M - 1 :]) / M
cumulative_slices_used = np.zeros(N, dtype=np.int32)
acc = 0
@@ -230,8 +224,7 @@ def group_based_adaptive_bloating_kernel(X, P, M, simulated_pieces,
slices_used_per_group = np.zeros(num_group, dtype=np.int32)
slices_used_per_group[0] = group_boundary_indices[0]
for i in range(1, num_group):
slices_used_per_group[
i] = group_boundary_indices[i] - group_boundary_indices[i - 1]
slices_used_per_group[i] = group_boundary_indices[i] - group_boundary_indices[i - 1]
slices_used_per_group = M - slices_used_per_group
loads = np.zeros(M, dtype=np.float32)
@@ -240,7 +233,7 @@ def group_based_adaptive_bloating_kernel(X, P, M, simulated_pieces,
current_idx = 0
for g in range(num_group):
window = X_sorted[:, current_idx:current_idx + 2 * M]
window = X_sorted[:, current_idx : current_idx + 2 * M]
low = max(0, current_idx + M - N)
high = min(num_remain_slice, M - 1)
@@ -248,8 +241,7 @@ def group_based_adaptive_bloating_kernel(X, P, M, simulated_pieces,
mid = int((high + low) // 2)
keep = M - mid
current_group = window[:, :keep]
current_pieces = compute_piece_counts(current_group, M,
stage_weights)
current_pieces = compute_piece_counts(current_group, M, stage_weights)
current_pieces = np.maximum(current_pieces, 1)
current_slice = slice_values(current_group.sum(0), current_pieces)
current_slice_sorted = np.sort(current_slice)
@@ -257,8 +249,7 @@ def group_based_adaptive_bloating_kernel(X, P, M, simulated_pieces,
current_max: np.float32 = np.max(current_loads)
current_min: np.float32 = np.min(current_loads)
current_slope = (current_max - current_min) / M
next_slope: np.float32 = np.max(simulated_slopes[current_idx +
keep:])
next_slope: np.float32 = np.max(simulated_slopes[current_idx + keep :])
if abs(current_slope) > abs(next_slope):
low = mid
@@ -327,9 +318,7 @@ def auto_fix_new_placement(old_placement, new_placement):
old_row = old_placement[rank_id]
new_row = new_placement[rank_id]
index_array = np.full((max_expert + 1, num_experts),
-1,
dtype=np.int32)
index_array = np.full((max_expert + 1, num_experts), -1, dtype=np.int32)
count_array = np.zeros(max_expert + 1, dtype=np.int32)
for idx in range(num_experts):
@@ -387,21 +376,17 @@ def auto_fix_new_placement(old_placement, new_placement):
class FlashLB(EplbPolicy):
def __init__(self, config: DynamicConfig):
super().__init__(config)
self.par_history: Dict[int, float] = {}
self.hotness_window: Dict[int, deque[float]] = {}
self.max_stage_window = (config.max_stage_window if hasattr(
config, "max_stage_window") else 1)
self.par_history: dict[int, float] = {}
self.hotness_window: dict[int, deque[float]] = {}
self.max_stage_window = config.max_stage_window if hasattr(config, "max_stage_window") else 1
self.buffer_expert_layer_num = (
config.buffer_expert_layer_num if hasattr(
config, "buffer_expert_layer_num") else 58)
self.threshold_ratio = (config.threshold_ratio if hasattr(
config, "threshold_ratio") else 0)
config.buffer_expert_layer_num if hasattr(config, "buffer_expert_layer_num") else 58
)
self.threshold_ratio = config.threshold_ratio if hasattr(config, "threshold_ratio") else 0
def compute_expert_hotness(self, num_of_expert: int,
deployment: np.ndarray, rank_load: np.ndarray):
def compute_expert_hotness(self, num_of_expert: int, deployment: np.ndarray, rank_load: np.ndarray):
hotness = np.zeros(num_of_expert, dtype=rank_load.dtype)
deployment_flat = deployment.ravel()
rank_load_flat = rank_load.ravel()
@@ -413,22 +398,14 @@ class FlashLB(EplbPolicy):
if np.any(deployment < 0):
raise ValueError("Deployment table contains negative values.")
counts = np.bincount(deployment.reshape(-1), minlength=N)
unit_hotness = np.divide(hotness,
counts,
out=np.zeros_like(hotness, dtype=float),
where=counts != 0)
unit_hotness = np.divide(hotness, counts, out=np.zeros_like(hotness, dtype=float), where=counts != 0)
stage_par = np.zeros(n_stage)
for i in range(n_stage):
stage_load = unit_hotness[i][deployment].sum(-1)
stage_par[i] = stage_load.max() / stage_load.mean()
return stage_par.mean()
def group_based_adaptive_bloating(self,
X,
P,
M,
stage_weights=None,
recorsive=False):
def group_based_adaptive_bloating(self, X, P, M, stage_weights=None, recorsive=False):
n_stage, N = X.shape
if stage_weights is None:
stage_weights = np.ones(n_stage, dtype=np.float32)
@@ -437,15 +414,10 @@ class FlashLB(EplbPolicy):
(
simulated_deployment,
simulated_pieces,
) = self.group_based_adaptive_bloating(X,
P,
M,
stage_weights,
recorsive=False)
) = self.group_based_adaptive_bloating(X, P, M, stage_weights, recorsive=False)
else:
simulated_pieces = compute_piece_counts(X, P, stage_weights)
simulated_deployment = jsq_placement(X, simulated_pieces, M,
stage_weights)
simulated_deployment = jsq_placement(X, simulated_pieces, M, stage_weights)
pieces = group_based_adaptive_bloating_kernel(
X.astype(np.float32),
@@ -459,10 +431,7 @@ class FlashLB(EplbPolicy):
deployment = jsq_placement(X, pieces, M, stage_weights)
X_all = X.sum(0)
unit_load = np.divide(X_all,
pieces,
out=np.zeros_like(X_all, dtype=float),
where=pieces != 0)
unit_load = np.divide(X_all, pieces, out=np.zeros_like(X_all, dtype=float), where=pieces != 0)
load = unit_load[deployment].sum(-1)
sim_unit_load = X_all / simulated_pieces
@@ -510,16 +479,13 @@ class FlashLB(EplbPolicy):
def register_hotness(self, deployment, rank_load, num_layer, num_expert):
for layer in range(num_layer):
if layer not in self.hotness_window:
self.hotness_window[layer] = deque(
maxlen=self.max_stage_window)
hotness = self.compute_expert_hotness(num_expert,
deployment[layer],
rank_load[layer])
self.hotness_window[layer] = deque(maxlen=self.max_stage_window)
hotness = self.compute_expert_hotness(num_expert, deployment[layer], rank_load[layer])
self.hotness_window[layer].append(hotness)
def compress_by_avg_pooling_fast_nd(self, arr, m):
n, d = arr.shape
idx = (np.arange(n) * m // n)
idx = np.arange(n) * m // n
result = np.zeros((m, d))
counts = np.zeros((m, 1))
np.add.at(result, idx, arr)
@@ -532,8 +498,7 @@ class FlashLB(EplbPolicy):
expert_workload += 1
num_layer = expert_workload.shape[0]
num_expert = np.unique(current_expert_table[0].reshape(-1)).shape[0]
self.register_hotness(current_deployment, expert_workload, num_layer,
num_expert)
self.register_hotness(current_deployment, expert_workload, num_layer, num_expert)
new_deployment = current_deployment.copy()
@@ -544,21 +509,17 @@ class FlashLB(EplbPolicy):
for i, layer in enumerate(layers_need_update):
hotness = np.array(self.hotness_window[layer])
if hotness.shape[0] > self.max_stage_window:
hotness = self.compress_by_avg_pooling_fast_nd(
hotness, self.max_stage_window)
hotness = self.compress_by_avg_pooling_fast_nd(hotness, self.max_stage_window)
(
new_deployment[layer],
new_par[i],
current_par[i],
) = self.rebalance_layer(current_deployment[layer],
hotness,
layer_id=layer)
) = self.rebalance_layer(current_deployment[layer], hotness, layer_id=layer)
priority = new_par / current_par
priority_idx = np.argsort(priority)
priority_idx = priority_idx[priority[priority_idx] <
1][:self.buffer_expert_layer_num]
priority_idx = priority_idx[priority[priority_idx] < 1][: self.buffer_expert_layer_num]
if np.all(expert_workload == 1):
for _, layer in enumerate(layers_need_update):
@@ -572,16 +533,12 @@ class FlashLB(EplbPolicy):
layers_need_update = priority_idx
deployment = current_deployment
for layer in layers_need_update:
deployment[layer] = auto_fix_new_placement(
current_deployment[layer], new_deployment[layer])
deployment[layer] = auto_fix_new_placement(current_deployment[layer], new_deployment[layer])
return change, layers_need_update, deployment
def generate_layered_experts(num_layers=58,
layer_shape=(32, 9),
expert_min=0,
expert_max=255):
def generate_layered_experts(num_layers=58, layer_shape=(32, 9), expert_min=0, expert_max=255):
"""
Generate expert deployment matrix meeting the following conditions:
- Total of num_layers layers
@@ -598,32 +555,25 @@ def generate_layered_experts(num_layers=58,
"""
# 1. Basic parameter calculation
expert_num = expert_max - expert_min + 1 # Total number of experts: 256 (0~255)
layer_total = layer_shape[0] * layer_shape[
1] # Total elements in a single layer: 32*9=288
layer_total = layer_shape[0] * layer_shape[1] # Total elements in a single layer: 32*9=288
extra_slots = layer_total - expert_num # Number of random positions to fill per layer: 288-256=32
# 2. Verify feasibility (total elements must be ≥ number of experts to cover all experts)
assert layer_total >= expert_num, (
f"Number of elements in a single layer {layer_total} < number of experts {expert_num}, "
"cannot cover all experts")
f"Number of elements in a single layer {layer_total} < number of experts {expert_num}, cannot cover all experts"
)
# 3. Generate layers one by one
layers = []
for _ in range(num_layers):
# 3.1 Generate "complete expert sequence" (ensure each expert from 0 to 255 is included)
full_experts = torch.arange(expert_min,
expert_max + 1,
dtype=torch.int64) # shape (256,)
full_experts = torch.arange(expert_min, expert_max + 1, dtype=torch.int64) # shape (256,)
# 3.2 Generate "supplementary random experts" (fill remaining 32 positions, randomly selected from 0~255)
extra_experts = torch.randint(expert_min,
expert_max + 1,
size=(extra_slots, ),
dtype=torch.int64) # shape (32,)
extra_experts = torch.randint(expert_min, expert_max + 1, size=(extra_slots,), dtype=torch.int64) # shape (32,)
# 3.3 Concatenate and shuffle (ensure random distribution of experts in each layer)
layer_flat = torch.cat([full_experts, extra_experts],
dim=0) # shape (288,)
layer_flat = torch.cat([full_experts, extra_experts], dim=0) # shape (288,)
# Shuffle order (use randperm to generate random indices to avoid repeated shuffling issues)
shuffle_idx = torch.randperm(layer_flat.shape[0])
layer_shuffled = layer_flat[shuffle_idx]
@@ -642,7 +592,6 @@ def warm_up():
exam_config.num_die_per_host = 16
algo = FlashLB(exam_config)
# Generate target tensor
expert_tensor = generate_layered_experts(num_layers=58,
layer_shape=(32, 9))
expert_tensor = generate_layered_experts(num_layers=58, layer_shape=(32, 9))
algo.rebalance_experts(expert_tensor, torch.randint(1, 1000, (58, 32, 9)))

1
vllm_ascend/eplb/core/policy/policy_random.py
View File

@@ -9,7 +9,6 @@ random.seed(42)
class RandomLoadBalance(EplbPolicy):
def __init__(self, config: DynamicConfig):
super().__init__(config)

543
vllm_ascend/eplb/core/policy/policy_swift_balancer.py
View File

@@ -16,7 +16,6 @@ class DynamicConfig:
class EplbPolicy:
def __init__(self, config: DynamicConfig):
self.config = config
@@ -63,7 +62,6 @@ class DynamicTable:
class SwiftBalanceEplb(EplbPolicy):
def __init__(self, config: DynamicConfig):
super().__init__(config)
@@ -89,8 +87,7 @@ class SwiftBalanceEplb(EplbPolicy):
return a % b
@staticmethod
def add_redundant(current_expert_table, expert_workload,
num_original_expert):
def add_redundant(current_expert_table, expert_workload, num_original_expert):
layer_num, npu_num, experts_per_npu = expert_workload.shape
workload_new = np.zeros((layer_num, num_original_expert))
for layer_idx in range(layer_num):
@@ -99,8 +96,7 @@ class SwiftBalanceEplb(EplbPolicy):
workload_layer = expert_workload[layer_idx].copy()
for npu_idx in range(npu_num):
for expert_idx in range(experts_per_npu):
workload_dict[placement_layer[npu_idx][
expert_idx]] += workload_layer[npu_idx][expert_idx]
workload_dict[placement_layer[npu_idx][expert_idx]] += workload_layer[npu_idx][expert_idx]
for expert_idx in range(num_original_expert):
workload_new[layer_idx][expert_idx] = workload_dict[expert_idx]
return workload_new
@@ -118,9 +114,7 @@ class SwiftBalanceEplb(EplbPolicy):
max_heat_per_layer.append(np.max(npu_heats_now))
return max_heat_per_layer
def calculate_initial_imbalance(self, global_deployment,
new_layer_workloads):
def calculate_initial_imbalance(self, global_deployment, new_layer_workloads):
device_num = global_deployment.shape[1]
layer_imbalance = []
expert_num = np.zeros_like(new_layer_workloads)
@@ -136,56 +130,54 @@ class SwiftBalanceEplb(EplbPolicy):
box_workload = 0
for expert_id in box:
update_workload = self.safe_divide(
new_layer_workloads[layer_id][expert_id],
expert_num[layer_id][expert_id])
new_layer_workloads[layer_id][expert_id], expert_num[layer_id][expert_id]
)
box_workload += update_workload
total_workload += update_workload
if cur_layer_max_workload < box_workload:
cur_layer_max_workload = box_workload
cur_layer_imbalance = self.safe_divide(
cur_layer_max_workload,
(self.safe_divide(total_workload, device_num)))
cur_layer_max_workload, (self.safe_divide(total_workload, device_num))
)
layer_imbalance.append(cur_layer_imbalance)
return layer_imbalance
def compute_redundant_assignments(self, base_experts,
num_redundant_experts, num_experts):
redundant_assignments: list[list[int]] = [[]
for _ in range(num_experts)]
def compute_redundant_assignments(self, base_experts, num_redundant_experts, num_experts):
redundant_assignments: list[list[int]] = [[] for _ in range(num_experts)]
current_weights = base_experts.copy()
for i in range(num_redundant_experts):
sorted_indices = np.argsort([w for _, w in current_weights],
kind='stable')[::-1]
sorted_indices = np.argsort([w for _, w in current_weights], kind="stable")[::-1]
sorted_weights = [current_weights[i] for i in sorted_indices]
target_expert = sorted_weights[0]
expert_id, original_weight = target_expert
current_redundancy = len(redundant_assignments[expert_id])
new_avg_weight = self.safe_divide(
original_weight * (current_redundancy + 1),
(current_redundancy + 2))
new_avg_weight = self.safe_divide(original_weight * (current_redundancy + 1), (current_redundancy + 2))
redundant_assignments[expert_id].append(num_experts + i)
current_weights[sorted_indices[0]] = (expert_id, new_avg_weight)
sorted_indices = np.argsort([w for _, w in current_weights],
kind='stable')[::-1]
sorted_indices = np.argsort([w for _, w in current_weights], kind="stable")[::-1]
sorted_weights = [current_weights[i] for i in sorted_indices]
return redundant_assignments, sorted_weights
def repeat_compute_redundant_assignments(self, layer_workloads, rendun_pos,
num_experts, num_exist_expert,
device_assignments, device_counts,
expert_from_device,
com_between_devices):
current_weights = np.zeros((num_experts, ), dtype='object')
def repeat_compute_redundant_assignments(
self,
layer_workloads,
rendun_pos,
num_experts,
num_exist_expert,
device_assignments,
device_counts,
expert_from_device,
com_between_devices,
):
current_weights = np.zeros((num_experts,), dtype="object")
for expert_id, workload_weight in enumerate(layer_workloads):
current_weights[expert_id] = (expert_id, workload_weight)
@@ -195,8 +187,7 @@ class SwiftBalanceEplb(EplbPolicy):
devices_with_slots.append(device_id)
while devices_with_slots:
sorted_indices = np.argsort([w for _, w in current_weights],
kind='stable')[::-1]
sorted_indices = np.argsort([w for _, w in current_weights], kind="stable")[::-1]
sorted_weights = [current_weights[i] for i in sorted_indices]
for index, target_weight in enumerate(sorted_weights):
@@ -210,17 +201,15 @@ class SwiftBalanceEplb(EplbPolicy):
pos = rendun_pos[cur_device_id].pop()
if len(rendun_pos[cur_device_id]) == 0:
devices_with_slots = [
device_id for device_id in devices_with_slots
if device_id != cur_device_id
device_id for device_id in devices_with_slots if device_id != cur_device_id
]
device_assignments[cur_device_id][pos] = expert_id
device_counts[cur_device_id] += 1
communication_box_index = expert_from_device[expert_id]
com_between_devices[cur_device_id][
communication_box_index] = expert_id
com_between_devices[cur_device_id][communication_box_index] = expert_id
new_weight = self.safe_divide(
(original_weight * num_exist_expert[expert_id]),
(num_exist_expert[expert_id] + 1))
(original_weight * num_exist_expert[expert_id]), (num_exist_expert[expert_id] + 1)
)
sorted_weights[index] = (expert_id, new_weight)
num_exist_expert[expert_id] += 1
redundancy_successful = True
@@ -228,41 +217,31 @@ class SwiftBalanceEplb(EplbPolicy):
if redundancy_successful:
break
sorted_indices = np.argsort([id for id, _ in sorted_weights],
kind='stable')
sorted_indices = np.argsort([id for id, _ in sorted_weights], kind="stable")
sorted_weights = [sorted_weights[i][1] for i in sorted_indices]
return sorted_weights, device_assignments, device_counts, com_between_devices
@staticmethod
def prepare_expert_list(base_experts, redundant_assignments,
num_redundant_experts):
def prepare_expert_list(base_experts, redundant_assignments, num_redundant_experts):
redundant_expert_list = np.empty(num_redundant_experts, dtype=object)
index = 0
num_experts = len(redundant_assignments)
for expert_id in range(num_experts):
for _ in redundant_assignments[expert_id]:
redundant_expert_list[index] = (expert_id,
next(w
for eid, w in base_experts
if eid == expert_id))
redundant_expert_list[index] = (expert_id, next(w for eid, w in base_experts if eid == expert_id))
index += 1
sorted_indices = np.argsort([w for _, w in redundant_expert_list],
kind='stable')[::-1]
sorted_indices = np.argsort([w for _, w in redundant_expert_list], kind="stable")[::-1]
return [redundant_expert_list[i] for i in sorted_indices]
@staticmethod
def non_redundant_expert_information(origin_deployment, updated_weights,
rendun_pos):
def non_redundant_expert_information(origin_deployment, updated_weights, rendun_pos):
device_num = len(origin_deployment)
num_experts_per_device = origin_deployment.shape[1]
device_assignments = [[-1 for _ in range(num_experts_per_device)]
for _ in range(device_num)]
device_weights = [[0 for _ in range(num_experts_per_device)]
for _ in range(device_num)]
device_assignments = [[-1 for _ in range(num_experts_per_device)] for _ in range(device_num)]
device_weights = [[0 for _ in range(num_experts_per_device)] for _ in range(device_num)]
device_loads = [0] * device_num
device_counts = [0] * device_num
@@ -272,8 +251,8 @@ class SwiftBalanceEplb(EplbPolicy):
continue
device_assignments[device_id][index] = expert_id
cur_weight = next(
weight for expert_id_of_weight, weight in updated_weights
if expert_id_of_weight == expert_id)
weight for expert_id_of_weight, weight in updated_weights if expert_id_of_weight == expert_id
)
device_weights[device_id][index] = cur_weight
device_loads[device_id] += cur_weight
device_counts[device_id] += 1
@@ -292,19 +271,24 @@ class SwiftBalanceEplb(EplbPolicy):
if num_all_experts[expert_id] == 0:
cur_layer_workload.append(-1)
else:
cur_layer_workload.append(
self.safe_divide(weight, num_all_experts[expert_id]))
cur_layer_workload.append(self.safe_divide(weight, num_all_experts[expert_id]))
return cur_layer_workload, num_all_experts
def distribute_redun_experts(self, layer_workloads, device_assignments,
device_weights, device_loads, device_counts,
redundant_expert_list, expert_from_device,
num_experts, rendun_pos):
def distribute_redun_experts(
self,
layer_workloads,
device_assignments,
device_weights,
device_loads,
device_counts,
redundant_expert_list,
expert_from_device,
num_experts,
rendun_pos,
):
num_devices = len(device_assignments)
com_between_devices: list[dict[int,
int]] = [{} for _ in range(num_devices)]
com_between_devices: list[dict[int, int]] = [{} for _ in range(num_devices)]
for expert_id, weight in redundant_expert_list:
candidate = -1
@@ -313,8 +297,7 @@ class SwiftBalanceEplb(EplbPolicy):
continue
if expert_id in device_assignments[dev_id]:
continue
if candidate == -1 or device_loads[dev_id] < device_loads[
candidate]:
if candidate == -1 or device_loads[dev_id] < device_loads[candidate]:
candidate = dev_id
if candidate != -1:
pos = rendun_pos[candidate].pop()
@@ -324,31 +307,42 @@ class SwiftBalanceEplb(EplbPolicy):
device_counts[candidate] += 1
communication_box_index = expert_from_device[expert_id]
com_between_devices[candidate][
communication_box_index] = expert_id
com_between_devices[candidate][communication_box_index] = expert_id
if any(sublist for sublist in rendun_pos):
cur_layer_workload, num_exist_expert = self.recomputing_initial_weight(
layer_workloads, device_assignments)
cur_layer_workload, num_exist_expert = self.recomputing_initial_weight(layer_workloads, device_assignments)
update_workload, device_assignments, device_counts, com_between_devices = self.repeat_compute_redundant_assignments(
cur_layer_workload, rendun_pos, num_experts, num_exist_expert,
device_assignments, device_loads, expert_from_device,
com_between_devices)
update_workload, device_assignments, device_counts, com_between_devices = (
self.repeat_compute_redundant_assignments(
cur_layer_workload,
rendun_pos,
num_experts,
num_exist_expert,
device_assignments,
device_loads,
expert_from_device,
com_between_devices,
)
)
device_loads = [0] * len(device_counts)
for device_id, device in enumerate(device_assignments):
for index, expert_id in enumerate(device):
device_weights[device_id][index] = update_workload[
expert_id]
device_weights[device_id][index] = update_workload[expert_id]
device_loads[device_id] += update_workload[expert_id]
return device_assignments, device_weights, device_loads, device_counts, com_between_devices
def redundancy_again(self, layer_workloads, origin_weights,
origin_deployment, expert_from_device, num_node,
is_node_redundant, rendun_pos):
def redundancy_again(
self,
layer_workloads,
origin_weights,
origin_deployment,
expert_from_device,
num_node,
is_node_redundant,
rendun_pos,
):
num_experts = len(origin_weights)
if is_node_redundant:
num_experts = num_experts * num_node
@@ -358,25 +352,33 @@ class SwiftBalanceEplb(EplbPolicy):
num_redundant_experts += len(rank_empty_pos)
redundant_assignments, updated_weights = self.compute_redundant_assignments(
origin_weights, num_redundant_experts, num_experts)
origin_weights, num_redundant_experts, num_experts
)
redundant_expert_list = self.prepare_expert_list(
updated_weights, redundant_assignments, num_redundant_experts)
redundant_expert_list = self.prepare_expert_list(updated_weights, redundant_assignments, num_redundant_experts)
device_assignments, device_weights, device_loads, device_counts = self.non_redundant_expert_information(
origin_deployment, updated_weights, rendun_pos)
origin_deployment, updated_weights, rendun_pos
)
device_assignments, device_weights, device_loads, device_counts, com_between_devices = self.distribute_redun_experts(
layer_workloads, device_assignments, device_weights, device_loads,
device_counts, redundant_expert_list, expert_from_device,
num_experts, rendun_pos)
device_assignments, device_weights, device_loads, device_counts, com_between_devices = (
self.distribute_redun_experts(
layer_workloads,
device_assignments,
device_weights,
device_loads,
device_counts,
redundant_expert_list,
expert_from_device,
num_experts,
rendun_pos,
)
)
return device_assignments, device_weights, device_loads, device_counts, com_between_devices
@staticmethod
def generate_allocation_report(device_assignments, device_weights,
device_loads, device_counts):
def generate_allocation_report(device_assignments, device_weights, device_loads, device_counts):
report = []
max_load = 0.0
@@ -384,27 +386,27 @@ class SwiftBalanceEplb(EplbPolicy):
current_load = device_loads[dev_id]
max_load = max(max_load, current_load)
report.append({
"device_id": dev_id + 1,
"assigned_experts": device_assignments[dev_id],
"expert_weights": device_weights[dev_id],
"total_load": current_load,
"expert_count": device_counts[dev_id]
})
report.append(
{
"device_id": dev_id + 1,
"assigned_experts": device_assignments[dev_id],
"expert_weights": device_weights[dev_id],
"total_load": current_load,
"expert_count": device_counts[dev_id],
}
)
return report, max_load
@staticmethod
def exchange_expert(cur_exchange_index, next_exchange_index, cur_device_id,
next_device_id, cur_layer_result, com_between_devices):
def exchange_expert(
cur_exchange_index, next_exchange_index, cur_device_id, next_device_id, cur_layer_result, com_between_devices
):
cur_device_deployment = cur_layer_result[cur_device_id]["assigned_experts"]
next_device_deployment = cur_layer_result[next_device_id]["assigned_experts"]
cur_device_deployment = cur_layer_result[cur_device_id][
'assigned_experts']
next_device_deployment = cur_layer_result[next_device_id][
'assigned_experts']
cur_device_weight = cur_layer_result[cur_device_id]['expert_weights']
next_device_weight = cur_layer_result[next_device_id]['expert_weights']
cur_device_weight = cur_layer_result[cur_device_id]["expert_weights"]
next_device_weight = cur_layer_result[next_device_id]["expert_weights"]
cur_expert_id = cur_device_deployment[cur_exchange_index]
next_expert_id = next_device_deployment[next_exchange_index]
@@ -416,29 +418,25 @@ class SwiftBalanceEplb(EplbPolicy):
cur_device_weight[cur_exchange_index] = next_expert_weight
next_device_weight[next_exchange_index] = cur_expert_weight
cur_layer_result[cur_device_id][
'total_load'] += next_expert_weight - cur_expert_weight
cur_layer_result[next_device_id][
'total_load'] += cur_expert_weight - next_expert_weight
cur_layer_result[cur_device_id]["total_load"] += next_expert_weight - cur_expert_weight
cur_layer_result[next_device_id]["total_load"] += cur_expert_weight - next_expert_weight
com_between_devices[cur_device_id][next_device_id] = next_expert_id
com_between_devices[next_device_id][cur_device_id] = cur_expert_id
def redundant_expert_deployment(self, layer_workloads, original_deployment,
expert_from_device, node_num,
is_node_redundant, rendun_pos):
def redundant_expert_deployment(
self, layer_workloads, original_deployment, expert_from_device, node_num, is_node_redundant, rendun_pos
):
device_num, per_device_expert_num = original_deployment.shape
route_expert_num = layer_workloads.shape[0]
per_node_device_num = self.safe_exact_divide(device_num, node_num)
per_node_route_expert_num = per_node_device_num * (
per_device_expert_num - 1)
per_node_route_expert_num = per_node_device_num * (per_device_expert_num - 1)
weights = np.zeros((route_expert_num, ), dtype='object')
weights = np.zeros((route_expert_num,), dtype="object")
for expert_id, workload_weight in enumerate(layer_workloads):
weights[expert_id] = (expert_id, workload_weight)
if is_node_redundant:
device_assignments = []
device_weights = []
device_loads = []
@@ -446,23 +444,30 @@ class SwiftBalanceEplb(EplbPolicy):
com_between_devices = []
for node_id in range(node_num):
cur_node_weights = weights[node_id *
per_node_route_expert_num:(node_id +
1) *
per_node_route_expert_num]
cur_node_weights = weights[
node_id * per_node_route_expert_num : (node_id + 1) * per_node_route_expert_num
]
cur_original_deployment = original_deployment[
node_id * per_node_device_num:(node_id + 1) *
per_node_device_num]
node_id * per_node_device_num : (node_id + 1) * per_node_device_num
]
cur_node_rendun_pos = rendun_pos[node_id *
per_node_device_num:(node_id +
1) *
per_node_device_num]
cur_node_rendun_pos = rendun_pos[node_id * per_node_device_num : (node_id + 1) * per_node_device_num]
cur_device_assignments, cur_device_weights, cur_device_loads, cur_device_counts, cur_com_between_devices = self.redundancy_again(
layer_workloads, cur_node_weights, cur_original_deployment,
expert_from_device, node_num, is_node_redundant,
cur_node_rendun_pos)
(
cur_device_assignments,
cur_device_weights,
cur_device_loads,
cur_device_counts,
cur_com_between_devices,
) = self.redundancy_again(
layer_workloads,
cur_node_weights,
cur_original_deployment,
expert_from_device,
node_num,
is_node_redundant,
cur_node_rendun_pos,
)
device_assignments += cur_device_assignments
device_weights += cur_device_weights
device_loads += cur_device_loads
@@ -470,28 +475,41 @@ class SwiftBalanceEplb(EplbPolicy):
com_between_devices += cur_com_between_devices
else:
device_assignments, device_weights, device_loads, device_counts, com_between_devices = self.redundancy_again(
layer_workloads, weights, original_deployment,
expert_from_device, node_num, is_node_redundant, rendun_pos)
device_assignments, device_weights, device_loads, device_counts, com_between_devices = (
self.redundancy_again(
layer_workloads,
weights,
original_deployment,
expert_from_device,
node_num,
is_node_redundant,
rendun_pos,
)
)
report, max_load = self.generate_allocation_report(
device_assignments, device_weights, device_loads, device_counts)
device_assignments, device_weights, device_loads, device_counts
)
return report, max_load, com_between_devices
@staticmethod
def two_device_exchange_experts(cur_device_result, exchange_device_result,
cur_exchanged_expert_id,
next_exchanged_expert_id, ave_workload,
increment, num_redundancy_expert):
def two_device_exchange_experts(
cur_device_result,
exchange_device_result,
cur_exchanged_expert_id,
next_exchanged_expert_id,
ave_workload,
increment,
num_redundancy_expert,
):
cur_device_weight = cur_device_result["expert_weights"]
next_device_weight = exchange_device_result["expert_weights"]
cur_device_weight = cur_device_result['expert_weights']
next_device_weight = exchange_device_result['expert_weights']
cur_device_expert_id = cur_device_result["assigned_experts"]
next_device_expert_id = exchange_device_result["assigned_experts"]
cur_device_expert_id = cur_device_result['assigned_experts']
next_device_expert_id = exchange_device_result['assigned_experts']
cur_device_total_weight = cur_device_result['total_load']
next_device_total_weight = exchange_device_result['total_load']
cur_device_total_weight = cur_device_result["total_load"]
next_device_total_weight = exchange_device_result["total_load"]
max_weight = max(cur_device_total_weight, next_device_total_weight)
cur_exchange_index = -1
@@ -500,50 +518,47 @@ class SwiftBalanceEplb(EplbPolicy):
for index, weight in enumerate(cur_device_weight):
for next_index, next_weight in enumerate(next_device_weight):
change_flag = True
if (cur_device_expert_id[index] in next_device_expert_id
or next_device_expert_id[next_index]
in cur_device_expert_id):
if (
cur_device_expert_id[index] in next_device_expert_id
or next_device_expert_id[next_index] in cur_device_expert_id
):
change_flag = False
if (cur_device_expert_id[index] not in cur_exchanged_expert_id
) and (next_device_expert_id[next_index]
not in next_exchanged_expert_id) and change_flag:
if (
(cur_device_expert_id[index] not in cur_exchanged_expert_id)
and (next_device_expert_id[next_index] not in next_exchanged_expert_id)
and change_flag
):
cur_total_weight_after_exchange = cur_device_total_weight - weight + next_weight
next_total_weight_after_exchange = next_device_total_weight - next_weight + weight
exchange_max_weight = max(
cur_total_weight_after_exchange,
next_total_weight_after_exchange)
if exchange_max_weight < max_weight and (
max_weight -
exchange_max_weight) >= (ave_workload * increment):
exchange_max_weight = max(cur_total_weight_after_exchange, next_total_weight_after_exchange)
if exchange_max_weight < max_weight and (max_weight - exchange_max_weight) >= (
ave_workload * increment
):
max_weight = exchange_max_weight
cur_exchange_index = index
next_exchange_index = next_index
return cur_exchange_index, next_exchange_index
def expert_exchange_between_devices(self,
ave_workload,
increment,
cur_layer_result,
com_between_devices,
num_redundancy_expert,
node_idx=0,
per_node_device_num=0,
is_node_redundant=False):
def expert_exchange_between_devices(
self,
ave_workload,
increment,
cur_layer_result,
com_between_devices,
num_redundancy_expert,
node_idx=0,
per_node_device_num=0,
is_node_redundant=False,
):
if is_node_redundant:
cur_devices_result = cur_layer_result[node_idx *
per_node_device_num:
(node_idx + 1) *
per_node_device_num]
cur_devices_result = cur_layer_result[node_idx * per_node_device_num : (node_idx + 1) * per_node_device_num]
else:
cur_devices_result = cur_layer_result
devices_total_weight = []
for device in cur_devices_result:
devices_total_weight.append(
(device['total_load'], device['device_id'] - 1))
devices_total_weight.append((device["total_load"], device["device_id"] - 1))
exchange_frequency = 100
while exchange_frequency > 0:
@@ -553,64 +568,81 @@ class SwiftBalanceEplb(EplbPolicy):
exchange = False
for index in range(0, len(devices_total_weight) - 1):
min_weight_device_id = devices_total_weight[index][1]
if min_weight_device_id not in com_between_devices[
max_weight_device_id]:
cur_exchanged_expert_id = list(
com_between_devices[max_weight_device_id].values())
next_exchanged_expert_id = list(
com_between_devices[min_weight_device_id].values())
if min_weight_device_id not in com_between_devices[max_weight_device_id]:
cur_exchanged_expert_id = list(com_between_devices[max_weight_device_id].values())
next_exchanged_expert_id = list(com_between_devices[min_weight_device_id].values())
cur_exchange_index, next_exchange_index = self.two_device_exchange_experts(
cur_layer_result[max_weight_device_id],
cur_layer_result[min_weight_device_id],
cur_exchanged_expert_id, next_exchanged_expert_id,
ave_workload, increment, num_redundancy_expert)
cur_exchanged_expert_id,
next_exchanged_expert_id,
ave_workload,
increment,
num_redundancy_expert,
)
if cur_exchange_index != -1:
self.exchange_expert(cur_exchange_index,
next_exchange_index,
max_weight_device_id,
min_weight_device_id,
cur_layer_result,
com_between_devices)
self.exchange_expert(
cur_exchange_index,
next_exchange_index,
max_weight_device_id,
min_weight_device_id,
cur_layer_result,
com_between_devices,
)
devices_total_weight[-1] = (
cur_layer_result[max_weight_device_id]
['total_load'], max_weight_device_id)
cur_layer_result[max_weight_device_id]["total_load"],
max_weight_device_id,
)
devices_total_weight[index] = (
cur_layer_result[min_weight_device_id]
['total_load'], min_weight_device_id)
cur_layer_result[min_weight_device_id]["total_load"],
min_weight_device_id,
)
exchange = True
break
if not exchange:
break
def exchange_experts(self, layer_result, layer_com_between_devices,
num_nodes, device_num, is_node_redundant,
ave_workload, increment, num_redundancy_expert,
org_deployment):
def exchange_experts(
self,
layer_result,
layer_com_between_devices,
num_nodes,
device_num,
is_node_redundant,
ave_workload,
increment,
num_redundancy_expert,
org_deployment,
):
global_deployment = []
if is_node_redundant:
per_node_device_num = self.safe_exact_divide(device_num, num_nodes)
for node_idx in range(num_nodes):
self.expert_exchange_between_devices(
ave_workload, increment, layer_result,
layer_com_between_devices, num_redundancy_expert, node_idx,
per_node_device_num, is_node_redundant)
ave_workload,
increment,
layer_result,
layer_com_between_devices,
num_redundancy_expert,
node_idx,
per_node_device_num,
is_node_redundant,
)
else:
self.expert_exchange_between_devices(ave_workload, increment,
layer_result,
layer_com_between_devices,
num_redundancy_expert)
self.expert_exchange_between_devices(
ave_workload, increment, layer_result, layer_com_between_devices, num_redundancy_expert
)
max_workload = 0
for box in layer_result:
global_deployment.append(box['assigned_experts'])
if max_workload < box['total_load']:
max_workload = box['total_load']
global_deployment.append(box["assigned_experts"])
if max_workload < box["total_load"]:
max_workload = box["total_load"]
global_deployment = np.array(global_deployment)
@@ -626,16 +658,11 @@ class SwiftBalanceEplb(EplbPolicy):
return count
@staticmethod
def constraint_expert_local_exchange(current_expert_table,
global_deployment):
def constraint_expert_local_exchange(current_expert_table, global_deployment):
for layer_id in range(len(global_deployment)):
for card_id in range(len(global_deployment[layer_id])):
current_list = [
int(x) for x in current_expert_table[layer_id][card_id]
]
new_list = [
int(x) for x in global_deployment[layer_id][card_id]
]
current_list = [int(x) for x in current_expert_table[layer_id][card_id]]
new_list = [int(x) for x in global_deployment[layer_id][card_id]]
num = len(new_list)
new_index = [-1] * num
@@ -645,8 +672,7 @@ class SwiftBalanceEplb(EplbPolicy):
for i in range(num):
flag = True
for j in range(num):
if new_list[i] == current_list[j] and new_index[
j] == -1:
if new_list[i] == current_list[j] and new_index[j] == -1:
new_index[j] = 0
new_result[j] = current_list[j]
flag = False
@@ -664,25 +690,17 @@ class SwiftBalanceEplb(EplbPolicy):
return global_deployment
def rebalance_experts(self,
current_expert_table,
expert_workload,
is_node_redundant=False,
increment=0.01):
def rebalance_experts(self, current_expert_table, expert_workload, is_node_redundant=False, increment=0.01):
info = DynamicTable()
info.workload_table = expert_workload.numpy()
info.placement_table = current_expert_table.numpy()
assert info.workload_table is not None
layer_num, num_npus, experts_per_npu = info.workload_table.shape
expert_ids, counts = np.unique(info.placement_table[0],
return_counts=True)
expert_ids, counts = np.unique(info.placement_table[0], return_counts=True)
num_redundancy_expert = self.get_redundant_num(num_npus, counts)
num_original_expert = len(expert_ids)
layer_workloads = self.add_redundant(info.placement_table,
info.workload_table,
num_original_expert)
max_heat_per_layer_before = self.calculate_max_heat_per_layer(
info.workload_table, layer_num)
layer_workloads = self.add_redundant(info.placement_table, info.workload_table, num_original_expert)
max_heat_per_layer_before = self.calculate_max_heat_per_layer(info.workload_table, layer_num)
npu_heat_all_origin = sum(max_heat_per_layer_before)
num_node = self.safe_exact_divide(num_npus, 8)
@@ -700,14 +718,13 @@ class SwiftBalanceEplb(EplbPolicy):
if num_npus < num_redundancy_expert:
raise ValueError(
f"The number of NPUs ({num_npus}) must be greater than or equal to the number of redundant experts ({num_redundancy_expert})"
"The number of NPUs "
f"({num_npus}) must be greater than or equal to the number of redundant experts "
f"({num_redundancy_expert})"
)
global_deployment: list[list[list[int]]] = [[[]
for _ in range(num_npus)]
for _ in range(layer_num)]
layer_initial_imbalance = self.calculate_initial_imbalance(
info.placement_table, layer_workloads)
global_deployment: list[list[list[int]]] = [[[] for _ in range(num_npus)] for _ in range(layer_num)]
layer_initial_imbalance = self.calculate_initial_imbalance(info.placement_table, layer_workloads)
max_heat_per_layer_after = np.zeros([layer_num])
sum_num = 0
for layer in range(layer_num):
@@ -715,8 +732,7 @@ class SwiftBalanceEplb(EplbPolicy):
global_deployment[layer] = info.placement_table[layer]
continue
ave_workload = self.safe_divide(np.sum(layer_workloads[layer]),
num_npus)
ave_workload = self.safe_divide(np.sum(layer_workloads[layer]), num_npus)
rendun_pos: list[list[int]] = [[] for _ in range(num_npus)]
existing_experts = set()
@@ -729,30 +745,37 @@ class SwiftBalanceEplb(EplbPolicy):
rendun_pos[device_id].append(index)
result, max_workload, com_between_devices = self.redundant_expert_deployment(
layer_workloads[layer], info.placement_table[layer],
expert_from_device[layer], num_node, is_node_redundant,
rendun_pos)
layer_workloads[layer],
info.placement_table[layer],
expert_from_device[layer],
num_node,
is_node_redundant,
rendun_pos,
)
global_deployment[layer], new_max_workload = self.exchange_experts(
result, com_between_devices, num_node, num_npus,
is_node_redundant, ave_workload, increment,
num_redundancy_expert, info.placement_table[layer])
result,
com_between_devices,
num_node,
num_npus,
is_node_redundant,
ave_workload,
increment,
num_redundancy_expert,
info.placement_table[layer],
)
for device_id in range(num_npus):
com_between_devices[device_id] = {
key: value
for key, value in com_between_devices[device_id].items()
}
com_between_devices[device_id] = {key: value for key, value in com_between_devices[device_id].items()}
sum_num += self.count_elements(com_between_devices[device_id])
max_heat_per_layer_after[layer] = max(
result, key=lambda x: x['total_load'])['total_load']
max_heat_per_layer_after[layer] = max(result, key=lambda x: x["total_load"])["total_load"]
layer_changed_ratio = []
for layer_idx in range(layer_num):
layer_changed_ratio.append(
self.safe_divide(max_heat_per_layer_after[layer_idx],
max_heat_per_layer_before[layer_idx]))
self.safe_divide(max_heat_per_layer_after[layer_idx], max_heat_per_layer_before[layer_idx])
)
per_layer_priority = np.argsort(layer_changed_ratio)
npu_heat_all_after = sum(max_heat_per_layer_after)
@@ -761,8 +784,6 @@ class SwiftBalanceEplb(EplbPolicy):
if npu_heat_all_after < 0.95 * npu_heat_all_origin:
change = 1
new_global_deployment = self.constraint_expert_local_exchange(
current_expert_table, global_deployment)
new_global_deployment = self.constraint_expert_local_exchange(current_expert_table, global_deployment)
return change, per_layer_priority, np.array(
new_global_deployment).tolist()
return change, per_layer_priority, np.array(new_global_deployment).tolist()