From 570e30a6bc5a1fb8ce9a51e498b37726154cd98c Mon Sep 17 00:00:00 2001 From: yingtongxiong <974106207@qq.com> Date: Mon, 31 Jul 2023 20:28:40 +0800 Subject: [PATCH] support float32 training --- internlm/initialize/launch.py | 10 +- internlm/model/linear.py | 59 +++++++-- internlm/model/multi_head_attention.py | 3 +- internlm/model/utils.py | 175 +++++++++++++++++++++++++ 4 files changed, 236 insertions(+), 11 deletions(-) diff --git a/internlm/initialize/launch.py b/internlm/initialize/launch.py index d2d61b1..4c7ea7a 100644 --- a/internlm/initialize/launch.py +++ b/internlm/initialize/launch.py @@ -154,8 +154,16 @@ def args_sanity_check(): gpc.config.model.dtype = torch.bfloat16 elif gpc.config.model.dtype in ("torch.float16", "torch.half"): gpc.config.model.dtype = torch.float16 + elif gpc.config.model.dtype == "torch.float32": + assert gpc.config.model.use_flash_attn == False, "when using float32, the use_flash_attn must be False" + gpc.config.model.dtype = torch.float32 + elif gpc.config.model.dtype == "torch.tf32": + assert gpc.config.model.use_flash_attn == False, "when using tf32, the use_flash_attn must be False" + torch.backends.cudnn.allow_tf32 = True + torch.backends.cuda.matmul.allow_tf32 = True + gpc.config.model.dtype = torch.float32 else: - assert gpc.config.model.dtype in ["torch.float16", "torch.half", "torch.bfloat16"] + assert gpc.config.model.dtype in ["torch.float16", "torch.half", "torch.bfloat16", "torch.float32", "torch.tf32"] if gpc.is_rank_for_log(): logger.info("+" * 15 + " Model Info " + "+" * 15) # pylint: disable=W1201 diff --git a/internlm/model/linear.py b/internlm/model/linear.py index 88129af..918fff6 100644 --- a/internlm/model/linear.py +++ b/internlm/model/linear.py @@ -5,16 +5,12 @@ from typing import Optional import torch import torch.nn.functional as F -from flash_attn.ops.fused_dense import ( - ColumnParallelLinear, - RowParallelLinear, - fused_dense_func, -) +from torch.distributed import ProcessGroup from torch import nn from internlm.core.context import IS_TENSOR_PARALLEL, ParallelMode from internlm.core.context import global_context as gpc - +from internlm.model.utils import fused_dense_func_torch, reduce_scatter, all_reduce class ScaleColumnParallelLinear(nn.Linear): """ @@ -61,7 +57,7 @@ class ScaleColumnParallelLinear(nn.Linear): weight = self.weight * self.weight_scale + (1 - self.weight_scale) * self.weight.detach() else: weight = self.weight - return fused_dense_func( + return fused_dense_func_torch( input, weight, self.bias, process_group=self.process_group, sequence_parallel=self.sequence_parallel ) @@ -107,11 +103,58 @@ class RewardModelLinear(ScaleColumnParallelLinear): weight = self.weight * self.weight_scale + (1 - self.weight_scale) * self.weight.detach() else: weight = self.weight - return fused_dense_func( + return fused_dense_func_torch( input, weight, self.bias, process_group=self.process_group, sequence_parallel=self.sequence_parallel ) +class ColumnParallelLinear(nn.Linear): + + def __init__(self, in_features: int, out_features: int, process_group: ProcessGroup, + bias: bool = True, sequence_parallel=True, device=None, dtype=None) -> None: + world_size = torch.distributed.get_world_size(process_group) + if out_features % world_size != 0: + raise ValueError(f'out_features ({out_features}) must be divisible by ' + f'world_size ({world_size})') + super().__init__(in_features, out_features // world_size, 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_torch(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, device=None, dtype=None) -> None: + world_size = torch.distributed.get_world_size(process_group) + rank = torch.distributed.get_rank(process_group) + if in_features % world_size != 0: + raise ValueError(f'in_features ({in_features}) must be divisible by ' + f'world_size ({world_size})') + # Only rank 0 will have bias + super().__init__(in_features // world_size, 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_torch(x, self.weight, self.bias) + reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce + return reduce_fn(out, self.process_group) + + class FeedForward(nn.Module): """ FeedForward. diff --git a/internlm/model/multi_head_attention.py b/internlm/model/multi_head_attention.py index 7513563..fe3e152 100644 --- a/internlm/model/multi_head_attention.py +++ b/internlm/model/multi_head_attention.py @@ -12,13 +12,12 @@ from flash_attn.modules.mha import ( SelfAttention, _update_kv_cache, ) -from flash_attn.ops.fused_dense import ColumnParallelLinear, RowParallelLinear from torch import nn from internlm.core.context import IS_TENSOR_PARALLEL, ParallelMode from internlm.core.context import global_context as gpc from internlm.model.embedding import RotaryEmbedding - +from internlm.model.linear import ColumnParallelLinear, RowParallelLinear class MHA(nn.Module): """ diff --git a/internlm/model/utils.py b/internlm/model/utils.py index b0d7264..9d430c7 100644 --- a/internlm/model/utils.py +++ b/internlm/model/utils.py @@ -2,6 +2,14 @@ # -*- encoding: utf-8 -*- import torch +import torch.nn.functional as F + +from typing import Optional +from torch import Tensor +from torch.cuda.amp import custom_bwd, custom_fwd +from torch.distributed import ProcessGroup + +from flash_attn.ops.fused_dense import FusedDenseFunc from internlm.core.context import global_context as gpc @@ -71,3 +79,170 @@ class _GatherForwardSplitBackward(torch.autograd.Function): def gather_forward_split_backward(input_, parallel_mode, dim): return _GatherForwardSplitBackward.apply(input_, parallel_mode, dim) + +def all_gather_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False): + world_size = torch.distributed.get_world_size(process_group) + output = torch.empty(world_size * input_.shape[0], *input_.shape[1:], + dtype=input_.dtype, device=input_.device) + handle = torch.distributed.all_gather_into_tensor(output, input_.contiguous(), + group=process_group, async_op=async_op) + return output, handle + +def reduce_scatter_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False): + world_size = torch.distributed.get_world_size(process_group) + assert input_.shape[0] % world_size == 0 + output = torch.empty(input_.shape[0] // world_size, *input_.shape[1:], + dtype=input_.dtype, device=input_.device) + handle = torch.distributed.reduce_scatter_tensor(output, input_.contiguous(), + group=process_group, + async_op=async_op) + return output, handle + +def all_reduce_raw(input_: Tensor, process_group: ProcessGroup, async_op: bool = False): + input_ = input_.contiguous() + handle = torch.distributed.all_reduce(input_, group=process_group, async_op=async_op) + return input_, handle + +class ReduceScatterFunc(torch.autograd.Function): + """Reduce scatter the input from the sequence parallel region and concatenate.""" + + @staticmethod + def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor: + ctx.process_group = process_group + output, _ = reduce_scatter_raw(input_, process_group) + return output + + @staticmethod + def backward(ctx, grad_output: Tensor): + grad_input, _ = all_gather_raw(grad_output, ctx.process_group) + return grad_input, None + + +class AllReduceFunc(torch.autograd.Function): + """Gather the input from sequence parallel region and concatenate.""" + + @staticmethod + def forward(ctx, input_: Tensor, process_group: ProcessGroup) -> Tensor: + ctx.process_group = process_group + output, _ = all_reduce_raw(input_, process_group) + return output + + @staticmethod + def backward(ctx, grad_output: Tensor): + return grad_output, None + + +# Supports autograd, but does not support async +reduce_scatter = ReduceScatterFunc.apply +# Supports autograd, but does not support async +all_reduce = AllReduceFunc.apply + +def linear_bias_wgrad_torch(input, grad_output, has_d_bias): + assert input.dtype == grad_output.dtype + grad_weight = torch.matmul(grad_output.t(), input) + grad_bias = grad_output.sum(dim=0) if has_d_bias else None + return grad_weight, grad_bias + + +class FusedDenseFuncTorch(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 = linear_bias_wgrad_torch( + 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_torch(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: + return FusedDenseFuncTorch.apply(x, weight, bias, return_residual, process_group, + sequence_parallel) \ No newline at end of file