transformer.py 17.1 KB
Newer Older
1
# coding=utf-8
Mohammad's avatar
Mohammad committed
2
# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
#
# 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.

"""Transformer."""

import math

import torch
from apex.normalization.fused_layer_norm import FusedLayerNorm as LayerNorm

Mohammad's avatar
Mohammad committed
23
from megatron import get_args
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
from megatron import mpu
from megatron.module import MegatronModule


""" We use the following notation throughout this file:
     h: hidden size
     n: number of attention heads
     p: number of model parallel partitions
     np: n/p
     hp: h/p
     hn: h/n
     b: batch size
     s: sequence length
     l: number of layers
    Transformer takes input of size [b, s, h] and returns a
    tensor of the same size. We use the following arguments:
        hyperparameters: transformer hyperparameters
        attention_mask_func: a function that takes `unmaksed-attention-scores`
            with size [b, np, s, s] and an `attention-mask` and will apply
            the masking. The function should return a masked score of the
            same size [b, np, s, s].
               masked-attention-scores = attention_mask_func(
                                     unmaksed-attention-scores, attention-mask)
"""

Neel Kant's avatar
Neel Kant committed
49

50
51
52
53
54
55
56
57
58
class ParallelMLP(MegatronModule):
    """MLP.

    MLP will take the input with h hidden state, project it to 4*h
    hidden dimension, perform nonlinear transformation, and project the
    state back into h hidden dimension. At the end, dropout is also
    applied.
    """

Mohammad's avatar
Mohammad committed
59
60
    def __init__(self, mlp_activation_func, init_method,
                 output_layer_init_method):
61
        super(ParallelMLP, self).__init__()
Mohammad's avatar
Mohammad committed
62
        args = get_args()
63
64
65

        # Project to 4h.
        self.dense_h_to_4h = mpu.ColumnParallelLinear(
Mohammad's avatar
Mohammad committed
66
            args.hidden_size,
Neel Kant's avatar
Neel Kant committed
67
            4 * args.hidden_size,
68
            gather_output=False,
Mohammad's avatar
Mohammad committed
69
            init_method=init_method)
70

Mohammad's avatar
Mohammad committed
71
        self.activation_func = mlp_activation_func
72
73
74

        # Project back to h.
        self.dense_4h_to_h = mpu.RowParallelLinear(
Neel Kant's avatar
Neel Kant committed
75
            4 * args.hidden_size,
Mohammad's avatar
Mohammad committed
76
            args.hidden_size,
77
            input_is_parallel=True,
Mohammad's avatar
Mohammad committed
78
            init_method=output_layer_init_method)
79

Mohammad's avatar
Mohammad committed
80
        self.dropout = torch.nn.Dropout(args.hidden_dropout)
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

    def forward(self, hidden_states):

        # [b, s, 4hp]
        intermediate_parallel = self.dense_h_to_4h(hidden_states)
        intermediate_parallel = self.activation_func(intermediate_parallel)

        # [b, s, h]
        output = self.dense_4h_to_h(intermediate_parallel)
        output = self.dropout(output)
        return output


class ParallelSelfAttention(MegatronModule):
    """Parallel self-attention layer abstract class.

    Self-attention layer takes input with size [b, s, h]
    and returns output of the same size.
    """
Neel Kant's avatar
Neel Kant committed
100

Mohammad's avatar
Mohammad committed
101
102
    def __init__(self, attention_mask_func, init_method,
                 output_layer_init_method, layer_number):
103
        super(ParallelSelfAttention, self).__init__()
Mohammad's avatar
Mohammad committed
104
        args = get_args()
Mohammad's avatar
Mohammad committed
105
        self.fp16 = args.fp16
106
107

        self.attention_mask_func = attention_mask_func
Mohammad's avatar
Mohammad committed
108
109
        self.apply_query_key_layer_scaling = args.apply_query_key_layer_scaling
        self.attention_softmax_in_fp32 = args.attention_softmax_in_fp32
110
111
112
        if self.apply_query_key_layer_scaling:
            self.attention_softmax_in_fp32 = True
        self.layer_number = max(1, layer_number)
113
114
115

        # Per attention head and per partition values.
        world_size = mpu.get_model_parallel_world_size()
Mohammad's avatar
Mohammad committed
116
117
        self.hidden_size_per_partition = mpu.divide(args.hidden_size,
                                                    world_size)
118
        self.hidden_size_per_attention_head = mpu.divide(
Mohammad's avatar
Mohammad committed
119
            args.hidden_size, args.num_attention_heads)
