github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-25 06:52:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix] issue #2388

Browse Source
This commit is contained in:
Jiarui Fang 2023-01-07 19:36:48 +08:00 committed by GitHub
commit 58b013060f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/kernel/cuda_native/layer_norm.py
View File

@ -16,17 +16,17 @@ class FusedLayerNormAffineFunction(torch.autograd.Function):
@custom_fwd(cast_inputs=torch.float32) @custom_fwd(cast_inputs=torch.float32)
def forward(ctx, input, weight, bias, normalized_shape, eps): def forward(ctx, input, weight, bias, normalized_shape, eps):
try: try:
import colossalai._C.layer_norm from colossalai._C import layer_norm
except ImportError: except ImportError:
raise RuntimeError('FusedLayerNormAffineFunction requires cuda extensions') from colossalai.kernel.op_builder.layernorm import LayerNormBuilder
layer_norm = LayerNormBuilder().load()
ctx.normalized_shape = normalized_shape ctx.normalized_shape = normalized_shape
ctx.eps = eps ctx.eps = eps
input_ = input.contiguous() input_ = input.contiguous()
weight_ = weight.contiguous() weight_ = weight.contiguous()
bias_ = bias.contiguous() bias_ = bias.contiguous()
output, mean, invvar = colossalai._C.layer_norm.forward_affine(input_, ctx.normalized_shape, weight_, bias_, output, mean, invvar = layer_norm.forward_affine(input_, ctx.normalized_shape, weight_, bias_, ctx.eps)
ctx.eps)
ctx.save_for_backward(input_, weight_, bias_, mean, invvar) ctx.save_for_backward(input_, weight_, bias_, mean, invvar)
return output return output
@ -35,14 +35,15 @@ class FusedLayerNormAffineFunction(torch.autograd.Function):
@custom_bwd @custom_bwd
def backward(ctx, grad_output): def backward(ctx, grad_output):
try: try:
import colossalai._C.layer_norm from colossalai._C import layer_norm
except ImportError: except ImportError:
raise RuntimeError('FusedLayerNormAffineFunction requires cuda extensions') from colossalai.kernel.op_builder.layernorm import LayerNormBuilder
layer_norm = LayerNormBuilder().load()
input_, weight_, bias_, mean, invvar = ctx.saved_tensors input_, weight_, bias_, mean, invvar = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None grad_input = grad_weight = grad_bias = None
grad_input, grad_weight, grad_bias \ grad_input, grad_weight, grad_bias \
= colossalai._C.layer_norm.backward_affine( = layer_norm.backward_affine(
grad_output.contiguous(), mean, invvar, grad_output.contiguous(), mean, invvar,
input_, ctx.normalized_shape, input_, ctx.normalized_shape,
weight_, bias_, ctx.eps) weight_, bias_, ctx.eps)

14
colossalai/kernel/cuda_native/scaled_softmax.py
View File

@ -53,26 +53,28 @@ class ScaledMaskedSoftmax(torch.autograd.Function):
@staticmethod @staticmethod
def forward(ctx, inputs, mask, scale): def forward(ctx, inputs, mask, scale):
try: try:
import colossalai._C.scaled_masked_softmax from colossalai._C import scaled_masked_softmax
except ImportError: except ImportError:
raise RuntimeError('ScaledMaskedSoftmax requires cuda extensions') from colossalai.kernel.op_builder.scaled_masked_softmax import ScaledMaskedSoftmaxBuilder
scaled_masked_softmax = ScaledMaskedSoftmaxBuilder().load()
scale_t = torch.tensor([scale]) scale_t = torch.tensor([scale])
softmax_results = colossalai._C.scaled_masked_softmax.forward(inputs, mask, scale_t[0]) softmax_results = scaled_masked_softmax.forward(inputs, mask, scale_t[0])
ctx.save_for_backward(softmax_results, scale_t) ctx.save_for_backward(softmax_results, scale_t)
return softmax_results return softmax_results
@staticmethod @staticmethod
def backward(ctx, output_grads): def backward(ctx, output_grads):
try: try:
import colossalai._C.scaled_masked_softmax from colossalai._C import scaled_masked_softmax
except ImportError: except ImportError:
raise RuntimeError('ScaledMaskedSoftmax requires cuda extensions') from colossalai.kernel.op_builder.scaled_masked_softmax import ScaledMaskedSoftmaxBuilder
scaled_masked_softmax = ScaledMaskedSoftmaxBuilder().load()
softmax_results, scale_t = ctx.saved_tensors softmax_results, scale_t = ctx.saved_tensors
input_grads = colossalai._C.scaled_masked_softmax.backward(output_grads, softmax_results, scale_t[0]) input_grads = scaled_masked_softmax.backward(output_grads, softmax_results, scale_t[0])
return input_grads, None, None return input_grads, None, None