using the fused dropout in the FusedMLP

facebookresearch · Nov 9, 2021 · cf69993 · cf69993
1 parent 5843434
commit cf69993
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 8 deletions.
diff --git a/xformers/components/feedforward/fused_mlp.py b/xformers/components/feedforward/fused_mlp.py
@@ -19,7 +19,7 @@
 
 if torch.cuda.is_available():
     try:
-        from xformers.triton import FusedLinear
+        from xformers.triton import FusedDropoutBias, FusedLinear
 
         @dataclass
         class FusedMlpConfig(FeedforwardConfig):
@@ -39,8 +39,8 @@ def __init__(
                 dropout: float,
                 activation: Activation,
                 hidden_layer_multiplier: int,
-                *args,
-                **kwargs,
+                *_,
+                **__,
             ):
                 super().__init__()
 
@@ -51,11 +51,13 @@ def __init__(
                         in_features=dim_model,
                         out_features=hidden_layer_multiplier * dim_model,
                         activation=activation,
-                        bias=True,
+                        bias=False,
                     ),
-                    nn.Dropout(dropout),
-                    nn.Linear(hidden_layer_multiplier * dim_model, dim_model),
-                    nn.Dropout(dropout),
+                    FusedDropoutBias(dropout, hidden_layer_multiplier * dim_model),
+                    nn.Linear(
+                        hidden_layer_multiplier * dim_model, dim_model, bias=False
+                    ),
+                    FusedDropoutBias(dropout, hidden_layer_multiplier * dim_model),
                 )
                 self.requires_cuda = True
 

diff --git a/xformers/triton/__init__.py b/xformers/triton/__init__.py
@@ -9,7 +9,7 @@
 _triton_available = torch.cuda.is_available()
 if _triton_available:
     try:
-        from .dropout import dropout  # noqa
+        from .dropout import FusedDropoutBias, dropout  # noqa
         from .fused_linear_layer import FusedLinear  # noqa
         from .layer_norm import FusedLayerNorm, layer_norm  # noqa
         from .softmax import log_softmax, softmax  # noqa
@@ -18,6 +18,7 @@
             "dropout",
             "softmax",
             "log_softmax",
+            "FusedDropoutBias",
             "FusedLinear",
             "FusedLayerNorm",
             "layer_norm",

diff --git a/xformers/triton/dropout.py b/xformers/triton/dropout.py
@@ -102,3 +102,13 @@ def dropout(x: torch.Tensor, p: float, bias: Optional[torch.Tensor] = None):
         return _dropout.apply(x, p, bias)
 
     return x + bias if bias is not None else x
+
+
+class FusedDropoutBias(torch.nn.Module):
+    def __init__(self, p: float, bias_shape: Optional[int]) -> None:
+        super().__init__()
+        self.p = p
+        self.bias = torch.zeros(bias_shape) if bias_shape is not None else None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return dropout(x, self.p, self.bias)