Upload NemotronFlashForCausalLM

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b15235dc83bb411387d9a15694e02d3697051e303b067336e17c302a17b6125d
-size 1930804368
+oid sha256:9581e287ebcd6a1f13c9564d9ae4d2e03d2d464032635af3b3be0de15dbb8939
+size 1930795952

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(