mirror of https://github.com/InternLM/InternLM
support float32 training
parent
5ee651c2f1
commit
570e30a6bc
|
@ -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
|
||||
|
|
|
@ -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.
|
||||
|
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
@ -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)
|
Loading…
Reference in New Issue