Repository: facebookresearch/xformers

Issue:
Can't deepcopy an xformer model with triton 2 update
# 🐛 Bug

Given the recent triton2 update an xformer model cannot be deep-copied. This is an important requirement for numerous tasks including EMA (without knowledge of the generating class / hyper-params).

## xformers ViT-B Config

```yaml
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
```

## To reproduce

```python
from copy import deepcopy

with open(transfomer_config_file, "rb") as fileptr:
    self.model_config = yaml.load(fileptr, Loader=yaml.FullLoader)

model = xFormer.from_config(xFormerConfig([self.model_config]))
deepcopy(model)
```

Error is:
```bash
TypeError: cannot pickle 'PyCapsule' object
```


Discussion:
Oh jeez.. OK, you can remove triton from your env and this should unlock short term, I'll have a look in the meantime :) 
For ref even the SWA impl in vanilla pytorch relies on deepcopy: https://github.com/pytorch/pytorch/blob/master/torch/optim/swa_utils.py#L100 -- not being able to EMA is a ~1-2% loss in overall perf 😬 

Resolving patch:
diff --git a/xformers/triton/dropout.py b/xformers/triton/dropout.py
--- a/xformers/triton/dropout.py
+++ b/xformers/triton/dropout.py
@@ -7,7 +7,7 @@
 # CREDITS: This comes almost as-is from the Triton dropout tutorial
 # https://raw.githubusercontent.com/openai/triton/master/python/tutorials/04-low-memory-dropout.py
 
-from typing import Optional
+from typing import Any, Optional
 
 import torch
 import triton
@@ -196,6 +196,11 @@ def dropout(
 
 
 class FusedDropoutBias(torch.nn.Module):
+    """
+    A layer which fuses the computation of Dropout(Activation(x))
+    in a single GPU kernel
+    """
+
     def __init__(
         self,
         p: float,
@@ -216,15 +221,24 @@ def __init__(
             if bias_shape is not None
             else None
         )
-        self.activation = get_triton_activation_kernel(activation)
-        self.pytorch_activation = build_activation(self.activation_type)
-        self.activation_grad = get_triton_activation_bwd_kernel(activation)
+
+        self.activation: Optional[Any] = None
+        self.activation_grad: Optional[Any] = None
+        self.activation_pytorch: Optional[Any] = None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Convenience, catch a possible type or device mismatch
         if self.bias is not None:
             self.bias = self.bias.to(dtype=x.dtype, device=x.device)  # type: ignore
 
+        # Lazy init (helps with pickling)
+        if self.activation is None:
+            self.activation = get_triton_activation_kernel(self.activation_type)
+            self.pytorch_activation = build_activation(self.activation_type)
+            self.activation_grad = get_triton_activation_bwd_kernel(
+                self.activation_type
+            )
+
         # Train/inference
         p = self.p if self.training else 0.0
 


Test patch:
diff --git a/tests/test_block_factory.py b/tests/test_block_factory.py
--- a/tests/test_block_factory.py
+++ b/tests/test_block_factory.py
@@ -27,13 +27,13 @@
 VOCAB_SIZE = 64
 
 
-@pytest.mark.parametrize("attn_dropout", [0.0, 0.1])
-@pytest.mark.parametrize("residual_dropout", [0.0, 0.1])
+@pytest.mark.parametrize("attn_dropout", [0.1])
+@pytest.mark.parametrize("residual_dropout", [0.1])
 @pytest.mark.parametrize("heads", [1, 2])
 @pytest.mark.parametrize("activation", [a.value for a in Activation])
 @pytest.mark.parametrize("attention_name", ATTENTION_REGISTRY.keys())
 @pytest.mark.parametrize("feedforward_name", FEEDFORWARD_REGISTRY.keys())
-@pytest.mark.parametrize("layer_norm_style", ["pre", "post"])
+@pytest.mark.parametrize("layer_norm_style", ["pre", "post", "deepnorm"])
 @pytest.mark.parametrize("device", DEVICES)
 @pytest.mark.parametrize("reversible", [True, False])
 @pytest.mark.skipif(
@@ -127,15 +127,15 @@ def test_xformer_encoder_block(
     _ = block(inputs, input_mask=input_mask)
 
 
-@pytest.mark.parametrize("attn_dropout", [0.0, 0.1])
-@pytest.mark.parametrize("residual_dropout", [0.0, 0.1])
+@pytest.mark.parametrize("attn_dropout", [0.1])
+@pytest.mark.parametrize("residual_dropout", [0.1])
 @pytest.mark.parametrize("causal", [True, False])
 @pytest.mark.parametrize("heads", [1, 2])
 @pytest.mark.parametrize("activation", [a.value for a in Activation])
 @pytest.mark.parametrize("rotary_embeddings", [False, True])
 @pytest.mark.parametrize("attention_name", ATTENTION_REGISTRY.keys())
 @pytest.mark.parametrize("feedforward_name", FEEDFORWARD_REGISTRY.keys())
-@pytest.mark.parametrize("layer_norm_style", ["pre", "post"])
+@pytest.mark.parametrize("layer_norm_style", ["pre", "post", "deepnorm"])
 @pytest.mark.parametrize("device", DEVICES)
 @pytest.mark.skipif(
     not torch.cuda.is_available(), reason="This test requires a CUDA device"
diff --git a/tests/test_pickling.py b/tests/test_pickling.py
--- a/tests/test_pickling.py
+++ b/tests/test_pickling.py
@@ -7,14 +7,17 @@
 # https://github.com/facebookresearch/xformers/issues/203
 
 import pickle
+from copy import deepcopy
 
+import pytest
 from torch import nn
 
+from xformers import _is_triton_available
 from xformers.factory import xFormer, xFormerConfig
 
 test_config = [
     {
-        "reversible": False,  # Turn on to test the effect of using reversible layers
+        "reversible": False,
         "block_type": "encoder",
         "num_layers": 2,
         "dim_model": 768,
@@ -30,7 +33,7 @@
             },
         },
         "feedforward_config": {
-            "name": "MLP",  # FIXME: Test with FusedMLP also
+            "name": "FusedMLP",
             "dropout": 0.1,
             "activation": "gelu",
             "hidden_layer_multiplier": 4,
@@ -40,11 +43,20 @@
 
 
 class ViT(nn.Module):
-    def __init__(self):
+    def __init__(self, mlp):
         super().__init__()
-        self.xformer = xFormer.from_config(xFormerConfig(test_config))
+        test_config[0]["feedforward_config"]["name"] = mlp
+        xformer_config = xFormerConfig(test_config)
+        self.xformer = xFormer.from_config(xformer_config)
 
 
-def test_pickling():
-    test = ViT()
-    pickle.dumps(test)
+MLPs = ["MLP"]
+if _is_triton_available:
+    MLPs.append("FusedMLP")
+
+
+@pytest.mark.parametrize("mlp", MLPs)
+def test_pickling(mlp):
+    test = ViT(mlp)
+    _ = pickle.dumps(test)
+    _ = deepcopy(test)