"vscode:/vscode.git/clone" did not exist on "c9e6fedb6b9d6ea1e5127add3254205b0d096d7d"
fused_layer_norm.py 4.12 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""This code is copied fron NVIDIA apex:
      https://github.com/NVIDIA/apex
18
   with some changes. """
19
20

import numbers
21
import torch
22
23
24
25
from torch.nn.parameter import Parameter
from torch.nn import init
import importlib

Lawrence McAfee's avatar
fixed.  
Lawrence McAfee committed
26
27
from megatron.mpu import make_viewless_tensor

28
29
30
31
32
try:
    from apex.contrib.layer_norm.layer_norm import FastLayerNormFN
    HAVE_PERSIST_LAYER_NORM = True
except:
    HAVE_PERSIST_LAYER_NORM = False
Sangkug Lym's avatar
Sangkug Lym committed
33

34
35
36
global fused_mix_prec_layer_norm_cuda
fused_mix_prec_layer_norm_cuda = None

37

38
39
40
41
class FusedLayerNormAffineFunction(torch.autograd.Function):

  @staticmethod
  def forward(ctx, input, weight, bias, normalized_shape, eps):
42

43
44
45
46
47
48
49
50
    ctx.normalized_shape = normalized_shape
    ctx.eps = eps
    input_ = input.contiguous()
    weight_ = weight.contiguous()
    bias_ = bias.contiguous()
    output, mean, invvar = fused_mix_prec_layer_norm_cuda.forward_affine(
        input_, ctx.normalized_shape, weight_, bias_, ctx.eps)
    ctx.save_for_backward(input_, weight_, bias_, mean, invvar)
51

52
53
    return output

54

55
56
  @staticmethod
  def backward(ctx, grad_output):
57

58
59
    input_, weight_, bias_, mean, invvar = ctx.saved_tensors
    grad_input = grad_weight = grad_bias = None
60
61
    grad_input, grad_weight, grad_bias \
      = fused_mix_prec_layer_norm_cuda.backward_affine(
62
63
64
65
        grad_output.contiguous(), mean, invvar,
        input_, ctx.normalized_shape,
        weight_, bias_, ctx.eps)

66
    return grad_input, grad_weight, grad_bias, None, None
67
68
69
70



class MixedFusedLayerNorm(torch.nn.Module):
71

Sangkug Lym's avatar
Sangkug Lym committed
72
  def __init__(self, normalized_shape, eps=1e-5, no_persist_layer_norm=True):
73
74
75
        super(MixedFusedLayerNorm, self).__init__()

        global fused_mix_prec_layer_norm_cuda
76
77
        fused_mix_prec_layer_norm_cuda = importlib.import_module(
          "fused_mix_prec_layer_norm_cuda")
78

Sangkug Lym's avatar
Sangkug Lym committed
79
80
81
82
83
84
        # List of hiddens sizes supported in the persistent layer norm kernel
        # If the hidden size is not supported, fall back to the non-persistent
        # kernel.
        persist_ln_hidden_sizes = [1024, 1536, 2048, 2304, 3072, 3840, 4096,
            5120, 6144, 8192, 10240, 12288, 12800, 15360, 16384, 18432, 20480,
            24576, 25600, 30720, 32768, 40960, 49152, 65536]
85
86
        if normalized_shape not in persist_ln_hidden_sizes or \
                not HAVE_PERSIST_LAYER_NORM:
Sangkug Lym's avatar
Sangkug Lym committed
87
88
            no_persist_layer_norm = True

89
90
91
92
        if isinstance(normalized_shape, numbers.Integral):
            normalized_shape = (normalized_shape,)
        self.normalized_shape = torch.Size(normalized_shape)
        self.eps = eps
93
94
        self.weight = Parameter(torch.Tensor(*normalized_shape))
        self.bias = Parameter(torch.Tensor(*normalized_shape))
95
        self.reset_parameters()
Sangkug Lym's avatar
Sangkug Lym committed
96
        self.no_persist_layer_norm = no_persist_layer_norm
97

98
99
100
101
102
103
104
105
106

  def reset_parameters(self):

    init.ones_(self.weight)
    init.zeros_(self.bias)


  def forward(self, input):

Sangkug Lym's avatar
Sangkug Lym committed
107
    if self.no_persist_layer_norm:
Lawrence McAfee's avatar
fixed.  
Lawrence McAfee committed
108
109
        return FusedLayerNormAffineFunction.apply(
          input, self.weight, self.bias, self.normalized_shape, self.eps)
Sangkug Lym's avatar
Sangkug Lym committed
110
    else:
Lawrence McAfee's avatar
fixed.  
Lawrence McAfee committed
111
112
113
114
115
116
117
118
119
120
121
122
        output = FastLayerNormFN.apply(
          input, self.weight, self.bias, self.eps)

        # Apex's fast layer norm function outputs a 'view' tensor (i.e., has
        # a populated '_base' field). This will result in schedule.py's
        # deallocate_output_tensor() throwing an error, so a viewless tensor is
        # created to prevent this.
        output = make_viewless_tensor(inp = output,
                                      requires_grad = input.requires_grad,
                                      keep_graph = True)

        return output