Repository: facebookresearch/xformers

Issue:
Config does not support model_dim that does not match attention model_dim
# 🐛 Bug

Adding a `dim_model` within a `position_encoding_config` that is not equal to the `dim_model` of the rest of the model results in an unusable model, as the result of the embedding layer is no longer compatible with the feedforward  layers within the attention blocks.

This looks like a bug to me, since it doesn't ever allow you to specify a `dim_model` in the vocabulary that is not equal to the `dim_model` of the global model. In my use-case, in which I would like to tie the input of embeddings of two different models with different hidden sizes, it would be nice to have this work by inserting a linear layer in between the embedding layer and the first attention layer. 

## Command

Using the following config results in a crash:

```python
config = [{'reversible': False,
  'block_type': 'encoder',
  'num_layers': 4,
  'dim_model': 64,
  'layer_norm_style': 'pre',
  'position_encoding_config': {'name': 'vocab',
   'seq_len': 512,
   'vocab_size': 30522,
   'dim_model': 256},
  'multi_head_config': {'num_heads': 1,
   'residual_dropout': 0,
   'use_rotary_embeddings': True,
   'attention': {'name': 'scaled_dot_product',
    'dropout': 0.1,
    'causal': False,
    'seq_len': 512,
    'num_rules': 4}},
  'feedforward_config': {'name': 'MLP',
   'dropout': 0.1,
   'activation': 'gelu',
   'hidden_layer_multiplier': 4}}]

import torch
from xformers.factory.model_factory import xFormer, xFormerConfig

cfg = xFormerConfig(config)
model = xFormer.from_config(cfg)

X = torch.randint(0, 1000, (10, 256))

model(X)
```

```
RuntimeError: Given normalized_shape=[64], expected input with shape [*, 64], but got input of size[10, 256, 25
6]
```


## To Reproduce

See above

## Expected behavior

I'd expect the model builder to complain that adding a `dim_model` within a vocab block does not make sense, or I expect it to work.




Discussion:
Hi @stephantul, thanks for the report ! I agree with the error flagging, it's done down the line since the dimensions do not match but we could either catch that at config building time (a) or add an error message on top of this 'RuntimeError' (b) to make it a little easier to debug. I think that we should add a dimensionality check, option (a), feels like a relatively easy take. 

Note that if I remember correctly you can just pass dim_model once, and the other fields are auto filled with this value.

On the other hand the architecture with an extra projection layer is not "standard" Transformer, I would label that as a feature, could be interesting indeed. Do you know if there are existing papers testing out this take, out of curiosity ? 
Hi @blefaudeux,

