Repository: facebookresearch/xformers

Issue:
xformers ViT-B ImageNet MAE + Deepnorm training instability
# 🐛 Bug

I'm trying to create a 1:1 config that can train a stable ViT-B with the MAE config (from appendix A.2). 

Maybe I'm missing something (highly plausible), but when I use xformers instead of timm it creates an unstable training scenario [over numerous trials] with exactly the same hyper-parameters (batch_size=4096 + cutmix + mixup + label smoothing + AdamW[0.9, 0.95], lr=1e-4 [with scaling rule ofc], lr warmup + cosine decay, skip bias/CLS/pos_embed weight decay, etc, etc).

![image](https://user-images.githubusercontent.com/8204807/156156416-ca13a631-d5ec-4662-a871-2796e82dddce.png)

## xformers ViT-B Config
```python
reversible: False
block_type: "encoder"
num_layers: 12
dim_model: 768
layer_norm_style: "pre"

multi_head_config:
  num_heads: 12
  residual_dropout: 0.1  # (1) tried without this, (2) swapping this for DropPath, (3) with regular dropout
  use_rotary_embeddings: False

  attention:
    name: "scaled_dot_product"
    dropout: 0.0
    causal: False

feedforward_config:
  name: "MLP"
  dropout: 0.0
  activation: "gelu"
  hidden_layer_multiplier: 4
```

## xformers ViT-B 
```python
"""A simple ViT-B in xformers."""

import typing as t
from pathlib import Path
import yaml

import torch
from torch import nn
from timm.models.vision_transformer import DropPath
from timm.models.layers.patch_embed import PatchEmbed
from timm.models.layers.weight_init import trunc_normal_
from xformers.factory import xFormer, xFormerConfig


def _init_vit_weights(module: nn.Module):
    """Transformer weight initialization from TIMM."""
    if isinstance(module, nn.Linear):
        trunc_normal_(module.weight, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
        nn.init.zeros_(module.bias)
        nn.init.ones_(module.weight)


class ViT(nn.Module):
    """Vision transformer + head module."""

    def __init__(
        self,
        img_size: int = 224,
        in_chans: int = 3,
        patch_size: int = 16,
        num_classes: int = 1000,
        drop_path_rate: float = 0,
        transfomer_config_file: t.Union[Path, str] = "configs/vit_b.yaml",
    ):
        """A standard ViT module."""
        super().__init__()

        # read the model config
        with open(transfomer_config_file, "rb") as fileptr:
            self.model_config = yaml.load(fileptr, Loader=yaml.FullLoader)

        # embed_dim = self.model_config["block_config"]["dim_model"]
        embed_dim = self.model_config["dim_model"]
        print(self.model_config)

        # build the patch embedding model
        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            embed_dim=embed_dim,
            flatten=True,
        )
        self.num_patches = self.patch_embed.num_patches

        # Build the tokens / position embeddings
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(  # +1 for CLS token
            torch.zeros(1, self.patch_embed.num_patches + 1, embed_dim)
        )

        # build the backbone
        self.backbone = xFormer.from_config(xFormerConfig([self.model_config]))

        # Swap dropout with drop-path
        # Also tried (1) without this, (2) without dropout.
        if drop_path_rate > 0:
            dpr_idx = 0
            dpr = [
                x.item()
                for x in torch.linspace(0, drop_path_rate, len(self.backbone.encoders))
            ]
            for layer in self.backbone.encoders:
                if hasattr(layer.mha, "resid_drop"):
                    setattr(layer.mha, "resid_drop", DropPath(dpr[dpr_idx]))
                    dpr_idx += 1

        # build the head network
        self.head = nn.Sequential(
            nn.LayerNorm(embed_dim, eps=1e-6),
            nn.Linear(embed_dim, num_classes)
        )

        # do ViT initializations
        self.init_weights()

    def init_weights(self):
        """Initialize layers, pos_embed and CLS for ViTs."""
        trunc_normal_(self.pos_embed, std=.02)
        trunc_normal_(self.cls_token, std=.02)
        self.apply(_init_vit_weights)

    def forward(self, inputs: torch.Tensor) -> t.Dict[str, torch.Tensor]:
        """Infer variates and return a dict with repr and logits.

        Example sizing:
        patches = [2, 196, 768] --> [2, 197, 768] (after CLS)
        representation = [2, 197, 768]
        logits = [2, 197, 1000]
        CLS = [2, 1000]  # select first of 197

        """
        patches = self.patch_embed(inputs)
        cls_token = self.cls_token.expand(inputs.shape[0], -1, -1)
        out = torch.cat((cls_token, patches), dim=1)
        out = out + self.pos_embed

        representation = self.backbone(out)
        logits = self.head(representation)
        return {
            "representation": representation.detach(),
            "logits": logits,
            "CLS": logits[:, 0],
        }
```

## Command

```ipython
vit_b = ViT()
vit_b_timm = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=1000)

number_of_parameters(vit_b)
Out[18]: 86567656

number_of_parameters(vit_b_timm)
Out[19]: 86567656
```
## To Reproduce

Steps to reproduce the behavior:

<!-- If you were running a command, post the exact command that you were running -->

1. Train a ViT-B with xformers --> unstable
2. Using same training setup with timm --> stable
3. 😢 



Discussion:
thanks for the issue, just saw that, looking into it !
@jramapuram can you elaborate on your config, do you have Triton for instance ? Could you share a print(model) here, to be sure of what parts are actually instantiated ? After a quick look it seems that there could be a part where the gradients are not handled up to the same precision as torch, at least that's my #1 hypothesis
many thanks for the detailed issue and code snippets, this is perfect
if you're using Triton, could you test out installing a recent dev package ? `pip install triton==1.1.2.dev20220106`
also @jramapuram could you confirm that this is with torch AMP ? (fp16)
cc @dianaml0, @fmassa, is that something that you've seen ? I remember @xwhan saw that at some point, but I thought that this was fixed. I just did a quick check in the triton code, and we're keeping the data type as fp32 in the softmax and layernorm case when AMP is activated, which should lead to a similar precision as pytorch (layernorm is a bit below). It looks like a vanishing gradient problem, and the parts here are very standard (MLP and scaled_dot_product attention), I'm wondering whether it could be somewhere else in the code, or if the timm ViT adds some parameter-less normalization for instance. I'm not seeing this on the Cifar example that we host

edit: adding some more context and info
@jramapuram the eps parameter for LayerNorm is not the same in between timm and xformers (1e-5 vs. 1e-6), it's a long shot but since your issue could be related to vanishing gradient, could explain. Fixing that
Filling in details:
- AMP FP16 ✅ 
- triton==1.1.1 ( can test 1.1.2.dev20220106 👍  )
- Will try the layernorm eps; good find! Might be relevant for AMP

Instantiated model print to STDOUT: https://gist.github.com/jramapuram/d284e0f261d3fdb15c213dd929d272b9

> Filling in details:
> 
>     * AMP FP16 white_check_mark
> 
>     * triton==1.1.1 ( can test 1.1.2.dev20220106 +1  )
> 
>     * Will try the layernorm eps; good find! Might be relevant for AMP
> 
> 
> Instantiated model print to STDOUT: https://gist.github.com/jramapuram/d284e0f261d3fdb15c213dd929d272b9

I can repro the problem with the minimal microViT example actually (prior to the linked PRs), just need to wait long enough. Testing right now with the changes from the linked PRs
seems fine with the updated eps @jramapuram, let me know if it fixes your issue ? 
Training now; will update here :)

