fused_dense.py

# Adapted from https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py
# We make it work with pytorch amp and with bfloat16.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.cuda.amp import custom_bwd, custom_fwd

# import fused_dense_cuda  # from apex
import fused_dense_lib as fused_dense_cuda
# from src.ops.triton.triton_matmul import matmul_dgelu
from flash_attn.ops.gelu_activation import gelu_bwd
# from src.ops.gelu_activation import gelu_bwd, bias_gelu, bias_gelu_back


# implements fused GEMM+bias in forward pass using mlp_cuda from apex
class FusedDenseFuncTD(torch.autograd.Function):

    @staticmethod
    @custom_fwd
    def forward(ctx, x, weight, bias):
        if torch.is_autocast_enabled():
            dtype = torch.get_autocast_gpu_dtype()
            x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]]
        x = x.contiguous()
        weight = weight.contiguous()
        bias = bias.contiguous()
        ctx.save_for_backward(x, weight)
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        assert batch_dim <= 64 * 1024, 'fused_dense only supports dimension at most 64k'
        output = fused_dense_cuda.linear_bias_forward(x.reshape(batch_dim, n), weight, bias)
        return output.reshape(*batch_shape, output.shape[-1])

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output):
        grad_output = grad_output.contiguous()
        x, weight = ctx.saved_tensors
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        if ctx.needs_input_grad[0]:
            grad_input, grad_weight, grad_bias = fused_dense_cuda.linear_bias_backward(
                x.reshape(batch_dim, n), weight, grad_output.reshape(batch_dim, grad_output.shape[-1])
            )
            grad_input = grad_input.reshape_as(x)
        else:
            grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad(
                x.reshape(batch_dim, n), grad_output.reshape(batch_dim, grad_output.shape[-1])
            )
            grad_input = None
        # print((grad_bias - grad_output.view(-1, grad_output.shape[-1]).sum(dim=0)).abs().max())
        return grad_input, grad_weight, grad_bias
        # grad_input, grad_weight = None, None
        # grad_output_reshaped = grad_output.reshape(batch_dim, grad_output.shape[-1])
        # if ctx.needs_input_grad[0]:
        #     grad_input = (grad_output_reshaped @ weight.conj()).reshape(*batch_shape, n)
        # if ctx.needs_input_grad[1]:
        #     grad_weight = grad_output_reshaped.t() @ x.conj().reshape(batch_dim, n)
        # # We don't need to compute grad_bias explicitly, when we return grad_out Pytorch
        # # will sum over the batch dimension to get grad_bias.
        # return grad_input, grad_weight, grad_output


fused_dense_function_td = FusedDenseFuncTD.apply


class FusedDenseTD(nn.Linear):

    def __init__(self, in_features: int, out_features: int, bias: bool = True,
                 device=None, dtype=None) -> None:
        super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype)

    def forward(self, x):
        if x.is_cuda and self.bias is not None:
            return fused_dense_function_td(x, self.weight, self.bias)
        else:
            return F.linear(x, self.weight, self.bias)


class FusedDenseResidualFunc(torch.autograd.Function):

    @staticmethod
    @custom_fwd
    def forward(ctx, x, weight, bias):
        if torch.is_autocast_enabled():
            dtype = torch.get_autocast_gpu_dtype()
            x, weight, bias = [a.to(dtype=dtype) for a in [x, weight, bias]]
        x = x.contiguous()
        x = x.contiguous()
        weight = weight.contiguous()
        bias = bias.contiguous()
        ctx.save_for_backward(x, weight)
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        assert batch_dim <= 64 * 1024, 'fused_dense only supports dimension at most 64k'
        output = fused_dense_cuda.linear_bias_forward(x.reshape(batch_dim, n), weight, bias)
        return output.reshape(*batch_shape, output.shape[-1]), x

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output, grad_input):
        grad_output = grad_output.contiguous()
        grad_input = grad_input.contiguous()
        x, weight = ctx.saved_tensors
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        grad_input, grad_weight, grad_bias = fused_dense_cuda.linear_bias_residual_backward(
            x.reshape(batch_dim, n), weight, grad_output.reshape(batch_dim, grad_output.shape[-1]),
            grad_input.reshape(batch_dim, n)
        )
        return grad_input.reshape_as(x), grad_weight, grad_bias


fused_dense_residual_function = FusedDenseResidualFunc.apply


class FusedDenseResidual(nn.Linear):
    """Similar to FusedDense, but we return both the output and the input.
    This is so that in the backward pass, we can combine the input gradient from the residual branch
    with the input gradient from the matrix multiply, without having to do a separate addition.
    """

    def forward(self, x):
        if x.is_cuda and self.bias is not None:
            return fused_dense_residual_function(x, self.weight, self.bias)
        else:
            return F.linear(x, self.weight, self.bias), x


