[BugFix] fix real_bsz in ep (#3366)

* Your commit message here

* fix ep

* delete cuda_graph
This commit is contained in:
lizexu123
2025-08-14 17:31:19 +08:00
committed by GitHub
parent 0ea8712018
commit 7b596d0877
3 changed files with 6 additions and 2 deletions

View File

@@ -450,7 +450,7 @@ class Ernie4_5_MoeForCausalLM(ModelForCasualLM):
self.fd_config.model_config.moe_layer_start_index,
self.fd_config.model_config.num_hidden_layers,
):
self.ernie.layers[i].mlp.expert(fake_hidden_states)
self.ernie.layers[i].mlp.experts(fake_hidden_states, self.ernie.layers[i].mlp.gate)
def forward(
self,