```python
def update_ln_eps(module: nn.Module, new_eps: float):
    """Recurse and update LN eps with this value."""
    from xformers.triton.layer_norm import FusedLayerNorm

    if isinstance(module, torch.nn.modules.LayerNorm):
        module.eps = new_eps

    if isinstance(module, FusedLayerNorm):
        module.epsilon = new_eps

    for _, child in module.named_children():
        update_ln_eps(child, new_eps)
```
@blefaudeux : Unfortunately this has not seemed to fix it for me 😬 . Not sure if the scaling from microViT --> ViT-B ImageNet might be causing some issues that are not easily evident. 

With LN fix using function above:
![image](https://user-images.githubusercontent.com/8204807/156363286-ae507064-9002-4aa3-af34-e95f2dd98272.png)


With Triton `1.1.2.dev20220106` (tested with pip freeze to validate)
![image](https://user-images.githubusercontent.com/8204807/156363402-fe1016b7-3b05-4747-86b4-d80a228f980d.png)


Commit d4c28fbbb881753e7855d08d121c85878a72b775 (tried with and without triton `1.1.2.dev20220106`):
![image](https://user-images.githubusercontent.com/8204807/156401623-af078d55-6ee3-470f-937e-65814db27a2e.png)

For sanity I also tried again swapping back to TIMM and it is still working 😬 
![image](https://user-images.githubusercontent.com/8204807/156401842-6778ab03-ee6d-416a-90f0-4a57734d3e9a.png)


ouch, this is not good.. the issue auto-closed it seems, but keeping it open, I'll try to dig a bit more
@jramapuram to try to pinpoint this a little better (and if you have time), could you try in an environment which does not have Triton ? a few parts will default switch to PyTorch, if you don't see an issue there then I would know where to look (well, softmax and layernorm)
Else I can think of 
- different init strategies for the weights (probable but kind of unlikely that it explains I think)
- shared weights in the projection, 
-  different pre/post normalization
looking into Timm's implementation. 

I can confirm that it does not happen on cifar and a smaller ViT unfortunately, would have been nice to have an easy repro

edit: adding more context
> @jramapuram to try to pinpoint this a little better (and if you have time), could you try in an environment which does not have Triton ? a few parts will default switch to PyTorch, if you don't see an issue there then I would know where to look (well, softmax and layernorm)

testing with pure pytorch layers right now, and I'm not seeing any difference so far, so might not be a good explanation
> Else I can think of
> 
>     * different init strategies for the weights (probable but kind of unlikely that it explains I think)

init is different indeed, [see ](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py#L393) for instance, while xformers mostly follows default pytorch

>     * shared weights in the projection,

the projection seems to follow the same structure, n x 3n matrix + bias, nothing different here

>     * different pre/post normalization

nope, Pre-norm in both  cases



in short I don't see much difference (provided my home test with pytorch vs. triton parts is confirmed on your end @jramapuram) except for weights init, since AMP training is notoriously a little finicky maybe that could explain ? Not super intuitive to me but having a deeper look
- Will try vanilla pytorch (without triton) on ImgNet for my own sanity as well 😅 
- PreNorm ✅ 
- I do the custom TIMM init already (see code above which distills [this](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py#L393) ; will also try a lower std (std=0.01) as well.):
```python
def _init_vit_weights(module: nn.Module):
    """Transformer weight initialization from TIMM."""
    if isinstance(module, nn.Linear):
        trunc_normal_(module.weight, std=0.02)
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
        nn.init.zeros_(module.bias)
        nn.init.ones_(module.weight)
```
@blefaudeux : are there any `xformers` linear layers that don't inherit `nn.Linear` that might be missed by this init function?

Thanks for the great suggestions btw!
> 



> * Will try vanilla pytorch (without triton) on ImgNet for my own sanity as well sweat_smile
> 
>     * PreNorm white_check_mark
> 
>     * I do the custom TIMM init already (see code above which distills [this](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py#L393) ; will also try a lower std (std=0.01) as well.):
> 
> 
> ```python
> def _init_vit_weights(module: nn.Module):
>     """Transformer weight initialization from TIMM."""
>     if isinstance(module, nn.Linear):
>         trunc_normal_(module.weight, std=0.02)
>         if module.bias is not None:
>             nn.init.zeros_(module.bias)
>     elif isinstance(module, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
>         nn.init.zeros_(module.bias)
>         nn.init.ones_(module.weight)
> ```
> 

ahh, I didn't know for the init on your side, so this rules it out also ! 


> @blefaudeux : are there any `xformers` linear layers that don't inherit `nn.Linear` that might be missed by this init function?

No I don't think so, although fused MLP uses a normal nn.Linear but fuses the dropout/bias/activation (so the bias init would be missed). It does not seem like you're using fusedMLP so it should not be the case

> 
> Thanks for the great suggestions btw!

No problem, this is a little perplexing to be honest, but we'll root it out !
seeing your curves, it does seem a little different from what I was seeing prior to the eps adjustment: validation accuracy was collapsing in the microViT example / CIFAR, but over many steps, while yours looks like a complete breakdown, one update completely breaks the model. Really looks like a raw fp16 representation problem, an underflow or overflow would look like that


this is what a faulty normalization floor looked like (eps = 1e-5, pre/post correction), not really what you're seeing, unless it's a logging issue (not logging often enough, but guess is no since I'm seeing your steps axis and you seem to log per step)
![gnome-shell-screenshot-a985s4](https://user-images.githubusercontent.com/3072861/156427920-f4f5fe51-eb51-46de-bbd7-cd03163c6766.png)


hmm turns out I was testing with rotary embeddings turned on, and they make a huge difference
Lower std on trunc normal init (0.01):
![image](https://user-images.githubusercontent.com/8204807/156474949-96eb8a32-3d5c-45be-866b-f0f1436dbe0c.png)

Without triton:
![image](https://user-images.githubusercontent.com/8204807/156545108-949eb862-e97a-4767-b0c3-d6a9ffe498b2.png)
```bash
WARNING:root:Triton is not available, some optimizations will not be enabled.
Error No module named 'triton'
```

FusedMLP: 
![image](https://user-images.githubusercontent.com/8204807/156541108-a18ab1e6-5f79-45d9-900c-06e6447fcdbe.png)
```bash
        (feedforward): FusedMLP(
          (mlp): Sequential(
            (0): Linear(in_features=768, out_features=3072, bias=True)
            (1): FusedDropoutBias(
              (pytorch_activation): GELU()
            )
            (2): Linear(in_features=3072, out_features=768, bias=True)
            (3): FusedDropoutBias(
              (pytorch_activation): Passthrough()
            )
          )
        )
```
thanks @jramapuram, it's very informative, so no issues with the triton layers whatsover, the problem is in a pure pytorch definition.. :/
checking for what could cause an underflow or overflow in between Timm's implementation and ours, looks like the sqrt(d) normalization is done post-hoc in timm ([see](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py#L222)) while we do it prior to computing the attention ([see](https://github.com/facebookresearch/xformers/blob/dc80aeb7990889752b67d07ba646a7218bc1bced/xformers/components/attention/core.py#L196)). cc @fmassa @dianaml0, thoughts ?

@jramapuram can you check out https://github.com/facebookresearch/xformers/pull/225 ? I'm trying to find a test to repro this issue
~~oh @jramapuram, something else which could be significant: you can check out [this paper](https://arxiv.org/pdf/2106.10270.pdf) (end of page 3) and the matching timm-vit model, Ross drops the hidden layer in the last block citing "training unstabilities" ! I'll add an option to xformers to be able to do this from the config side (pass "multiplier" == -1), it could well be the reason~~ 

nevermind, I just realized that you were already doing that..
Running with `issue_219` branch 😢  :

<img width="1204" alt="image" src="https://user-images.githubusercontent.com/8204807/156891031-6f02d960-b971-4f8e-b6c0-d4945cd4db4f.png">

```bash
Branch 'issue_219' set up to track remote branch 'issue_219' from 'origin'.
...
Building wheels for collected packages: xformers
  Building wheel for xformers (setup.py): started
  Building wheel for xformers (setup.py): still running...
  Building wheel for xformers (setup.py): finished with status 'done'
  Created wheel for xformers: filename=xformers-0.0.9-cp39-cp39-linux_x86_64.whl size=1337425 sha256=8b0825f014b9859ad3786042973af9d389b1031c68a1eb734cb7c4b78050ce08
  Stored in directory: /mnt/tmp/pip-ephem-wheel-cache-a4wszhv8/wheels/10/9b/72/7597306f87828c97afa86adf56f4c78db426e51c2ee02d2f66
Successfully built xformers
Installing collected packages: xformers
  Attempting uninstall: xformers
    Found existing installation: xformers 0.0.7
    Uninstalling xformers-0.0.7:
      Successfully uninstalled xformers-0.0.7
Successfully installed xformers-0.0.9
```
> Running with `issue_219` branch cry :
> 
> <img alt="image" width="1204" src="https://user-images.githubusercontent.com/8204807/156891031-6f02d960-b971-4f8e-b6c0-d4945cd4db4f.png">
> 
> ```shell
> Branch 'issue_219' set up to track remote branch 'issue_219' from 'origin'.
> ...
> Building wheels for collected packages: xformers
>   Building wheel for xformers (setup.py): started
>   Building wheel for xformers (setup.py): still running...
>   Building wheel for xformers (setup.py): finished with status 'done'
>   Created wheel for xformers: filename=xformers-0.0.9-cp39-cp39-linux_x86_64.whl size=1337425 sha256=8b0825f014b9859ad3786042973af9d389b1031c68a1eb734cb7c4b78050ce08
>   Stored in directory: /mnt/tmp/pip-ephem-wheel-cache-a4wszhv8/wheels/10/9b/72/7597306f87828c97afa86adf56f4c78db426e51c2ee02d2f66
> Successfully built xformers
> Installing collected packages: xformers
>   Attempting uninstall: xformers
>     Found existing installation: xformers 0.0.7
>     Uninstalling xformers-0.0.7:
>       Successfully uninstalled xformers-0.0.7
> Successfully installed xformers-0.0.9
> ```

thanks a lot for the test, it's really perplexing.. @fmassa has been using some xformers block a lot with imagenet, but not the whole model, I don't remember him getting this. I'll check the weight inits as soon as I get the time, sorry for the delay
Hey @jramapuram, back to you ! We can exchange by mail if that helps, I'd really love to get to the bottom of this. I'm adding DeepNet to xformers, some of the issues mentioned with pre-ln look like they could apply here (although it does not exaplain why Timm's implementation does not face the same instabilities). Current thoughts are that it could be related to weight init (I know that you're handling that already, but bug/issue ?) or different LR layer wise for instance, and a grep which fails on the xformers model because of different names ? Would it be possible to share more of the training code ? 

Resolving patch:
diff --git a/xformers/factory/model_factory.py b/xformers/factory/model_factory.py
--- a/xformers/factory/model_factory.py
+++ b/xformers/factory/model_factory.py
@@ -184,22 +184,10 @@ def __init__(
         )
         self.decoders = torch.nn.ModuleList(decoders)
 
-        if len(self.decoders) > 0:
-            # Use Xavier init for encoding/decoding tasks
-            self._reset_parameters()
-
     @classmethod
     def from_config(cls, config: xFormerConfig):
         return cls(config.stack_configs, config.tie_embedding_weights)
 
-    def _reset_parameters(self):
-        r"""Initiate parameters in the transformer model
-        following the Xavier distribution."""
-
-        for p in self.parameters():
-            if p.dim() > 1:
-                torch.nn.init.xavier_uniform_(p)
-
     def _verify_reversible(self, stack_configs: List[xFormerBlockConfig]):
         reversible = [
             c.reversible


Test patch: