feat: add cogview

README.md

@@ -124,6 +124,7 @@ Sequence to sequence model based on "[BART: Denoising Sequence-to-Sequence Pre-t
 - "[Deepnet: Scaling Transformers to 1,000 Layers](https://arxiv.org/abs/2203.00555)"
 - "[NormFormer: Improved Transformer Pretraining with Extra Normalization](https://arxiv.org/abs/2110.09456)"
 - "[Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://arxiv.org/abs/2103.14030)"
+- "[CogView: Mastering Text-to-Image Generation via Transformers](https://arxiv.org/abs/2105.13290v2)
 - "[Root Mean Square Layer Normalization](https://arxiv.org/abs/1910.07467)"
 
 Main optimizer (Distributed Shampoo) from "[Scalable Second Order Optimization for Deep Learning](https://arxiv.org/abs/2002.09018)".
@@ -225,6 +226,17 @@ Main optimizer (Distributed Shampoo) from "[Scalable Second Order Optimization f
 }
 ```
 
+```text
+@misc{ding2021cogview,
+      title = {CogView: Mastering Text-to-Image Generation via Transformers},
+      author = {Ming Ding and Zhuoyi Yang and Wenyi Hong and Wendi Zheng and Chang Zhou and Da Yin and Junyang Lin and Xu Zou and Zhou Shao and Hongxia Yang and Jie Tang},
+      year = {2021},
+      eprint = {2105.13290},
+      archivePrefix = {arXiv},
+      primaryClass = {cs.CV}
+}
+```
+
 ```text
 @misc{zhang2019root,
       title = {Root Mean Square Layer Normalization},

src/dalle_mini/model/configuration.py

@@ -60,7 +60,7 @@ class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
         # transformer variants
         head_scale=False,  # used in NormFormer
         ln_type="layernorm",  # layer normalization type, "rmsnorm", "layernorm"
-        ln_positions="deepnet",  # layer normalization positions, "normformer", "swinv2", "deepnet" (same as post-ln)
+        ln_positions="deepnet",  # layer normalization positions, "normformer", "swinv2", "cogview", "deepnet" (same as post-ln)
         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
@@ -80,6 +80,7 @@ class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
         assert ln_positions in [
             "normformer",
             "swinv2",
+            "cogview",
             "deepnet",
         ], "ln_positions must be 'normformer', 'swinv2' or 'deepnet'"
         self.ln_positions = ln_positions

src/dalle_mini/model/modeling.py

@@ -373,7 +373,7 @@ class GLU(nn.Module):
             self.config
         )
 
-        if self.config.ln_positions in ["normformer"]:
+        if self.config.ln_positions in ["normformer", "cogview"]:
             x = norm(
                 self.config.ln_type, dtype=self.dtype, epsilon=1e-05, use_scale=False
             )(x)
@@ -411,7 +411,7 @@ class GLU(nn.Module):
             if self.config.use_deepnet_scaling
             else jax.nn.initializers.normal(self.config.init_std),
         )(x)
-        if self.config.ln_positions in ["swinv2"]:
+        if self.config.ln_positions in ["swinv2", "cogview"]:
             x = norm(self.config.ln_type, dtype=self.dtype, epsilon=1e-05)(x)
         x = nn.Dropout(rate=self.config.dropout)(x, deterministic=deterministic)
         return x
@@ -432,7 +432,7 @@ class FFN(nn.Module):
         gain = deepnet_gain["encoder" if self.is_encoder else "decoder"]["beta"](
             self.config
         )
-        if self.config.ln_positions in ["normformer"]:
+        if self.config.ln_positions in ["normformer", "cogview"]:
             x = norm(
                 self.config.ln_type, dtype=self.dtype, epsilon=1e-05, use_scale=False
             )(x)
@@ -460,7 +460,7 @@ class FFN(nn.Module):
             if self.config.use_deepnet_scaling
             else jax.nn.initializers.normal(self.config.init_std),
         )(x)
-        if self.config.ln_positions in ["swinv2"]:
+        if self.config.ln_positions in ["swinv2", "cogview"]:
             x = norm(self.config.ln_type, dtype=self.dtype, epsilon=1e-05)(x)
         x = nn.Dropout(rate=self.config.dropout)(x, deterministic=deterministic)
         return x
@@ -593,7 +593,7 @@ class FlaxBartDecoderLayer(nn.Module):
         residual = hidden_states
 
         # Self Attention
-        if self.config.ln_positions in ["normformer"]:
+        if self.config.ln_positions in ["normformer", "cogview"]:
             hidden_states = norm(
                 self.config.ln_type,
                 dtype=self.dtype,
@@ -615,7 +615,7 @@ class FlaxBartDecoderLayer(nn.Module):
             init_cache=init_cache,
         )
 
-        if self.config.ln_positions in ["normformer", "swinv2"]:
+        if self.config.ln_positions in ["normformer", "swinv2", "cogview"]:
             hidden_states = norm(self.config.ln_type, dtype=self.dtype, epsilon=1e-05)(
                 hidden_states
             )
@@ -632,7 +632,7 @@ class FlaxBartDecoderLayer(nn.Module):
         cross_attn_weights = None
         if encoder_hidden_states is not None:
             residual = hidden_states
-            if self.config.ln_positions in ["normformer"]:
+            if self.config.ln_positions in ["normformer", "cogview"]:
                 hidden_states = norm(
                     self.config.ln_type,
                     dtype=self.dtype,
@@ -652,7 +652,7 @@ class FlaxBartDecoderLayer(nn.Module):
                 key_value_states=encoder_hidden_states,
                 attention_mask=encoder_attention_mask,
             )
-            if self.config.ln_positions in ["normformer", "swinv2"]:
+            if self.config.ln_positions in ["normformer", "swinv2", "cogview"]:
                 hidden_states = norm(
                     self.config.ln_type, dtype=self.dtype, epsilon=1e-05
                 )(hidden_states)