Skip to content

Commit edb38e4

Browse files
committed
fix bug in compute moe load
1 parent 9e2cca1 commit edb38e4

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

vllm_ascend/eplb/eplb_updator.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ def forward_before(self):
131131
self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
132132

133133
def forward_end(self, dummy_run=False):
134-
self.adaptor.get_rank_expert_workload(self.num_moe_layers)
134+
self.adaptor.get_rank_expert_workload(self.num_moe_layers,dummy_run)
135135
if not self.update_in_flight and self.get_update_iteration():
136136
moe_load = self.compute_and_set_moe_load(dummy_run)
137137
self.wakeup_eplb_worker()
@@ -144,8 +144,8 @@ def forward_end(self, dummy_run=False):
144144

145145
self.eplb_loader.update_expert_map_and_weight(self.reqs, self.redundant_enable)
146146

147-
def compute_and_set_moe_load(self):
148-
local_load = self.adaptor.get_rank_expert_workload(self.num_moe_layers)
147+
def compute_and_set_moe_load(self,dummy_run=False):
148+
local_load = self.adaptor.get_rank_expert_workload(self.num_moe_layers,dummy_run)
149149
self._gather_buffer = None
150150
if dist.is_initialized():
151151
self.world_size = dist.get_world_size()

0 commit comments

Comments
 (0)