Skip to content

Commit c2014a5

Browse files
committed
fix mixer name
Signed-off-by: jenchen13 <[email protected]>
1 parent fa46075 commit c2014a5

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

modelopt/torch/export/plugins/mcore_nemotron.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@
6767
"linear_fc2": NameRemapping("backbone.layers.{}.mixer.down_proj.", ROW_TP),
6868
# MoE
6969
"router": NameRemapping(
70-
"backbone.layers.{}.mlp.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
70+
"backbone.layers.{}.mixer.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
7171
),
7272

7373
"local_experts.linear_fc1": NameRemapping(
@@ -108,7 +108,7 @@
108108
"linear_fc2": NameRemapping("backbone.layers.{}.mixer.down_proj."),
109109
# MoE
110110
"router": NameRemapping(
111-
"backbone.layers.{}.mlp.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
111+
"backbone.layers.{}.mixer.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
112112
),
113113
"local_experts.linear_fc1": NameRemapping("backbone.layers.{}.mixer.experts.{}.up_proj."),
114114
"local_experts.linear_fc2": NameRemapping("backbone.layers.{}.mixer.experts.{}.down_proj."),

0 commit comments

Comments
 (0)