feat(model): allow bias (borisdayma#152)

Author: borisdayma

src/dalle_mini/__init__.py

@@ -1,3 +1,3 @@
-__version__ = "0.0.3"
+__version__ = "0.0.4"
 
 from .model import DalleBart, DalleBartProcessor

src/dalle_mini/model/configuration.py

@@ -58,14 +58,15 @@ def __init__(
         tie_word_embeddings=False,  # different modalities and sizes
         do_sample=True,
         # transformer variants
+        use_bias=False,  # use bias in attention and dense layers (except for lm_head)
         ln_type="layernorm",  # layer normalization type, "rmsnorm", "layernorm"
         ln_positions="normformer",  # layer normalization positions, "normformer", "swinv2", "cogview", "postln", "preln", "deepnet" (same as postln)
         use_head_scale=False,  # used in NormFormer
         use_cosine_attention=False,  # used in Swin v2
         tau_init=0.05,  # used only in cosine attention (Swin v2)
         use_deepnet_scaling=False,  # used in Deepnet
         use_glu=False,  # "GLU Variants Improve Transformer"
-        use_alibi=False,  # from "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation"
+        use_alibi=False,  # Not implemented yet - from "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation"
         sinkhorn_iters=1,  # used in SinkFormers
         use_final_ln_encoder=False,  # final layer normalization in encoder
         use_final_ln_decoder=False,  # final layer normalization in decoder
@@ -77,7 +78,7 @@ def __init__(
         self.normalize_text = normalize_text
 
         # transformer variants
-        self.use_head_scale = use_head_scale  # per Normformer
+        self.use_bias = use_bias
         assert ln_type in [
             "rmsnorm",
             "layernorm",
@@ -92,6 +93,7 @@ def __init__(
             "postln",
             "preln",
         ], "ln_positions must be 'normformer', 'swinv2', 'cogview', 'postln', 'preln'"
+        self.use_head_scale = use_head_scale
         assert use_alibi is False, "use_alibi is not supported yet"
         self.ln_positions = ln_positions
         self.use_cosine_attention = use_cosine_attention

src/dalle_mini/model/modeling.py

@@ -444,7 +444,7 @@ def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:
         w = nn.Dense(
             self.ffn_dim,
             dtype=self.dtype,
-            use_bias=False,
+            use_bias=self.config.use_bias,
             kernel_init=deepnet_init(gain)
             if self.config.use_deepnet_scaling
             else jax.nn.initializers.normal(self.config.init_std),
@@ -453,7 +453,7 @@ def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:
         v = nn.Dense(
             self.ffn_dim,
             dtype=self.dtype,
-            use_bias=False,
+            use_bias=self.config.use_bias,
             kernel_init=deepnet_init(gain)
             if self.config.use_deepnet_scaling
             else jax.nn.initializers.normal(self.config.init_std),
@@ -473,7 +473,7 @@ def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:
         x = nn.Dense(
             self.embed_dim,
             dtype=self.dtype,
-            use_bias=False,
+            use_bias=self.config.use_bias,
             kernel_init=deepnet_init(gain)
             if self.config.use_deepnet_scaling
             else jax.nn.initializers.normal(self.config.init_std),
@@ -509,7 +509,7 @@ def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:
         x = nn.Dense(
             self.ffn_dim,
             dtype=self.dtype,
-            use_bias=False,
+            use_bias=self.config.use_bias,
             kernel_init=deepnet_init(gain)
             if self.config.use_deepnet_scaling
             else jax.nn.initializers.normal(self.config.init_std),
@@ -528,7 +528,7 @@ def __call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray:
         x = nn.Dense(
             self.embed_dim,
             dtype=self.dtype,
-            use_bias=False,
+            use_bias=self.config.use_bias,
             kernel_init=deepnet_init(gain)
             if self.config.use_deepnet_scaling
             else jax.nn.initializers.normal(self.config.init_std),
@@ -580,7 +580,7 @@ def __call__(
             embed_dim=embed_dim,
             num_heads=self.config.encoder_attention_heads,
             dropout=self.config.attention_dropout,
-            bias=False,
+            bias=self.config.use_bias,
             dtype=self.dtype,
             is_encoder=True,
         )(hidden_states=hidden_states, attention_mask=attention_mask)
@@ -686,7 +686,7 @@ def __call__(
             num_heads=self.config.decoder_attention_heads,
             dropout=self.config.attention_dropout,
             causal=True,
-            bias=False,
+            bias=self.config.use_bias,
             dtype=self.dtype,
             is_encoder=False,
         )(
@@ -724,7 +724,7 @@ def __call__(
                 embed_dim=embed_dim,
                 num_heads=self.config.decoder_attention_heads,
                 dropout=self.config.attention_dropout,
-                bias=False,
+                bias=self.config.use_bias,
                 dtype=self.dtype,
                 is_encoder=False,
             )(

tools/train/train.py

@@ -49,6 +49,7 @@
 from tqdm import tqdm
 from transformers import HfArgumentParser
 
+import dalle_mini
 from dalle_mini.data import Dataset
 from dalle_mini.model import (
     DalleBart,
@@ -675,6 +676,7 @@ def main():
                     "transformers": transformers.__version__,
                     "datasets": datasets.__version__,
                     "wandb": wandb.__version__,
+                    "dalle_mini": dalle_mini.__version__,
                 },
             }
         )