Upload NemotronFlashForCausalLM

model-00001-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6cea1b75a41cb4881552e82a27b03b380185cd2afc5604623ec7a6f9e7b71eb5
-size 4987938992
+oid sha256:816ac0d20ac6856b93bbcd72e700d87175cb42685508cd43362ab36e3ee1db54
+size 4987920224

model.safetensors.index.json

@@ -1,7 +1,7 @@
 {
   "metadata": {
-    "total_parameters": 2750007840,
-    "total_size": 5500015680
+    "total_parameters": 2749998624,
+    "total_size": 5499997248
   },
   "weight_map": {
     "model.embed_tokens.weight": "model-00001-of-00002.safetensors",
@@ -34,7 +34,6 @@
     "model.layers.11.ffn.up_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.11.pre_ffn_layernorm.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.input_layernorm.weight": "model-00001-of-00002.safetensors",
-    "model.layers.12.pre_ffn_layernorm.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.12.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
@@ -93,7 +92,6 @@
     "model.layers.2.mamba.norm.weight": "model-00001-of-00002.safetensors",
     "model.layers.2.mamba.out_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.input_layernorm.weight": "model-00001-of-00002.safetensors",
-    "model.layers.20.pre_ffn_layernorm.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.20.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",
@@ -201,7 +199,6 @@
     "model.layers.35.ffn.up_proj.weight": "model-00002-of-00002.safetensors",
     "model.layers.35.pre_ffn_layernorm.weight": "model-00002-of-00002.safetensors",
     "model.layers.4.input_layernorm.weight": "model-00001-of-00002.safetensors",
-    "model.layers.4.pre_ffn_layernorm.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00002.safetensors",
     "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00002.safetensors",

modeling_nemotron_flash.py

@@ -918,11 +918,12 @@ class NemotronFlashAttentionDecoderLayer(nn.Module):
 
         if self.config.intermediate_size > 0:
             self.ffn = NemotronFlashMLP(config, layer_idx=layer_idx)
+            self.pre_ffn_layernorm = NemotronFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
             self.ffn = None
+            self.pre_ffn_layernorm = None
 
         self.input_layernorm = NemotronFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.pre_ffn_layernorm = NemotronFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(
             self,
@@ -1037,13 +1038,12 @@ class NemotronFlashMambaDecoderLayer(nn.Module):
         self.intermediate_size = config.intermediate_size
         if self.intermediate_size > 0:
             self.ffn = NemotronFlashMLP(config, layer_idx=layer_idx)
-
-        self.input_layernorm = NemotronFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        
-        if self.intermediate_size > 0:
             self.pre_ffn_layernorm = NemotronFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         else:
+            self.ffn = None
             self.pre_ffn_layernorm = None
+
+        self.input_layernorm = NemotronFlashRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
                 
 
     def forward(