fix naming

Signed-off-by: Jennifer Chen <[email protected]>

Author: jenchen13

modelopt/torch/export/plugins/mcore_nemotron.py

@@ -67,11 +67,10 @@
     "linear_fc2": NameRemapping("backbone.layers.{}.mixer.down_proj.", ROW_TP),
     # MoE
     "router": NameRemapping(
-        "backbone.layers.{}.mlp.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
+        "backbone.layers.{}.mixer.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
     ),
-
     "local_experts.linear_fc1": NameRemapping(
-        "backbone.layers.{}.mixer.experts.{}.up_proj", COL_ETP
+        "backbone.layers.{}.mixer.experts.{}.up_proj.", COL_ETP
     ),
     "local_experts.linear_fc2": NameRemapping(
         "backbone.layers.{}.mixer.experts.{}.down_proj.", ROW_ETP
@@ -108,7 +107,7 @@
     "linear_fc2": NameRemapping("backbone.layers.{}.mixer.down_proj."),
     # MoE
     "router": NameRemapping(
-        "backbone.layers.{}.mlp.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
+        "backbone.layers.{}.mixer.gate.", {"mapping": {"expert_bias": "e_score_correction_bias"}}
     ),
     "local_experts.linear_fc1": NameRemapping("backbone.layers.{}.mixer.experts.{}.up_proj."),
     "local_experts.linear_fc2": NameRemapping("backbone.layers.{}.mixer.experts.{}.down_proj."),