debuging falcon new arc model on TPUs

erfanzar · May 26, 2024 · 61e11ee · 61e11ee
1 parent f4d95af
commit 61e11ee
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/src/python/easydel/modules/falcon/modelling_falcon_flax.py b/src/python/easydel/modules/falcon/modelling_falcon_flax.py
@@ -342,15 +342,13 @@ def setup(self) -> None:
             self.config.num_ln_in_parallel_attn = 2
         config = self.config
 
-        if not config.parallel_attn:
-            self.input_layernorm = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
-            self.post_attention_layernorm = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
+        if config.new_decoder_architecture and config.num_ln_in_parallel_attn == 2:
+            self.ln_attn = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
+            self.ln_mlp = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
         else:
-            if config.num_ln_in_parallel_attn == 2:
-                self.ln_attn = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
-                self.ln_mlp = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
-            else:
-                self.input_layernorm = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
+            self.input_layernorm = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
+            if not config.parallel_attn:
+                self.post_attention_layernorm = nn.LayerNorm(epsilon=config.layer_norm_epsilon, dtype=self.dtype)
         attn_block = FlaxFalconAttention
         mlp_block = FlaxFalconMlp
         if self.config.gradient_checkpointing != "":