120
        self.num_attention_heads_per_partition = mpu.divide(
Mohammad's avatar
Mohammad committed
121
            args.num_attention_heads, world_size)
122
123
124

        # Strided linear layer.
        self.query_key_value = mpu.ColumnParallelLinear(
Mohammad's avatar
Mohammad committed
125
            args.hidden_size,
Neel Kant's avatar
Neel Kant committed
126
            3 * args.hidden_size,
127
128
            stride=3,
            gather_output=False,
Mohammad's avatar
Mohammad committed
129
            init_method=init_method)
130
131
132
133

        # Dropout. Note that for a single iteration, this layer will generate
        # different outputs on different number of parallel partitions but
        # on average it should not be partition dependent.
Mohammad's avatar
Mohammad committed
134
        self.attention_dropout = torch.nn.Dropout(args.attention_dropout)
135
136
137

        # Output.
        self.dense = mpu.RowParallelLinear(
Mohammad's avatar
Mohammad committed
138
139
            args.hidden_size,
            args.hidden_size,
140
            input_is_parallel=True,
Mohammad's avatar
Mohammad committed
141
142
            init_method=output_layer_init_method)
        self.output_dropout = torch.nn.Dropout(args.hidden_dropout)
143
144
145
146
147
148

    def _transpose_for_scores(self, tensor):
        """Transpose a 3D tensor [b, s, np*hn] into a 4D tensor with
        size [b, np, s, hn].
        """
        new_tensor_shape = tensor.size()[:-1] + \
Neel Kant's avatar
Neel Kant committed
149
150
            (self.num_attention_heads_per_partition,
             self.hidden_size_per_attention_head)
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
        tensor = tensor.view(*new_tensor_shape)
        return tensor.permute(0, 2, 1, 3)

    def _get_query_key_value(self, hidden_states):
        """Get query, key, and value and transpose to
        get size [b, np, s, hn].
        """
        # Attention heads. [b, s, hp]
        mixed_x_layer = self.query_key_value(hidden_states)
        (mixed_query_layer,
         mixed_key_layer,
         mixed_value_layer) = mpu.split_tensor_along_last_dim(mixed_x_layer, 3)

        # Reshape and transpose [b, np, s, hn]
        query_layer = self._transpose_for_scores(mixed_query_layer)
        key_layer = self._transpose_for_scores(mixed_key_layer)
        value_layer = self._transpose_for_scores(mixed_value_layer)

        return query_layer, key_layer, value_layer

    def _get_unmasked_attention_scores(self, query_layer, key_layer):
        """Unmasked attention scores with size [b, np, s, s]."""
173
174
175
176
177
        coeff = 1
        if self.apply_query_key_layer_scaling:
            coeff = self.layer_number
        norm_factor = math.sqrt(coeff *
                                math.sqrt(self.hidden_size_per_attention_head))
178
        # Raw attention scores. [b, np, s, s]
Neel Kant's avatar
Neel Kant committed
179
180
        return torch.matmul(query_layer / norm_factor,
                            key_layer.transpose(-1, -2) / norm_factor)
181
182
183
184
185
186

    def _get_attention_probs(self, attention_scores):
        """Attention probabilies with dropout. The output has
        the size [b, np, s, s].
        """
        # Attention probabilities. [b, np, s, s]
187
188
        if self.apply_query_key_layer_scaling:
            attention_scores = attention_scores * self.layer_number
189
        attention_probs = torch.nn.Softmax(dim=-1)(attention_scores)
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        with mpu.get_cuda_rng_tracker().fork():
            attention_probs = self.attention_dropout(attention_probs)

        return attention_probs

    def _get_attended_context(self, attention_probs, value_layer):
        """Final attended tesnor and transposed back to [b, s, hp]."""
        # Context layer.
        # [b, np, s, hn]
        context_layer = torch.matmul(attention_probs, value_layer)
        # [b, s, np, hn]
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + \
Neel Kant's avatar
Neel Kant committed
205
            (self.hidden_size_per_partition,)
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
        # [b, s, hp]
        context_layer = context_layer.view(*new_context_layer_shape)

        return context_layer

    def _get_output(self, context_layer):
        """Output layer with dropout."""
        # Output. [b, s, h]
        output = self.dense(context_layer)
        output = self.output_dropout(output)

        return output

    def forward(self, hidden_states, attention_mask, layer_past=None,
                get_key_value=False):
        # hidden_states: [b, s, h]

        # Attention heads. [b, np, s, hn]
        query_layer, key_layer, value_layer = self._get_query_key_value(
            hidden_states)

        if layer_past is not None:
            past_key, past_value = layer_past
            key_layer = torch.cat((past_key.type_as(key_layer),
                                   key_layer), dim=-2)
            value_layer = torch.cat((past_value.type_as(value_layer),
                                     value_layer), dim=-2)
        if get_key_value:
            present = (key_layer, value_layer)

        # Raw attention scores. [b, np, s, s]
        attention_scores = self._get_unmasked_attention_scores(
            query_layer, key_layer)

240
        # fp32 conversion.
Mohammad's avatar
Mohammad committed
241
        if self.fp16 and self.attention_softmax_in_fp32:
242
243
            attention_scores = attention_scores.float()

244
245
246
247
248
249
        # Apply attention mask. [b, np, s, s]
        if get_key_value:
            with torch.no_grad():
                if layer_past is not None:
                    attention_mask = attention_mask[
                        ...,
Neel Kant's avatar
Neel Kant committed
250
                        attention_scores.size(3) - 1,
251
252
253
254
255
256
257
258
259
260
261
262
                        :attention_scores.size(3)].unsqueeze(2)
                else:
                    attention_mask = attention_mask[
                        ...,
                        :attention_scores.size(3),
                        :attention_scores.size(3)]
        attention_scores = self.attention_mask_func(attention_scores,
                                                    attention_mask)

        # Attention probabilities. [b, np, s, s]
        attention_probs = self._get_attention_probs(attention_scores)

263
        # fp16 conversion
Mohammad's avatar
Mohammad committed
264
        if self.fp16 and self.attention_softmax_in_fp32:
265
266
            attention_probs = attention_probs.half()

267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
        # Context layer. [b, s, hp]
        context_layer = self._get_attended_context(attention_probs, value_layer)

        # Output. [b, s, h]
        output = self._get_output(context_layer)

        if get_key_value:
            output = [output, present]

        return output


class ParallelTransformerLayer(MegatronModule):
    """A single transformer layer.

    Transformore layer takes input with size [b, s, h] and returns an
    output of the same size.
    """
Neel Kant's avatar
Neel Kant committed
285

Mohammad's avatar
Mohammad committed
286
287
288
    def __init__(self, attention_mask_func, mlp_activation_func,
                 init_method, output_layer_init_method, layer_number):
        args = get_args()
289
290

        super(ParallelTransformerLayer, self).__init__()
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
291
        self.layer_number = layer_number
292
293

        self.apply_residual_connection_post_layernorm \
Mohammad's avatar
Mohammad committed
294
            = args.apply_residual_connection_post_layernorm
295
296
297

        # Layernorm on the input data.
        self.input_layernorm = LayerNorm(
Mohammad's avatar
Mohammad committed
298
299
            args.hidden_size,
            eps=args.layernorm_epsilon)
300
301

        # Self attention.
Mohammad's avatar
Mohammad committed
302
303
304
        self.attention = ParallelSelfAttention(attention_mask_func, init_method,
                                               output_layer_init_method,
                                               layer_number)
305
306
307

        # Layernorm on the input data.
        self.post_attention_layernorm = LayerNorm(
Mohammad's avatar
Mohammad committed
308
309
            args.hidden_size,
            eps=args.layernorm_epsilon)
310
311

        # MLP
Mohammad's avatar
Mohammad committed
312
313
        self.mlp = ParallelMLP(mlp_activation_func, init_method,
                               output_layer_init_method)
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353

    def forward(self, hidden_states, attention_mask, layer_past=None,
                get_key_value=False):
        # hidden_states: [b, s, h]

        # Layer norm at the begining of the transformer layer.
        layernorm_output = self.input_layernorm(hidden_states)
        # Self attention.
        attention_output = self.attention(layernorm_output,
                                          attention_mask,
                                          layer_past=layer_past,
                                          get_key_value=get_key_value)
        if get_key_value:
            attention_output, presents = attention_output

        # Residual connection.
        if self.apply_residual_connection_post_layernorm:
            layernorm_input = layernorm_output + attention_output
        else:
            layernorm_input = hidden_states + attention_output
        # Layer norm post the self attention.
        layernorm_output = self.post_attention_layernorm(layernorm_input)

        # MLP.
        mlp_output = self.mlp(layernorm_output)
        # Second residual connection.
        if self.apply_residual_connection_post_layernorm:
            output = layernorm_output + mlp_output
        else:
            output = layernorm_input + mlp_output

        if get_key_value:
            output = [output, presents]

        return output


class ParallelTransformer(MegatronModule):
    """Transformer class."""

Mohammad's avatar
Mohammad committed
354
355
    def __init__(self, attention_mask_func, mlp_activation_func,
                 init_method, output_layer_init_method):
356
        super(ParallelTransformer, self).__init__()
Mohammad's avatar
Mohammad committed
357
        args = get_args()
358
359

        # Store activation checkpoiting flag.
Mohammad's avatar
Mohammad committed
360
361
        self.checkpoint_activations = args.checkpoint_activations
        self.checkpoint_num_layers = args.checkpoint_num_layers
362

Mohammad's avatar
Mohammad committed
363
364
365
366
367
368
369
370
371
372
373
374
        # Number of layers:
        self.num_layers = args.num_layers
        self.num_unique_layers = args.num_unique_layers
        if self.num_unique_layers is None:
            self.num_unique_layers = self.num_layers
        assert self.num_layers % self.num_unique_layers == 0, \
            'number of layers should be divisible by number of unique layers'
        self.param_sharing_style = args.param_sharing_style
        assert self.param_sharing_style in ['grouped', 'spaced']

        # Transformer layers.
        def build_layer(layer_number):
375
            return ParallelTransformerLayer(
Mohammad's avatar
Mohammad committed
376
377
                attention_mask_func, mlp_activation_func,
                init_method, output_layer_init_method, layer_number)
378
        self.layers = torch.nn.ModuleList(
Mohammad's avatar
Mohammad committed
379
380
381
382
383
384
385
386
387
            [build_layer(i + 1) for i in range(self.num_unique_layers)])

        # Print layer ordering.
        if self.num_layers != self.num_unique_layers:
            if torch.distributed.get_rank() == 0:
                print('> will be using the following layer ordering:')
                for i in range(self.num_layers):
                    print('   layer: {:3d} --> unique layer: {:3d}'.format(
                        i, self._get_layer_index(i)), flush=True)
388
389
390

        # Final layer norm before output.
        self.final_layernorm = LayerNorm(
Mohammad's avatar
Mohammad committed
391
392
            args.hidden_size,
            eps=args.layernorm_epsilon)
393

Mohammad's avatar
Mohammad committed
394
395
396
397
398
399
400
401
402
403
    def _get_layer_index(self, layer_number):
        if self.param_sharing_style == 'grouped':
            return layer_number % self.num_unique_layers
        if self.param_sharing_style == 'spaced':
            return layer_number // (self.num_layers // self.num_unique_layers) 
        assert False, 'should not be here'

    def _get_layer(self, layer_number):
        return self.layers[self._get_layer_index(layer_number)]

404
405
406
407
408
    def _checkpointed_forward(self, hidden_states, attention_mask):
        """Forward method with activation checkpointing."""
        def custom(start, end):
            def custom_forward(*inputs):
                x_ = inputs[0]
Mohammad's avatar
Mohammad committed
409
410
                for index in range(start, end):
                    layer = self._get_layer(index)
411
412
413
414
415
                    x_ = layer(x_, inputs[1])
                return x_
            return custom_forward

        l = 0
Mohammad's avatar
Mohammad committed
416
        while l < self.num_layers:
417
            hidden_states = mpu.checkpoint(
Neel Kant's avatar
Neel Kant committed
418
                custom(l, l + self.checkpoint_num_layers),
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
                hidden_states, attention_mask)
            l += self.checkpoint_num_layers

        return hidden_states

    def forward(self, hidden_states, attention_mask, layer_past=None,
                get_key_value=False):

        # Checks
        if layer_past is not None:
            assert get_key_value, \
                'for not None values in layer_past, ' \
                'expected get_key_value to be set'
        if get_key_value:
            assert not self.checkpoint_activations, \
                'get_key_value does not work with ' \
                'activation checkpointing'

        if self.checkpoint_activations:
            hidden_states = self._checkpointed_forward(hidden_states,
                                                       attention_mask)
        else:
            if get_key_value:
                presents = []
Mohammad's avatar
Mohammad committed
443
444
            for index in range(self.num_layers):
                layer = self._get_layer(index)
445
446
                past = None
                if layer_past is not None:
Mohammad's avatar
Mohammad committed
447
                    past = layer_past[index]
448
449
450
451
452
453
454
455
456
457
458
459
460
461
                hidden_states = layer(hidden_states,
                                      attention_mask,
                                      layer_past=past,
                                      get_key_value=get_key_value)
                if get_key_value:
                    hidden_states, present = hidden_states
                    presents.append(present)

        # Final layer norm.
        output = self.final_layernorm(hidden_states)
        if get_key_value:
            output = [output, presents]

        return output