Repository: facebookresearch/xformers

Issue:
feat(triton): InplaceNorm + InstanceNorm
I'd love to run LayerNorm in place and ideally also add InstanceNorm (by extracting the core normalization from LayerNorm) as HomebrewNLP is currently using a slow PyTorch-level [implementation](https://github.com/HomebrewNLP/HomebrewNLP/blob/master/src/model.py#L32-L53) with a correct backward pass.

While we're at it, optionally fusing GLU and [GLUv2](https://arxiv.org/abs/2102.11856) (`gelu(f(x)) * g(x) + gelu(h(x))`) with various activation functions and normalization might give another speed boost.

To add this myself, I'd need to fully understand triton's pointers and how to access the output instead of input in your LayerNorm implementation. Could you help me with that? or would you instead implement this yourself? Is this even in the scope of xformers?


Discussion:
that's a great question ! definitely open for contributions, and this looks very reasonable, there's a good chance that Triton gives something a lot faster than pytorch there. In terms of scope I think that it's very much ok, as xformers to me is also an optimized parts zoo (with some automatic builders for them, but that's optional).

Just a couple of caveats to begin with:
- while some Triton subparts will probably work with a gtx 1xxx (if you're not using the tensor cores basically, no dot product), it's not guaranteed, so having a pytorch fallback is good to keep in mind. 

- we keep all our triton code [here](https://github.com/facebookresearch/xformers/tree/main/xformers/triton), `k_xx` means `kernel` (the @jit code), put aside so that it does not distort the code coverage metrics. LayerNorm is [there](https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_layer_norm.py#L23) (basically from the tutorial from @ptillet), making the norm parametrizable is a great idea, it's how the fused linear layer [works](https://github.com/facebookresearch/xformers/blob/main/xformers/triton/k_fused_matmul.py#L114) (activations are [here](https://github.com/facebookresearch/xformers/blob/main/xformers/triton/activations.py#L60)), Triton is awesome for that since it will generate a fused kernel on the fly for anything you pass (so this becomes easily extensible)

- who does what ? PR most certainly welcome if you feel like it, else I can put up a PR which does not change layernorm per say, but makes the norm a @jit function that you can pass in, like it's done for the fused linear layer. This way you could easily experiment with other norms and submit them in PRs. As you prefer ! 

- Mini heads up in that for now gelu is a little on the slow side with Triton (80% of pytorch, give or take), could be that this improves over time, and that the combination is still worth it. In any case it would be easy enough to try


Resolving patch:
diff --git a/xformers/triton/k_layer_norm.py b/xformers/triton/k_layer_norm.py
--- a/xformers/triton/k_layer_norm.py
+++ b/xformers/triton/k_layer_norm.py
@@ -18,9 +18,32 @@
 _triton_registered_warnings = False
 
 
-# fmt: off
 @triton.jit
-def layer_norm_fw(X, Y, W, B, M, V, stride, N, eps, **META):
+def _affine(W, B, N, x, META):
+    cols = tl.arange(0, META["BLOCK_SIZE_N"])
+
+    w = tl.load(W + cols, mask=cols < N, other=1.0)
+    zero = 0.0
+    zero = zero.to(w.dtype)  # Triton bug workarounds
+    w = tl.where(cols < N, w, zero)
+
+    b = tl.load(B + cols, mask=cols < N, other=0.0)
+    b = tl.where(cols < N, b, zero)
+    y = x * w + b
+    return y
+
+
+@triton.jit
+def _store(y, Y, stride, N, META):
+    row = tl.program_id(0)
+    cols = tl.arange(0, META["BLOCK_SIZE_N"])
+
+    y_ptrs = Y + row * stride + cols
+    tl.store(y_ptrs, y, mask=cols < N)
+
+
+@triton.jit
+def _layer_norm_non_affine(X, M, V, stride, N, eps, META):
     # fmt: on
     """
     Fused layernorm kernel over a 3d tensor.
@@ -36,7 +59,7 @@ def layer_norm_fw(X, Y, W, B, M, V, stride, N, eps, **META):
     # Move to this row
     x_ptrs = X + row * stride + cols
     x = tl.load(x_ptrs, mask=cols < N, other=0.0).to(tl.float32)
-    x = tl.where(cols < N, x, 0.)  # Triton bug workarounds
+    x = tl.where(cols < N, x, 0.0)  # Triton bug workarounds
 
     # Compute variance
     x_mean = tl.sum(x, axis=0) / N
@@ -49,29 +72,40 @@ def layer_norm_fw(X, Y, W, B, M, V, stride, N, eps, **META):
     tl.store(M + row, x_mean)
     tl.store(V + row, x_inv_sigma)
 
-    # Normalize the inputs
-    w = tl.load(W + cols, mask=cols < N, other=1.0)
-    zero = 0.
-    zero = zero.to(w.dtype)  # Triton bug workarounds
-    w = tl.where(cols < N, w, zero)
+    return x_zm * x_inv_sigma
 
-    b = tl.load(B + cols, mask=cols < N, other=0.0)
-    b = tl.where(cols < N, b, zero)
-    y = x_zm * x_inv_sigma * w + b
 
-    # write back to Y.
-    y_ptrs = Y + row * stride + cols
-    tl.store(y_ptrs, y, mask=cols < N)
+# fmt: off
+@triton.jit
+def _layer_norm_non_affine_fw(X, Y, M, V, stride, N, eps, **META):
+    _store(_layer_norm_non_affine(X, M, V, stride, N, eps, META), Y, stride, N, META)
+
+
+# fmt: off
+@triton.jit
+def _layer_norm_fw(X, Y, W, B, M, V, stride, N, eps, **META):
+    # fmt: on
+    """
+    Fused layernorm kernel over a 3d tensor.
+    The layer norm is applied over the last dimension.
+
+    Compute
+        y = (x - E(x))/(sqrt(var(x) + epsilon)) * gamma + beta
+    """
+    y = _layer_norm_non_affine(X, M, V, stride, N, eps, META)
+    y = _affine(W, B, N, y, META)
+
+    _store(y, Y, stride, N, META)
 
 
 # Backward pass (DX + partial DW + partial DB)
 # fmt: off
 @triton.jit
 def _layer_norm_bwd_dx_fused(
-        DX, DY, DW, DB,
-        X, W, M, V,
-        Lock, stride, N,
-        **META
+    DX, DY, DW, DB,
+    Y, W, B, V,
+    Lock, stride, N,
+    **META
 ):
     # fmt: on
 
@@ -83,9 +117,10 @@ def _layer_norm_bwd_dx_fused(
     cols = tl.arange(0, BLOCK_SIZE_N)
 
     # offset data pointers to start at the row of interest
-    x_ptrs = X + row * stride + cols
+    y_ptrs = Y + row * stride + cols
     dy_ptrs = DY + row * stride + cols
     w_ptrs = W + cols
+    b_ptrs = B + cols
 
     # offset locks and weight/bias gradient pointer
     # each kernel instance accumulates partial sums for
@@ -97,15 +132,15 @@ def _layer_norm_bwd_dx_fused(
     Count = Lock + GROUP_SIZE_M
 
     # load data to SRAM
-    x = tl.load(x_ptrs, mask=cols < N, other=0).to(tl.float32)
+    y = tl.load(y_ptrs, mask=cols < N, other=0).to(tl.float32)
     dy = tl.load(dy_ptrs, mask=cols < N, other=0).to(tl.float32)
     w = tl.load(w_ptrs, mask=cols < N, other=0).to(tl.float32)
+    b = tl.load(b_ptrs, mask=cols < N, other=0).to(tl.float32)
 
-    mean = tl.load(M + row)
     rstd = tl.load(V + row)
 
     # compute dx
-    xhat = (x - mean) * rstd
+    xhat = (y - b) / w
     wdy = w * dy
     xhat = tl.where(cols < N, xhat, 0.0)
     wdy = tl.where(cols < N, wdy, 0.0)
@@ -114,13 +149,11 @@ def _layer_norm_bwd_dx_fused(
     dx = (wdy - (xhat * mean1 + mean2)) * rstd
 
     # write-back dx
-    cols = tl.arange(0, BLOCK_SIZE_N)
-    dx_ptrs = DX + row * stride + cols
-    tl.store(dx_ptrs, dx, mask=cols < N)
+    _store(dx, DX, stride, N, META)
 
     # accumulate partial sums for dw/db
     partial_dw = (dy * xhat).to(w.dtype)
-    partial_db = (dy).to(w.dtype)
+    partial_db = dy.to(w.dtype)
 
     # - wait for a lock on the accumulated dw/db
     while tl.atomic_cas(Lock, 0, 1) == 1:
@@ -146,6 +179,40 @@ def _layer_norm_bwd_dx_fused(
     tl.atomic_xchg(Lock, 0)
 
 
+@triton.jit
+def _layer_norm_no_affine_bwd(
+    DX, DY,
+    Y, V,
+    stride, N,
+    **META
+):
+    # fmt: on
+
+    # position of elements processed by this program
+    row = tl.program_id(0)
+    cols = tl.arange(0, META["BLOCK_SIZE_N"])
+
+    # offset data pointers to start at the row of interest
+    y_ptrs = Y + row * stride + cols
+    dy_ptrs = DY + row * stride + cols
+
+    # load data to SRAM
+    y = tl.load(y_ptrs, mask=cols < N, other=0).to(tl.float32)
+    dy = tl.load(dy_ptrs, mask=cols < N, other=0).to(tl.float32)
+
+    rstd = tl.load(V + row)
+
+    # compute dx
+    xhat = tl.where(cols < N, y, 0.0)
+    wdy = tl.where(cols < N, dy, 0.0)
+    mean1 = tl.sum(xhat * wdy, axis=0) / N
+    mean2 = tl.sum(wdy, axis=0) / N
+    dx = (wdy - (xhat * mean1 + mean2)) * rstd
+
+    # write-back dx
+    _store(dx, DX, stride, N, META)
+
+
 # Backward pass (total DW + total DB)
 # fmt: off
 @triton.jit
@@ -178,7 +245,6 @@ class _LayerNorm(torch.autograd.Function):
     @staticmethod
     @custom_fwd(cast_inputs=torch.float16 if _triton_layernorm_fp16_enabled else None)
     def forward(ctx, x, weight, bias, eps):
-
         # allocate output
         y = torch.empty_like(x)
 
@@ -212,30 +278,41 @@ def forward(ctx, x, weight, bias, eps):
 
         # enqueue kernel
         # fmt: off
-        layer_norm_fw[(M,)](
-            x_arg, y, weight, bias, mean, rstd,
-            x_arg.stride(0),
-            N,
-            eps,
-            num_warps=num_warps,
-            BLOCK_SIZE_N=BLOCK_SIZE_N
-        )
+        if weight is None:
+            _layer_norm_non_affine_fw[(M,)](
+                x_arg, y, mean, rstd,
+                x_arg.stride(0),
+                N,
+                eps,
+                num_warps=num_warps,
+                BLOCK_SIZE_N=BLOCK_SIZE_N
+            )
+        else:
+            _layer_norm_fw[(M,)](
+                x_arg, y, weight, bias, mean, rstd,
+                x_arg.stride(0),
+                N,
+                eps,
+                num_warps=num_warps,
+                BLOCK_SIZE_N=BLOCK_SIZE_N
+            )
         # fmt: on
 
-        ctx.save_for_backward(x, weight, mean, rstd)
+        ctx.save_for_backward(y, rstd, weight, bias)
         ctx.BLOCK_SIZE_N = BLOCK_SIZE_N
         ctx.num_warps = num_warps
         ctx.eps = eps
+        ctx.N = N
 
         return y.reshape_as(x)
 
     @staticmethod
     @custom_bwd
     def backward(ctx, dy):
-        x, weight, mean, var = ctx.saved_tensors
+        y, var, weight, bias = ctx.saved_tensors
 
         # heuristics for amount of parallel reduction stream for DG/DB
-        N = weight.shape[0]
+        N = y.size(-1)
         GROUP_SIZE_M = 64
         if N <= 8192:
             GROUP_SIZE_M = 96
@@ -246,38 +323,44 @@ def backward(ctx, dy):
 
         # flatten the batch dimension, if any.
         # We're interested in 'samples' x norm_dimension
-        x_arg = x.reshape(-1, x.shape[-1])
-        M, N = x_arg.shape
+        y = y.reshape(-1, y.size(-1))
+        M, N = y.size()
 
         # allocate output
         locks = torch.zeros(2 * GROUP_SIZE_M, dtype=torch.int32, device="cuda")
-        t_args = {"dtype" : x.dtype, "device" : x.device}
-        _dw = torch.empty((GROUP_SIZE_M, weight.shape[0]), **t_args)
-        _db = torch.empty((GROUP_SIZE_M, weight.shape[0]), **t_args)
-        dw = torch.empty((weight.shape[0],), **t_args)
-        db = torch.empty((weight.shape[0],), **t_args)
+        t_args = {"dtype": y.dtype, "device": y.device}
+        _dw = torch.empty((GROUP_SIZE_M, y.size(-1)), **t_args)
+        _db = torch.empty((GROUP_SIZE_M, y.size(-1)), **t_args)
+        dw = torch.empty((y.size(-1),), **t_args)
+        db = torch.empty((y.size(-1),), **t_args)
         dy = dy.contiguous()
         dx = torch.empty_like(dy)
 
         # Check the tensor shapes and layouts
         # we suppose in the kernel that they have the same size and are contiguous
-        assert dx.numel() == x.numel(), \
+        assert dx.numel() == y.numel(), \
             "Something is wrong in the backward graph, possibly because of an inplace operation after the layernorm"
 
         # enqueue kernel using forward pass heuristics
         # also compute partial sums for DW and DB
 
         # fmt: off
+        meta = {"BLOCK_SIZE_N": ctx.BLOCK_SIZE_N,
+                "GROUP_SIZE_M": GROUP_SIZE_M,
+                "num_warps": ctx.num_warps}
+        if weight is None:
+            _layer_norm_no_affine_bwd[(M,)](dx, dy, y, var, y.stride(0), N, **meta)
+            return dx, None, None, None
+
         _layer_norm_bwd_dx_fused[(M,)](
             dx, dy, _dw, _db,
-            x_arg, weight, mean, var,
+            y, weight, bias, var,
             locks,
-            x_arg.stride(0),
+            y.stride(0),
             N,
-            BLOCK_SIZE_N=ctx.BLOCK_SIZE_N,
-            GROUP_SIZE_M=GROUP_SIZE_M,
-            num_warps=ctx.num_warps
+            **meta
         )
+
         # fmt: on
 
         def grid(meta):
@@ -294,5 +377,5 @@ def grid(meta):
         )
         # fmt: on
 
-        dx = dx.reshape_as(x)
+        dx = dx.reshape_as(dy)
         return dx, dw, db, None
diff --git a/xformers/triton/layer_norm.py b/xformers/triton/layer_norm.py
--- a/xformers/triton/layer_norm.py
+++ b/xformers/triton/layer_norm.py
@@ -28,10 +28,13 @@ class FusedLayerNorm(nn.Module):
 
     """
 
-    def __init__(self, normalized_shape, eps=1e-05):
+    def __init__(self, normalized_shape, affine=True, eps=1e-05):
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(normalized_shape))
-        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        if affine:
+            self.weight = nn.Parameter(torch.ones(normalized_shape))
+            self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        else:
+            self.weight = self.bias = None
         self.epsilon = eps
 
     def forward(self, x):


Test patch: