0.23.2 release - fix export for nemo models

Author: kevalmorabia97

modelopt/torch/export/layer_utils.py

@@ -957,7 +957,7 @@ def build_stacked_experts(
     """Builds the experts_weight_1 and experts_weight_2 configs for the experts."""
     # Resmooth all experts
     preprocess_linear_fusion(
-        [expert_getter(experts, i, linear_names[2]) for i in range(num_experts)],
+        [expert_getter(experts, i, linear_names[1]) for i in range(num_experts)],
         resmooth_only=True,
     )
 

modelopt/torch/export/tensorrt_llm_utils.py

@@ -415,7 +415,9 @@ def convert_to_tensorrt_llm_config(
     layernorm_type_map = {i.name: i.value for i in LayerNormType}
     layernorm_position_map = {i.name: i.value for i in LayerNormPositionType}
 
-    if decoder_type == "mpt":
+    if decoder_type in ["gpt", "gemma", "llama"]:
+        pass
+    elif decoder_type == "mpt":
         config.update(
             {
                 "clip_qkv": first_attention_config.clip_qkv,
@@ -506,7 +508,7 @@ def convert_to_tensorrt_llm_config(
         config["vision_output_dim"] = vision_output_dim if vision_output_dim != 0 else 7680
     else:
         raise NotImplementedError(
-            f"Cannot export tensorrt_llm checkpoint for model {config_architecture}. "
+            f"Cannot export tensorrt_llm checkpoint for model {decoder_type}: {config_architecture}. "
             "It's not supported by TensorRT-LLM yet."
         )