fix

mosaicml · Apr 12, 2024 · 92a9ded · 92a9ded
1 parent 5938048
commit 92a9ded
Showing 1 changed file with 14 additions and 13 deletions.
diff --git a/llmfoundry/models/layers/blocks.py b/llmfoundry/models/layers/blocks.py
@@ -74,22 +74,14 @@ def __init__(
         super().__init__()
 
         ffn_type = ffn_config['ffn_type']
-
-        self.ffn = build_ffn(
-            name=ffn_type,
-            d_model=d_model,
-            expansion_ratio=expansion_ratio,
-            device=device,
-            bias=not no_bias,
-            ffn_kwargs=ffn_config,
-        )
+        ffn_has_norm = not ffn_type in ffns_with_norm
 
         if self.fuse_norm_attn_norm:
             self.norm_attn_norm = FusedNormAttentionNorm(
                 d_model=d_model,
                 n_heads=n_heads,
                 attn_config=attn_config,
-                ffn_type=ffn_type,
+                ffn_has_norm=ffn_has_norm,
                 fc_type=fc_type,
                 resid_pdrop=resid_pdrop,
                 norm_type=norm_type,
@@ -127,13 +119,22 @@ def __init__(
                 },
             )
             self.norm_2 = None
-            if not getattr(self.ffn, '_has_norm', False):
+            if not ffn_has_norm:
                 self.norm_2 = build_norm(
                     name=norm_type.lower(),
                     normalized_shape=d_model,
                     device=device,
                 )
 
+        self.ffn = build_ffn(
+            name=ffn_type,
+            d_model=d_model,
+            expansion_ratio=expansion_ratio,
+            device=device,
+            bias=not no_bias,
+            ffn_kwargs=ffn_config,
+        )
+
         self.resid_attn_dropout = nn.Dropout(resid_pdrop)
         self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
         self.use_pad_tok_in_ffn = use_pad_tok_in_ffn
@@ -201,7 +202,7 @@ def __init__(
         d_model: int,
         n_heads: int,
         attn_config: Optional[Dict] = None,
-        ffn_type: str = 'mptmlp',
+        ffn_has_norm: bool = False,
         fc_type: str = 'torch',
         resid_pdrop: float = 0.0,
         norm_type: str = 'low_precision_layernorm',
@@ -242,7 +243,7 @@ def __init__(
         )
 
         self.norm_2 = None
-        if not ffn_type in ffns_with_norm:
+        if not ffn_has_norm:
             self.norm_2 = build_norm(
                 name=norm_type.lower(),
                 normalized_shape=d_model,