class FusedDenseGeluDenseFuncTD(torch.autograd.Function):

    @staticmethod
    @custom_fwd
    def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0, heuristic=0):
        """checkpoint_lvl:
        0: no recomputation in the bwd
        1: recompute gelu_out in the bwd
        2: recompute gelu_in and gelu_out in the bwd
        """
        assert -1 <= heuristic <= 4
        if torch.is_autocast_enabled():
            dtype = torch.get_autocast_gpu_dtype()
            x, weight1, bias1, weight2, bias2 = [a.to(dtype=dtype)
                                                 for a in [x, weight1, bias1, weight2, bias2]]
        assert checkpoint_lvl in [0, 1, 2]
        x = x.contiguous()
        weight1 = weight1.contiguous()
        bias1 = bias1.contiguous()
        weight2 = weight2.contiguous()
        bias2 = bias2.contiguous()
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        assert batch_dim <= 64 * 1024, 'fused_dense only supports dimension at most 64k'
        # output1, output2, gelu_in = fused_dense_cuda.linear_gelu_linear_forward(
        #     x.reshape(batch_dim, n), weight1, bias1, weight2, bias2
        # )
        if heuristic == -1:
            gelu_in = fused_dense_cuda.linear_bias_forward(x.reshape(batch_dim, n), weight1, bias1)
            output1 = F.gelu(gelu_in, approximate='tanh')
            # gelu_in = F.linear(x.reshape(batch_dim, n), weight1)  # This is before adding bias1
            # with torch.jit.fuser('fuser2'):
            #     output1 = bias_gelu(gelu_in, bias1)
        else:
            save_gelu_in = checkpoint_lvl != 2
            output1, *rest = fused_dense_cuda.linear_gelu_forward(x.reshape(batch_dim, n), weight1,
                                                                  bias1, save_gelu_in, heuristic)
            if save_gelu_in:
                gelu_in = rest[0]
        output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2)
        ctx.checkpoint_lvl = checkpoint_lvl
        ctx.heuristic = heuristic
        if checkpoint_lvl == 0:
            ctx.save_for_backward(x, weight1, bias1, weight2, gelu_in, output1)
        elif checkpoint_lvl == 1:
            ctx.save_for_backward(x, weight1, bias1, weight2, gelu_in)
        elif checkpoint_lvl == 2:
            ctx.save_for_backward(x, weight1, bias1, weight2)
        return output2.reshape(*batch_shape, output2.shape[-1])

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output):
        grad_output = grad_output.contiguous()
        checkpoint_lvl = ctx.checkpoint_lvl
        x, weight1, bias1, weight2, *rest = ctx.saved_tensors
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        if checkpoint_lvl == 0:
            gelu_in, output1 = rest
        elif checkpoint_lvl == 1:
            gelu_in, = rest
            output1 = F.gelu(gelu_in, approximate='tanh')
        elif checkpoint_lvl == 2:
            # bias1, = rest
            if ctx.heuristic == -1:
                gelu_in = fused_dense_cuda.linear_bias_forward(x.reshape(batch_dim, n), weight1, bias1)
                output1 = F.gelu(gelu_in, approximate='tanh')
            else:
                output1, gelu_in = fused_dense_cuda.linear_gelu_forward(x.reshape(batch_dim, n),
                                                                        weight1, bias1, True, ctx.heuristic)

        if ctx.heuristic == -1:
            grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
            # grad_output1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_backward(output1, weight2, grad_output)
            grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output)
            # grad_gelu = matmul_dgelu(grad_output, weight2, gelu_in)
            grad_output1 = grad_output @ weight2
            with torch.jit.fuser('fuser2'):
                grad_gelu = gelu_bwd(grad_output1, gelu_in)
            grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward(
                x.reshape(batch_dim, n), weight1, grad_gelu
            )
            # with torch.jit.fuser('fuser2'):
            #     grad_gelu, grad_bias1 = bias_gelu_back(grad_output1, gelu_in, bias1)
            # grad_input = grad_gelu @ weight1
            # grad_weight1 = grad_gelu.reshape(batch_dim, -1).T @ x.reshape(batch_dim, n)
            # grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward(
            #     x.reshape(batch_dim, n), weight1, grad_gelu
            # )
        else:
            grad_input, grad_weight1, grad_bias1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_gelu_linear_backward(
                x.reshape(batch_dim, n), gelu_in, output1, weight1, weight2,
                grad_output.reshape(batch_dim, grad_output.shape[-1]),
                ctx.heuristic
            )
        # grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
        # # grad_output1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_backward(output1, weight2, grad_output)
        # grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output)
        # grad_gelu = matmul_dgelu(grad_output, weight2, gelu_in)
        # grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward(
        #     x.reshape(batch_dim, n), weight1, grad_gelu
        # )
        return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None, None


fused_dense_gelu_dense_function_td = FusedDenseGeluDenseFuncTD.apply


class FusedDenseGeluDenseTD(nn.Module):

    def __init__(self, in_features, intermediate_features, out_features=None, bias=True,
                 checkpoint_lvl=0, heuristic=0, device=None, dtype=None):
        """
        checkpoint_lvl (increasing lvl means slower but more memory saving):
            0: no recomputation in the bwd
            1: recompute gelu_out in the bwd
            2: recompute gelu_in and gelu_out in the bwd
        heuristic:
            -1: don't fuse gemm + gelu (separate kernel)
            0..4: use this heuristic for the algo section in the fused gemm + gelu
        """
        assert checkpoint_lvl in [0, 1, 2]
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        if out_features is None:
            out_features = in_features
        assert bias == True, "DenseGeluDense module without bias is currently not supported"
        self.checkpoint_lvl = checkpoint_lvl
        self.heuristic = heuristic
        self.fc1 = nn.Linear(in_features, intermediate_features, bias=bias, **factory_kwargs)
        self.fc2 = nn.Linear(intermediate_features, out_features, bias=bias, **factory_kwargs)

    def forward(self, x):
        return fused_dense_gelu_dense_function_td(x, self.fc1.weight, self.fc1.bias,
                                                  self.fc2.weight, self.fc2.bias,
                                                  self.checkpoint_lvl, self.heuristic)


class FusedDenseResGeluDenseFunc(torch.autograd.Function):

    @staticmethod
    @custom_fwd
    def forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl=0, heuristic=0):
        """checkpoint_lvl:
        0: no recomputation in the bwd
        1: recompute gelu_out in the bwd
        2: recompute gelu_in and gelu_out in the bwd
        """
        assert -1 <= heuristic <= 4
        if torch.is_autocast_enabled():
            dtype = torch.get_autocast_gpu_dtype()
            x, weight1, bias1, weight2, bias2 = [a.to(dtype=dtype)
                                                 for a in [x, weight1, bias1, weight2, bias2]]
        assert checkpoint_lvl in [0, 1, 2]
        x = x.contiguous()
        weight1 = weight1.contiguous()
        bias1 = bias1.contiguous()
        weight2 = weight2.contiguous()
        bias2 = bias2.contiguous()
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        assert batch_dim <= 64 * 1024, 'fused_dense only supports dimension at most 64k'
        # output1, output2, gelu_in = fused_dense_cuda.linear_gelu_linear_forward(
        #     x.reshape(batch_dim, n), weight1, bias1, weight2, bias2
        # )
        # gelu_in = fused_dense_cuda.linear_bias_forward(x.reshape(batch_dim, n), weight1, bias1)
        # output1 = F.gelu(gelu_in, approximate='tanh')
        save_gelu_in = checkpoint_lvl != 2
        output1, *rest = fused_dense_cuda.linear_gelu_forward(x.reshape(batch_dim, n), weight1,
                                                              bias1, save_gelu_in, heuristic)
        if save_gelu_in:
            gelu_in = rest[0]
        output2 = fused_dense_cuda.linear_bias_forward(output1, weight2, bias2)
        ctx.checkpoint_lvl = checkpoint_lvl
        ctx.heuristic = heuristic
        if checkpoint_lvl == 0:
            ctx.save_for_backward(x, weight1, weight2, gelu_in, output1)
        elif checkpoint_lvl == 1:
            ctx.save_for_backward(x, weight1, weight2, gelu_in)
        elif checkpoint_lvl == 2:
            ctx.save_for_backward(x, weight1, weight2, bias1)
        return output2.reshape(*batch_shape, output2.shape[-1]), x

    @staticmethod
    @custom_bwd
    def backward(ctx, grad_output, grad_input):
        grad_output = grad_output.contiguous()
        grad_input = grad_input.contiguous()
        checkpoint_lvl = ctx.checkpoint_lvl
        x, weight1, weight2, *rest = ctx.saved_tensors
        batch_shape, n = x.shape[:-1], x.shape[-1]
        batch_dim = batch_shape.numel()
        if checkpoint_lvl == 0:
            gelu_in, output1 = rest
        elif checkpoint_lvl == 1:
            gelu_in, = rest
            output1 = F.gelu(gelu_in, approximate='tanh')
        elif checkpoint_lvl == 2:
            bias1, = rest
            output1, gelu_in = fused_dense_cuda.linear_gelu_forward(x.reshape(batch_dim, n),
                                                                    weight1, bias1, True, ctx.heuristic)
        grad_input, grad_weight1, grad_bias1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_residual_gelu_linear_backward(
            x.reshape(batch_dim, n), gelu_in, output1, weight1, weight2,
            grad_output.reshape(batch_dim, grad_output.shape[-1]),
            grad_input.reshape(batch_dim, n),
            ctx.heuristic
        )
        # grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
        # # grad_output1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_backward(output1, weight2, grad_output)
        # grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output)
        # grad_gelu = matmul_dgelu(grad_output, weight2, gelu_in)
        # grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_residual_backward(
        #     x.reshape(batch_dim, n), weight1, grad_gelu,
        #     grad_input.reshape(batch_dim, n)
        # )
        return grad_input.reshape_as(x), grad_weight1, grad_bias1, grad_weight2, grad_bias2, None, None


fused_dense_res_gelu_dense_function_td = FusedDenseResGeluDenseFunc.apply


class FusedDenseResGeluDense(FusedDenseGeluDenseTD):

    def forward(self, x):
        return fused_dense_res_gelu_dense_function_td(x, self.fc1.weight, self.fc1.bias,
                                                      self.fc2.weight, self.fc2.bias,
                                                      self.checkpoint_lvl, False, self.heuristic)