lint

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

Author: jenchen13

modelopt/torch/export/plugins/mcore_nemotron.py

@@ -19,8 +19,8 @@
 from .mcore_custom import (
     COL_ETP,
     COL_TP,
-    ROW_ETP,
     REPLICATE,
+    ROW_ETP,
     ROW_TP,
     CustomModuleMapping,
     NameRemapping,
@@ -41,7 +41,7 @@
     "linear_fc2": NameRemapping("model.layers.{}.mlp.down_proj."),
     "final_layernorm": NameRemapping("model.norm."),
     "output_layer": NameRemapping("lm_head."),
-    }
+}
 
 
 nemotron_h_causal_lm_import: dict[str, CustomModuleMapping] = {
@@ -67,13 +67,18 @@
     "linear_fc2": NameRemapping("backbone.layers.{}.mixer.down_proj.", ROW_TP),
     # MoE
     "router": NameRemapping("model.layers.{}.mlp.gate.", REPLICATE),
-    "local_experts.linear_fc1": NameRemapping("backbone.layers.{}.mixer.experts.{}.up_proj", COL_ETP),
-    "local_experts.linear_fc2": NameRemapping("backbone.layers.{}.mixer.experts.{}.down_proj.", ROW_ETP),
-    "shared_experts.linear_fc1": NameRemapping("backbone.layers.{}.mixer.shared_experts.up_proj.", COL_TP),
+    "local_experts.linear_fc1": NameRemapping(
+        "backbone.layers.{}.mixer.experts.{}.up_proj", COL_ETP
+    ),
+    "local_experts.linear_fc2": NameRemapping(
+        "backbone.layers.{}.mixer.experts.{}.down_proj.", ROW_ETP
+    ),
+    "shared_experts.linear_fc1": NameRemapping(
+        "backbone.layers.{}.mixer.shared_experts.up_proj.", COL_TP
+    ),
     "shared_experts.linear_fc2": NameRemapping(
         "backbone.layers.{}.mixer.shared_experts.down_proj.", ROW_TP
     ),
-    
 }
 
 
@@ -103,6 +108,7 @@
     "local_experts.linear_fc1": NameRemapping("backbone.layers.{}.mixer.experts.{}.up_proj."),
     "local_experts.linear_fc2": NameRemapping("backbone.layers.{}.mixer.experts.{}.down_proj."),
     "shared_experts.linear_fc1": NameRemapping("backbone.layers.{}.mixer.shared_experts.up_proj."),
-    "shared_experts.linear_fc2": NameRemapping("backbone.layers.{}.mixer.shared_experts.down_proj."),
-
+    "shared_experts.linear_fc2": NameRemapping(
+        "backbone.layers.{}.mixer.shared_experts.down_proj."
+    ),
 }