Thanks for the quick response! I got the idea from [this repo](https://github.com/lucidrains/electra-pytorch/blob/master/pretraining/openwebtext/pretrain.py#L129). In this electra implementation, the embeddings of the discriminator and the generator are tied, but the hidden dimension of the discriminator is 256, while the hidden size of the generator is 64. This size discrepancy is necessary for electra pretraining to work (i.e., the generator needs to be about 1/4 of the discriminator in size).

The repo I linked works with the huggingface implementation, which solves it by using an embeddings projector, see [here](https://github.com/huggingface/transformers/blob/c30798ec9dcc8813bbe20b04ddabca5ba9fc5c30/src/transformers/models/electra/modeling_electra.py#L914).

I'm sure we can work around this somehow, but it would be a nice feature to perhaps add "projector layers" in between parts of the model that have a dimensionality mismatch when loading a config (maybe with an optional boolean flag to not let users accidentally add these projector layers?)
Sounds like a plan, I like the idea actually, seems like a good fit with the existing codebase, can be extended to do just that J believe

I can have a look later (time permitting) or feel free to submit a PR if you'd like

Resolving patch:
diff --git a/xformers/factory/block_factory.py b/xformers/factory/block_factory.py
--- a/xformers/factory/block_factory.py
+++ b/xformers/factory/block_factory.py
@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import logging
 from dataclasses import asdict, dataclass
 from enum import Enum
 from typing import Any, Dict, Optional, Tuple, Union
@@ -39,7 +40,7 @@ class LayerPositionBitmask(int, Enum):
 
 
 class LayerPosition:
-    """ Bitmask to mark this layer as first, last, nothing or both"""
+    """Bitmask to mark this layer as first, last, nothing or both"""
 
     def __init__(self):
         self.bitmask = LayerPositionBitmask.Default
@@ -286,7 +287,7 @@ def __init__(
 
 
 class xFormerEncoderBlock(torch.nn.Module):
-    r""" A vanilla Transformer Encoder block """
+    r"""A vanilla Transformer Encoder block"""
 
     def __init__(self, config: xFormerEncoderConfig, **kwargs):
         super().__init__()
@@ -297,11 +298,26 @@ def __init__(self, config: xFormerEncoderConfig, **kwargs):
         self.dim_model = config.dim_model
 
         # If this layer is the first one, and a pose encoding has been requested
-        self.pose_encoding = (
-            build_positional_embedding(asdict(config.position_encoding_config))
-            if config.position_encoding_config and config.layer_position.is_first()
-            else None
-        )
+        if (
+            config.position_encoding_config is not None
+            and config.layer_position.is_first()
+        ):
+            self.pose_encoding = build_positional_embedding(
+                asdict(config.position_encoding_config)
+            )
+
+            pos_encoding_dim = config.position_encoding_config.dim_model
+            mha_dim = config.multi_head_config["dim_model"]
+
+            if pos_encoding_dim != mha_dim:
+
+                logging.warning(
+                    f"The embedding dim and model dim do not match ({pos_encoding_dim} vs {mha_dim}), adding a projector layer."  # noqa
+                )
+
+                self.embedding_projector = nn.Linear(pos_encoding_dim, mha_dim)
+        else:
+            self.pose_encoding = None
 
         if config.layer_norm_style == LayerNormStyle.DeepNorm:
             # Just use the layer norm coefficient here,
@@ -361,9 +377,12 @@ def forward(
         att_mask: Optional[torch.Tensor] = None,
         input_mask: Optional[torch.Tensor] = None,
     ):
-        if self.pose_encoding:
+        if self.pose_encoding is not None:
             x = self.pose_encoding(x)
 
+            if hasattr(self, "embedding_projector"):
+                x = self.embedding_projector(x)
+
         # Handle the optional input masking, differs on Q, K, V
         if input_mask is not None:
             q = x
@@ -388,11 +407,26 @@ def __init__(self, config: xFormerDecoderConfig, **kwargs):
         super().__init__()
 
         # If this layer is the first one, and a pose encoding as been requested
-        self.pose_encoding = (
-            build_positional_embedding(config.position_encoding_config)
-            if config.position_encoding_config and config.layer_position.is_first()
-            else None
-        )
+        if (
+            config.position_encoding_config is not None
+            and config.layer_position.is_first()
+        ):
+            self.pose_encoding = build_positional_embedding(
+                config.position_encoding_config
+            )
+
+            pos_encoding_dim = config.position_encoding_config.dim_model
+            mha_dim = config.multi_head_config_masked["dim_model"]
+
+            if pos_encoding_dim != mha_dim:
+
+                logging.warning(
+                    f"The embedding dim and model dim do not match ({pos_encoding_dim} vs {mha_dim}), adding a projector layer."  # noqa
+                )
+
+                self.embedding_projector = nn.Linear(pos_encoding_dim, mha_dim)
+        else:
+            self.pose_encoding = None
 
         if config.layer_norm_style == LayerNormStyle.DeepNorm:
             # Just use the layer norm coefficient here,
@@ -439,9 +473,12 @@ def forward(
         decoder_att_mask: Optional[torch.Tensor] = None,
         input_mask: Optional[torch.Tensor] = None,
     ):
-        if self.pose_encoding:
+        if self.pose_encoding is not None:
             target = self.pose_encoding(target)
 
+            if hasattr(self, "embedding_projector"):
+                target = self.embedding_projector(target)
+
         # Handle the optional input masking, differs on Q, K, V
         if input_mask is not None:
             target_q = target


Test patch: