0cb595ad94
when using checkpoint_lvl=2, we all_gather_raw(x) without async_op=True. So we don't need to wait for handle. Just skip.
689 lines
27 KiB
Python
689 lines
27 KiB
Python
# Copyright (c) 2023, Tri Dao.
|
|
# Inspired by https://github.com/NVIDIA/apex/blob/master/apex/fused_dense/fused_dense.py
|
|
# We make it work with pytorch amp and with bfloat16.
|
|
# The TensorParallel linear modules are inspired by https://github.com/NVIDIA/apex/blob/master/apex/transformer/tensor_parallel/layers.py
|
|
from functools import partial
|
|
from typing import Optional
|
|
|
|
# import fused_dense_cuda # from apex
|
|
import fused_dense_lib as fused_dense_cuda
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch import Tensor
|
|
from torch.cuda.amp import custom_bwd, custom_fwd
|
|
from torch.distributed import ProcessGroup
|
|
|
|
from flash_attn.ops.activations import gelu_bwd, relu_bwd, sqrelu_bwd, sqrelu_fwd
|
|
from flash_attn.utils.distributed import (
|
|
all_gather_raw,
|
|
all_reduce,
|
|
all_reduce_raw,
|
|
reduce_scatter,
|
|
reduce_scatter_raw,
|
|
)
|
|
|
|
|
|
class FusedDenseFunc(torch.autograd.Function):
|
|
@staticmethod
|
|
@custom_fwd
|
|
def forward(
|
|
ctx, x, weight, bias, return_residual=False, process_group=None, sequence_parallel=True
|
|
):
|
|
"""
|
|
If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel
|
|
with sequence parallelism: we do an all_gather_raw of x before doing the matmul.
|
|
"""
|
|
ctx.compute_weight_gradient = weight.requires_grad
|
|
ctx.return_residual = return_residual
|
|
ctx.process_group = process_group
|
|
ctx.sequence_parallel = sequence_parallel
|
|
|
|
if torch.is_autocast_enabled():
|
|
x = x.to(dtype=torch.get_autocast_gpu_dtype())
|
|
x = x.contiguous()
|
|
if process_group is not None and sequence_parallel:
|
|
# We want to kick off the all_gather early, before weight dtype conversion
|
|
total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
|
|
else:
|
|
total_x = x
|
|
|
|
if torch.is_autocast_enabled():
|
|
weight = weight.to(dtype=torch.get_autocast_gpu_dtype())
|
|
bias = bias.to(dtype=torch.get_autocast_gpu_dtype()) if bias is not None else None
|
|
weight = weight.contiguous()
|
|
if process_group is not None and sequence_parallel:
|
|
handle_x.wait()
|
|
batch_shape, n = total_x.shape[:-1], total_x.shape[-1]
|
|
batch_dim = batch_shape.numel()
|
|
# https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174
|
|
if min(batch_dim, n, *weight.shape) > 65535 * 32:
|
|
raise RuntimeError("fused_dense only supports matrix dims <= 2M")
|
|
output = F.linear(total_x, weight, bias)
|
|
if ctx.compute_weight_gradient:
|
|
ctx.save_for_backward(x, weight)
|
|
else:
|
|
ctx.save_for_backward(weight)
|
|
return output if not return_residual else (output, x)
|
|
|
|
@staticmethod
|
|
@custom_bwd
|
|
def backward(ctx, grad_output, *args):
|
|
grad_output = grad_output.contiguous()
|
|
if ctx.return_residual:
|
|
(grad_input,) = args
|
|
grad_input = grad_input.contiguous()
|
|
process_group = ctx.process_group
|
|
sequence_parallel = ctx.sequence_parallel
|
|
if ctx.compute_weight_gradient:
|
|
x, weight = ctx.saved_tensors
|
|
if process_group is not None and sequence_parallel:
|
|
total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
|
|
else:
|
|
total_x = x
|
|
else:
|
|
(weight,) = ctx.saved_tensors
|
|
total_x = None
|
|
batch_shape = grad_output.shape[:-1]
|
|
batch_dim = batch_shape.numel()
|
|
grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
|
|
if ctx.needs_input_grad[0]:
|
|
if not ctx.return_residual:
|
|
grad_input = F.linear(grad_output, weight.t())
|
|
else:
|
|
grad_input = torch.addmm(
|
|
grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_output, weight
|
|
)
|
|
grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1])
|
|
if process_group is not None:
|
|
reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw
|
|
grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True)
|
|
else:
|
|
grad_input = None
|
|
if ctx.needs_input_grad[1]:
|
|
assert ctx.compute_weight_gradient
|
|
if process_group is not None and sequence_parallel:
|
|
handle_x.wait()
|
|
grad_weight, grad_bias = fused_dense_cuda.linear_bias_wgrad(
|
|
total_x.reshape(batch_dim, total_x.shape[-1]), grad_output, ctx.needs_input_grad[2]
|
|
)
|
|
else:
|
|
grad_weight = None
|
|
grad_bias = grad_output if ctx.needs_input_grad[2] else None
|
|
if process_group is not None and ctx.needs_input_grad[0]:
|
|
handle_grad_input.wait()
|
|
return grad_input, grad_weight, grad_bias, None, None, None
|
|
|
|
|
|
def fused_dense_func(
|
|
x: Tensor,
|
|
weight: Tensor,
|
|
bias: Optional[Tensor] = None,
|
|
return_residual: bool = False,
|
|
process_group: Optional[ProcessGroup] = None,
|
|
sequence_parallel: bool = True,
|
|
):
|
|
dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or (
|
|
x.dtype == torch.float32 and torch.is_autocast_enabled()
|
|
)
|
|
if x.is_cuda and weight.is_cuda and (bias is None or bias.is_cuda) and dtype_eligible:
|
|
return FusedDenseFunc.apply(
|
|
x, weight, bias, return_residual, process_group, sequence_parallel
|
|
)
|
|
else:
|
|
assert process_group is None
|
|
out = F.linear(x, weight, bias)
|
|
return out if not return_residual else (out, x)
|
|
|
|
|
|
class FusedDense(nn.Linear):
|
|
def __init__(
|
|
self,
|
|
in_features: int,
|
|
out_features: int,
|
|
bias: bool = True,
|
|
return_residual: bool = False,
|
|
device=None,
|
|
dtype=None,
|
|
) -> None:
|
|
super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype)
|
|
self.return_residual = return_residual
|
|
|
|
def forward(self, x, process_group=None):
|
|
"""
|
|
If process_group is not None, we're doing Tensor Parallel with sequence parallelism:
|
|
we do an all_gather of x before doing the matmul.
|
|
"""
|
|
return fused_dense_func(
|
|
x,
|
|
self.weight,
|
|
self.bias,
|
|
return_residual=self.return_residual,
|
|
process_group=process_group,
|
|
)
|
|
|
|
|
|
class ColumnParallelLinear(nn.Linear):
|
|
def __init__(
|
|
self,
|
|
in_features: int,
|
|
out_features: int,
|
|
process_group: ProcessGroup,
|
|
bias: bool = True,
|
|
sequence_parallel=True,
|
|
multiple_of=1,
|
|
device=None,
|
|
dtype=None,
|
|
) -> None:
|
|
world_size = torch.distributed.get_world_size(process_group)
|
|
if out_features % multiple_of:
|
|
raise ValueError(f"out_features ({out_features}) must be a multiple of {multiple_of}")
|
|
multiple = out_features // multiple_of
|
|
# We want to split @multiple across world_size, but it could be an uneven split
|
|
div = multiple // world_size
|
|
mod = multiple % world_size
|
|
# The first @mod ranks get @div + 1 copies, the rest get @div copies
|
|
local_multiple = div + int(torch.distributed.get_rank(process_group) < mod)
|
|
super().__init__(
|
|
in_features, local_multiple * multiple_of, bias=bias, device=device, dtype=dtype
|
|
)
|
|
self.process_group = process_group
|
|
self.sequence_parallel = sequence_parallel
|
|
|
|
def forward(self, x):
|
|
# If self.sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism:
|
|
# we do an all_gather of x before doing the matmul.
|
|
# If not, then the input is already gathered.
|
|
return fused_dense_func(
|
|
x,
|
|
self.weight,
|
|
self.bias,
|
|
process_group=self.process_group,
|
|
sequence_parallel=self.sequence_parallel,
|
|
)
|
|
|
|
|
|
class RowParallelLinear(nn.Linear):
|
|
def __init__(
|
|
self,
|
|
in_features: int,
|
|
out_features: int,
|
|
process_group: ProcessGroup,
|
|
bias: bool = True,
|
|
sequence_parallel=True,
|
|
multiple_of=1,
|
|
device=None,
|
|
dtype=None,
|
|
) -> None:
|
|
world_size = torch.distributed.get_world_size(process_group)
|
|
rank = torch.distributed.get_rank(process_group)
|
|
if in_features % multiple_of:
|
|
raise ValueError(f"in_features ({in_features}) must be a multiple of {multiple_of}")
|
|
multiple = in_features // multiple_of
|
|
# We want to split @multiple across world_size, but it could be an uneven split
|
|
div = multiple // world_size
|
|
mod = multiple % world_size
|
|
# The first @mod ranks get @div + 1 copies, the rest get @div copies
|
|
local_multiple = div + int(torch.distributed.get_rank(process_group) < mod)
|
|
# Only rank 0 will have bias
|
|
super().__init__(
|
|
local_multiple * multiple_of,
|
|
out_features,
|
|
bias=bias and rank == 0,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
self.process_group = process_group
|
|
self.sequence_parallel = sequence_parallel
|
|
|
|
def forward(self, x):
|
|
"""
|
|
We're doing Tensor Parallel with sequence parallelism: we do the matmul and then
|
|
a reduce_scatter of the result.
|
|
"""
|
|
out = fused_dense_func(x, self.weight, self.bias)
|
|
reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
|
|
return reduce_fn(out, self.process_group)
|
|
|
|
|
|
class FusedMLPFunc(torch.autograd.Function):
|
|
@staticmethod
|
|
@custom_fwd
|
|
def forward(
|
|
ctx,
|
|
x,
|
|
weight1,
|
|
bias1,
|
|
weight2,
|
|
bias2,
|
|
activation="gelu_approx",
|
|
save_pre_act=True,
|
|
return_residual=False,
|
|
checkpoint_lvl=0,
|
|
heuristic=0,
|
|
process_group=None,
|
|
sequence_parallel=True,
|
|
):
|
|
"""
|
|
If process_group is not None and sequence_parallel=True, we're doing Tensor Parallel
|
|
with sequence parallelism: we do an all_gather of x before doing the matmul.
|
|
If sequence_parallel=False, then the input is already gathered.
|
|
|
|
checkpoint_lvl:
|
|
0: no recomputation in the bwd
|
|
1: recompute gelu_out / relu_out in the bwd
|
|
2: recompute pre_act and gelu_out / relu_out in the bwd
|
|
"""
|
|
assert -1 <= heuristic <= 4
|
|
assert activation in ["gelu_approx", "relu", "sqrelu"]
|
|
if activation == "sqrelu":
|
|
assert heuristic == -1
|
|
if not save_pre_act:
|
|
checkpoint_lvl = 2
|
|
assert checkpoint_lvl in [0, 1, 2]
|
|
ctx.return_residual = return_residual
|
|
ctx.process_group = process_group
|
|
ctx.sequence_parallel = sequence_parallel
|
|
ctx.checkpoint_lvl = checkpoint_lvl
|
|
ctx.activation = activation
|
|
ctx.heuristic = heuristic
|
|
|
|
if torch.is_autocast_enabled():
|
|
x = x.to(dtype=torch.get_autocast_gpu_dtype())
|
|
x = x.contiguous()
|
|
if process_group is not None and sequence_parallel:
|
|
# We want to kick off the all_gather early, before weight dtype conversion
|
|
total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
|
|
else:
|
|
total_x = x
|
|
|
|
if torch.is_autocast_enabled():
|
|
dtype = torch.get_autocast_gpu_dtype()
|
|
weight1, weight2 = [a.to(dtype=dtype) for a in [weight1, weight2]]
|
|
bias1 = bias1.to(dtype=dtype) if bias1 is not None else None
|
|
bias2 = bias2.to(dtype=dtype) if bias2 is not None else None
|
|
weight1 = weight1.contiguous()
|
|
bias1 = bias1.contiguous() if bias1 is not None else None
|
|
weight2 = weight2.contiguous()
|
|
bias2 = bias2.contiguous() if bias2 is not None else None
|
|
if process_group is not None and sequence_parallel:
|
|
handle_x.wait()
|
|
batch_shape, n = total_x.shape[:-1], total_x.shape[-1]
|
|
batch_dim = batch_shape.numel()
|
|
# https://github.com/pytorch/pytorch/blob/5b51849b48a7dbccd297286cc0110def4706f9e7/aten/src/ATen/native/cuda/Blas.cpp#L174
|
|
if min(batch_dim, n, *weight1.shape, *weight2.shape) > 65535 * 32:
|
|
raise RuntimeError("fused_dense only supports matrix dims <= 2M")
|
|
if heuristic == -1:
|
|
pre_act = F.linear(total_x, weight1, bias1)
|
|
activation_fn = (
|
|
partial(F.gelu, approximate="tanh")
|
|
if activation == "gelu_approx"
|
|
else (sqrelu_fwd if activation == "sqrelu" else F.relu)
|
|
)
|
|
with torch.jit.fuser("fuser2"):
|
|
output1 = activation_fn(pre_act)
|
|
# This is before adding bias1
|
|
# pre_act = F.linear(total_x.reshape(batch_dim, n), weight1)
|
|
# with torch.jit.fuser('fuser2'):
|
|
# output1 = bias_gelu(pre_act, bias1)
|
|
else:
|
|
is_gelu = activation == "gelu_approx"
|
|
output1, *rest = fused_dense_cuda.linear_act_forward(
|
|
total_x.reshape(batch_dim, n), weight1, bias1, is_gelu, save_pre_act, heuristic
|
|
)
|
|
if save_pre_act:
|
|
pre_act = rest[0]
|
|
output2 = F.linear(output1, weight2, bias2)
|
|
if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"):
|
|
# For RELU the pre_act is very small (just a bit-mask) so we just save it
|
|
ctx.save_for_backward(x, weight1, weight2, pre_act, output1)
|
|
elif checkpoint_lvl == 1:
|
|
ctx.save_for_backward(x, weight1, weight2, pre_act)
|
|
elif checkpoint_lvl == 2:
|
|
ctx.save_for_backward(x, weight1, weight2, bias1)
|
|
output2 = output2.reshape(*batch_shape, output2.shape[-1])
|
|
return output2 if not return_residual else (output2, x)
|
|
|
|
@staticmethod
|
|
@custom_bwd
|
|
def backward(ctx, grad_output, *args):
|
|
grad_output = grad_output.contiguous()
|
|
checkpoint_lvl = ctx.checkpoint_lvl
|
|
activation = ctx.activation
|
|
activation_fn = (
|
|
partial(F.gelu, approximate="tanh")
|
|
if activation == "gelu_approx"
|
|
else (sqrelu_fwd if activation == "sqrelu" else F.relu)
|
|
)
|
|
if ctx.return_residual:
|
|
(grad_input,) = args
|
|
grad_input = grad_input.contiguous()
|
|
process_group = ctx.process_group
|
|
sequence_parallel = ctx.sequence_parallel
|
|
x, weight1, weight2, *rest = ctx.saved_tensors
|
|
if process_group is None or not sequence_parallel:
|
|
total_x = x
|
|
batch_shape = grad_output.shape[:-1]
|
|
batch_dim = batch_shape.numel()
|
|
if checkpoint_lvl in [0, 1]:
|
|
if process_group is not None and sequence_parallel:
|
|
total_x, handle_x = all_gather_raw(x, process_group, async_op=True)
|
|
if checkpoint_lvl == 0 or (checkpoint_lvl == 1 and activation == "relu"):
|
|
pre_act, output1 = rest
|
|
elif checkpoint_lvl == 1:
|
|
(pre_act,) = rest
|
|
with torch.jit.fuser("fuser2"):
|
|
output1 = activation_fn(pre_act)
|
|
elif checkpoint_lvl == 2:
|
|
(bias1,) = rest
|
|
if process_group is not None and sequence_parallel:
|
|
total_x, _ = all_gather_raw(x, process_group)
|
|
if ctx.heuristic == -1:
|
|
pre_act = F.linear(total_x, weight1, bias1)
|
|
with torch.jit.fuser("fuser2"):
|
|
output1 = activation_fn(pre_act)
|
|
else:
|
|
output1, pre_act = fused_dense_cuda.linear_act_forward(
|
|
total_x.reshape(batch_dim, total_x.shape[-1]),
|
|
weight1,
|
|
bias1,
|
|
activation == "gelu_approx",
|
|
True,
|
|
ctx.heuristic,
|
|
)
|
|
|
|
grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
|
|
output1 = output1.reshape(batch_dim, output1.shape[-1])
|
|
pre_act = pre_act.reshape(batch_dim, pre_act.shape[-1])
|
|
if ctx.needs_input_grad[3]:
|
|
grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(
|
|
output1, grad_output, ctx.needs_input_grad[4]
|
|
)
|
|
else:
|
|
grad_weight2 = None
|
|
grad_bias2 = grad_output if ctx.needs_input_grad[4] else None
|
|
if ctx.heuristic == -1:
|
|
# grad_pre_act = matmul_dgelu(grad_output, weight2, pre_act)
|
|
grad_output1 = F.linear(grad_output, weight2.t())
|
|
activation_grad_fn = (
|
|
gelu_bwd
|
|
if activation == "gelu_approx"
|
|
else (sqrelu_bwd if activation == "sqrelu" else relu_bwd)
|
|
)
|
|
with torch.jit.fuser("fuser2"):
|
|
grad_pre_act = activation_grad_fn(grad_output1, pre_act)
|
|
else:
|
|
# The cublasLt epilogue has to compute both gelu/relu grad and bias grad, we can't
|
|
# just compute gelu/relu grad
|
|
grad_pre_act, grad_bias1 = fused_dense_cuda.bias_act_linear_dgrad_bgrad(
|
|
weight2, grad_output, pre_act, activation == "gelu_approx", ctx.heuristic
|
|
)
|
|
if not ctx.needs_input_grad[2]:
|
|
grad_bias1 = None
|
|
if ctx.needs_input_grad[0]:
|
|
if not ctx.return_residual:
|
|
grad_input = F.linear(grad_pre_act, weight1.t())
|
|
else:
|
|
grad_input = torch.addmm(
|
|
grad_input.reshape(batch_dim, grad_input.shape[-1]), grad_pre_act, weight1
|
|
)
|
|
grad_input = grad_input.reshape(*batch_shape, grad_input.shape[-1])
|
|
if process_group is not None:
|
|
reduce_fn = reduce_scatter_raw if sequence_parallel else all_reduce_raw
|
|
grad_input, handle_grad_input = reduce_fn(grad_input, process_group, async_op=True)
|
|
else:
|
|
grad_input = None
|
|
if ctx.heuristic == -1:
|
|
if ctx.needs_input_grad[1]:
|
|
if process_group is not None and sequence_parallel and checkpoint_lvl != 2:
|
|
handle_x.wait()
|
|
grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_wgrad(
|
|
total_x.reshape(batch_dim, total_x.shape[-1]),
|
|
grad_pre_act,
|
|
ctx.needs_input_grad[2],
|
|
)
|
|
else:
|
|
grad_weight1 = None
|
|
grad_bias1 = grad_pre_act if ctx.needs_input_grad[2] else None
|
|
else:
|
|
if ctx.needs_input_grad[1]:
|
|
if process_group is not None and sequence_parallel and checkpoint_lvl != 2:
|
|
handle_x.wait()
|
|
grad_weight1 = F.linear(
|
|
grad_pre_act.t(), total_x.reshape(batch_dim, total_x.shape[-1]).t()
|
|
)
|
|
else:
|
|
grad_weight1 = None
|
|
if process_group is not None and ctx.needs_input_grad[0]:
|
|
handle_grad_input.wait()
|
|
return (
|
|
grad_input,
|
|
grad_weight1,
|
|
grad_bias1,
|
|
grad_weight2,
|
|
grad_bias2,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
|
|
|
|
def fused_mlp_func(
|
|
x: Tensor,
|
|
weight1: Tensor,
|
|
weight2: Tensor,
|
|
bias1: Optional[Tensor] = None,
|
|
bias2: Optional[Tensor] = None,
|
|
activation: str = "gelu_approx",
|
|
save_pre_act: bool = True,
|
|
return_residual: bool = False,
|
|
checkpoint_lvl: int = 0,
|
|
heuristic: int = 0,
|
|
process_group: Optional[ProcessGroup] = None,
|
|
sequence_parallel: bool = True,
|
|
):
|
|
assert activation in ["gelu_approx", "relu", "sqrelu"]
|
|
dtype_eligible = x.dtype in [torch.float16, torch.bfloat16] or (
|
|
x.dtype == torch.float32 and torch.is_autocast_enabled()
|
|
)
|
|
# If we save pre-activation, dimension must be divisible by 128 (relu) or 8 (gelu)
|
|
dim_eligible = not save_pre_act or (x.shape[-1] % (128 if activation == "relu" else 8) == 0)
|
|
if (
|
|
x.is_cuda
|
|
and weight1.is_cuda
|
|
and weight2.is_cuda
|
|
and (bias1 is None or bias1.is_cuda)
|
|
and (bias2 is None or bias2.is_cuda)
|
|
and dtype_eligible
|
|
and dim_eligible
|
|
):
|
|
return FusedMLPFunc.apply(
|
|
x,
|
|
weight1,
|
|
bias1,
|
|
weight2,
|
|
bias2,
|
|
activation,
|
|
save_pre_act,
|
|
return_residual,
|
|
checkpoint_lvl,
|
|
heuristic,
|
|
process_group,
|
|
sequence_parallel,
|
|
)
|
|
else:
|
|
assert process_group is None
|
|
pre_act = F.linear(x, weight1, bias1)
|
|
activation_fn = (
|
|
partial(F.gelu, approximate="tanh")
|
|
if activation == "gelu_approx"
|
|
else partial(F.relu, inplace=True)
|
|
)
|
|
output1 = activation_fn(pre_act)
|
|
output2 = F.linear(output1, weight2, bias2)
|
|
return output2 if not return_residual else (output2, x)
|
|
|
|
|
|
class FusedMLP(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_features,
|
|
hidden_features=None,
|
|
out_features=None,
|
|
bias1=True,
|
|
bias2=True,
|
|
activation="gelu_approx",
|
|
return_residual=False,
|
|
checkpoint_lvl=0,
|
|
heuristic="auto",
|
|
device=None,
|
|
dtype=None,
|
|
):
|
|
"""
|
|
If process_group is not None, we're doing Tensor Parallel with sequence parallelism:
|
|
we do an all_gather of x before doing the matmul, gelu, then matmul.
|
|
Finally we do a reduce_scatter of the output.
|
|
|
|
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 pre_act 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
|
|
'auto': heuristic will be picked automatically:
|
|
For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf.
|
|
For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16.
|
|
For H100, we set heuristic=-1 for both fp16 and bf16 as the fused cuBlasLt implementation
|
|
is slower than the unfused version.
|
|
return_residual: whether to return the input x along with the output. This is for
|
|
performance reason: for post-norm architecture, returning the input allows us
|
|
to fuse the backward of nn.Linear with the residual connection.
|
|
"""
|
|
assert checkpoint_lvl in [0, 1, 2]
|
|
assert activation in ["gelu_approx", "relu", "sqrelu"]
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
|
super().__init__()
|
|
out_features = out_features or in_features
|
|
hidden_features = hidden_features or in_features * 4
|
|
self.activation = activation
|
|
self.return_residual = return_residual
|
|
self.checkpoint_lvl = checkpoint_lvl
|
|
self.heuristic = heuristic if activation != "sqrelu" else -1
|
|
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs)
|
|
self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs)
|
|
|
|
def forward(self, x, process_group=None):
|
|
dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype()
|
|
if self.heuristic == "auto":
|
|
if self.activation == "gelu_approx":
|
|
if torch.cuda.get_device_capability("cuda") == (9, 0):
|
|
heuristic = -1
|
|
else:
|
|
cuda_ver = tuple(map(int, torch.version.cuda.split(".")))
|
|
heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1)
|
|
else:
|
|
heuristic = 0
|
|
else:
|
|
heuristic = self.heuristic
|
|
out = fused_mlp_func(
|
|
x,
|
|
self.fc1.weight,
|
|
self.fc2.weight,
|
|
self.fc1.bias,
|
|
self.fc2.bias,
|
|
activation=self.activation,
|
|
save_pre_act=self.training,
|
|
return_residual=self.return_residual,
|
|
checkpoint_lvl=self.checkpoint_lvl,
|
|
heuristic=heuristic,
|
|
process_group=process_group,
|
|
)
|
|
if self.return_residual:
|
|
out, x = out
|
|
if process_group is not None:
|
|
out = reduce_scatter(out, process_group)
|
|
return out if not self.return_residual else (out, x)
|
|
|
|
|
|
class ParallelFusedMLP(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_features,
|
|
hidden_features=None,
|
|
out_features=None,
|
|
activation="gelu_approx",
|
|
process_group: ProcessGroup = None,
|
|
bias1=True,
|
|
bias2=True,
|
|
sequence_parallel=True,
|
|
checkpoint_lvl=0,
|
|
heuristic="auto",
|
|
device=None,
|
|
dtype=None,
|
|
):
|
|
"""
|
|
process_group is required. We're doing Tensor Parallel with sequence parallelism:
|
|
we do an all_gather of x before doing the matmul, gelu, then matmul.
|
|
Finally we do a reduce_scatter of the output.
|
|
|
|
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 pre_act 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
|
|
'auto': heuristic will be picked automatically:
|
|
For CUDA >= 11.8, we set heuristic=0 for both fp16 and bf16 for best perf.
|
|
For CUDA <= 11.7, we set heuristic=1 for fp16 and heuristic=-1 for bf16.
|
|
"""
|
|
assert checkpoint_lvl in [0, 1, 2]
|
|
assert activation in ["gelu_approx", "relu", "sqrelu"]
|
|
assert process_group is not None
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
|
super().__init__()
|
|
out_features = out_features or in_features
|
|
hidden_features = hidden_features or in_features * 4
|
|
self.activation = activation
|
|
self.process_group = process_group
|
|
self.sequence_parallel = sequence_parallel
|
|
self.checkpoint_lvl = checkpoint_lvl
|
|
self.heuristic = heuristic if activation != "sqrelu" else -1
|
|
self.fc1 = ColumnParallelLinear(
|
|
in_features, hidden_features, process_group, bias=bias1, **factory_kwargs
|
|
)
|
|
self.fc2 = RowParallelLinear(
|
|
hidden_features, out_features, process_group, bias=bias2, **factory_kwargs
|
|
)
|
|
|
|
def forward(self, x):
|
|
dtype = x.dtype if not torch.is_autocast_enabled() else torch.get_autocast_gpu_dtype()
|
|
if self.heuristic == "auto":
|
|
if self.activation == "gelu_approx":
|
|
cuda_ver = tuple(map(int, torch.version.cuda.split(".")))
|
|
heuristic = 0 if cuda_ver >= (11, 8) else (1 if dtype == torch.float16 else -1)
|
|
else:
|
|
heuristic = 0
|
|
else:
|
|
heuristic = self.heuristic
|
|
out = fused_mlp_func(
|
|
x,
|
|
self.fc1.weight,
|
|
self.fc2.weight,
|
|
self.fc1.bias,
|
|
self.fc2.bias,
|
|
activation=self.activation,
|
|
save_pre_act=self.training,
|
|
checkpoint_lvl=self.checkpoint_lvl,
|
|
heuristic=heuristic,
|
|
process_group=self.process_group,
|
|
sequence_parallel=self.sequence_parallel,
|
|
)
|
|
reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce
|
|
return reduce_fn(out, self.process_group)
|