Llama4 Export: Remove outdated MLP weight transform (NVIDIA-NeMo#14297)

* Update HFLlamaExporter to remove outdated MLP weight transform for Llama4 model

Signed-off-by: Ao Tang <[email protected]>

* Apply isort and black reformatting

Signed-off-by: suiyoubi <[email protected]>

---------

Signed-off-by: Ao Tang <[email protected]>
Signed-off-by: suiyoubi <[email protected]>
Co-authored-by: suiyoubi <[email protected]>
Signed-off-by: Guyue Huang <[email protected]>

Author: suiyoubi

nemo/collections/llm/gpt/model/llama.py

@@ -903,6 +903,13 @@ def convert_state(self, source, target, source_config=None):
                     "decoder.layers.*.mlp.experts.linear_fc1.weight": "model.layers.*.feed_forward.experts.gate_up_proj",
                 }
             )
+
+            # Remove the transform with source_key "decoder.layers.*.mlp.linear_fc1.weight" from transforms
+            # Llama4's HF model has a different mapping for the MLP weights (map to feed_forward instead of mlp)
+            transforms = [
+                t for t in transforms if getattr(t, "source_key", None) != "decoder.layers.*.mlp.linear_fc1.weight"
+            ]
+
             transforms.extend(
                 [
                     io.state_transform(