transformer.py 40 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
#
# 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
19
import torch.nn.functional as F
20

Mohammad's avatar
Mohammad committed
21
from megatron import get_args
22
from megatron import mpu
23
from .module import MegatronModule
24
from megatron.model.enums import AttnMaskType, ModelType, LayerType, AttnType
25
from megatron.model import LayerNorm
26
27
from megatron.model.fused_softmax import FusedScaleMaskSoftmax
from megatron.model.fused_bias_gelu import bias_gelu_impl
28
from megatron.model.utils import attention_mask_func, openai_gelu, erf_gelu
29
30
31
32
33
34
35
36
37
38
39

""" 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
40
    Transformer takes input of size [s, b, h] and returns a
41
42
43
44
    tensor of the same size. We use the following arguments:
        hyperparameters: transformer hyperparameters
"""

45
46
47
48
49
50

class DropPath(MegatronModule):
    """Drop paths (Stochastic Depth) per sample 
    (when applied in main path of residual blocks).
    """

Vijay Korthikanti's avatar
Vijay Korthikanti committed
51
    def __init__(self, drop_prob=0.):
52
53
54
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

Vijay Korthikanti's avatar
Vijay Korthikanti committed
55
    def forward(self, hidden_state):
56
        if self.drop_prob == 0. or not self.training:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
57
            return hidden_state
58
59
        keep_prob = 1 - self.drop_prob
        # work with diff dim tensors, not just 2D ConvNets
Vijay Korthikanti's avatar
Vijay Korthikanti committed
60
        shape = (hidden_state.shape[0],) + (1,) * (hidden_state.ndim - 1)
61
        random_tensor = keep_prob + \
Vijay Korthikanti's avatar
Vijay Korthikanti committed
62
            torch.rand(shape, dtype=hidden_state.dtype, device=hidden_state.device)
63
        random_tensor.floor_()  # binarize
Vijay Korthikanti's avatar
Vijay Korthikanti committed
64
        output = hidden_state.div(keep_prob) * random_tensor
65
66
67
        return output


68
69
70
71
72
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
hwijeen's avatar
hwijeen committed
73
    state back into h hidden dimension.
74
75
    """

76
    def __init__(self, init_method, output_layer_init_method):
77
        super(ParallelMLP, self).__init__()
Mohammad's avatar
Mohammad committed
78
        args = get_args()
79
80
81

        # Project to 4h.
        self.dense_h_to_4h = mpu.ColumnParallelLinear(
Mohammad's avatar
Mohammad committed
82
            args.hidden_size,
83
            args.ffn_hidden_size,
84
            gather_output=False,
85
86
            init_method=init_method,
            skip_bias_add=True)
87

88
89
90
91
92
93
        self.bias_gelu_fusion = args.bias_gelu_fusion
        self.activation_func = F.gelu
        if args.openai_gelu:
            self.activation_func = openai_gelu
        elif args.onnx_safe:
            self.activation_func = erf_gelu
94
95
96

        # Project back to h.
        self.dense_4h_to_h = mpu.RowParallelLinear(
97
            args.ffn_hidden_size,
Mohammad's avatar
Mohammad committed
98
            args.hidden_size,
99
            input_is_parallel=True,
100
101
            init_method=output_layer_init_method,
            skip_bias_add=True)
102

103
104
    def forward(self, hidden_states):

105
106
        # [s, b, 4hp]
        intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states)
107

108
109
110
111
112
113
114
115
116
117
        if self.bias_gelu_fusion:
             intermediate_parallel = \
                     bias_gelu_impl(intermediate_parallel, bias_parallel)
        else:
            intermediate_parallel = \
                self.activation_func(intermediate_parallel + bias_parallel)

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

rprenger's avatar
rprenger committed
119
120
121
122
class SwitchMLP(MegatronModule):
    """
    Routes input to one of N MLP "experts"
    """
rprenger's avatar
rprenger committed
123
    def __init__(self, init_method, output_layer_init_method):
rprenger's avatar
rprenger committed
124
125
        super(SwitchMLP, self).__init__()
        args = get_args()
rprenger's avatar
rprenger committed
126
        self.router = torch.nn.Linear(args.hidden_size, args.num_experts)
rprenger's avatar
rprenger committed
127
        self.experts = torch.nn.ModuleList()
rprenger's avatar
rprenger committed
128
        for i in range(args.num_experts):
rprenger's avatar
rprenger committed
129
            self.experts.append(ParallelMLP(init_method, output_layer_init_method))
130

rprenger's avatar
rprenger committed
131
132
133
134
135
136
    def forward(self, hidden_states):
        # hidden_states: [b, s, h]
        b = hidden_states.size(0)
        s = hidden_states.size(1)
        h = hidden_states.size(2)
        route = self.router(hidden_states)
rprenger's avatar
rprenger committed
137
        route = torch.nn.functional.softmax(route, dim=2)
rprenger's avatar
rprenger committed
138
        max_prob, max_ind = torch.max(route, dim=2)
139
140
        max_prob = torch.unsqueeze(max_prob, 2) # [b s 1]

rprenger's avatar
rprenger committed
141
142
        # TODO (rprenger) TODO this could be made easier to read
        # Converting [b, s, h] to [b*s, h].
143
144
145
146
        # Each vector could be routed differently
        hidden_states = hidden_states.view(-1, hidden_states.size(2)) # [b*s h]
        max_prob = max_prob.view(-1, max_prob.size(2)) # [b*s 1]
        max_ind = max_ind.view(-1) # [b*s]
rprenger's avatar
rprenger committed
147
148
149

        output_total = torch.empty_like(hidden_states)
        output_bias_total = torch.empty_like(hidden_states)
rprenger's avatar
rprenger committed
150
        #TODO (rprenger) This does each expert in serial, but it could be parallelized
151
        
rprenger's avatar
rprenger committed
152
        for expert_num, expert in enumerate(self.experts):
153
154
            local_indices = (max_ind == expert_num).nonzero()
            hidden = hidden_states[local_indices,:]
rprenger's avatar
rprenger committed
155
156
            output, output_bias = expert(hidden)
            output_bias = output_bias.expand_as(output)
157
158
159
            output_total[local_indices,:] = output
            output_bias_total[local_indices,:] = output_bias

rprenger's avatar
rprenger committed
160
161
        output_total = output_total*max_prob
        output_bias_total = output_bias_total*max_prob
162
163
        output_total = output_total.view(b, s, h)
        output_bias_total = output_bias_total.view(b, s, h)
rprenger's avatar
rprenger committed
164
165

        return output_total, output_bias_total
166

167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
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
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304

class CoreAttention(MegatronModule):
    def __init__(self, layer_number,
                 attn_mask_type=AttnMaskType.padding):
        super(CoreAttention, self).__init__()
        args = get_args()
        self.fp16 = args.fp16
        self.bf16 = args.bf16

        self.apply_query_key_layer_scaling = args.apply_query_key_layer_scaling
        self.attention_softmax_in_fp32 = args.attention_softmax_in_fp32
        if self.apply_query_key_layer_scaling:
            self.attention_softmax_in_fp32 = True
        self.layer_number = max(1, layer_number)
        self.attn_mask_type = attn_mask_type
        self.model_parallel_memory_opt = args.model_parallel_memory_opt

        projection_size = args.kv_channels * args.num_attention_heads

        # Per attention head and per partition values.
        world_size = mpu.get_tensor_model_parallel_world_size()
        self.hidden_size_per_partition = mpu.divide(projection_size,
                                                    world_size)
        self.hidden_size_per_attention_head = mpu.divide(
            projection_size, args.num_attention_heads)

        coeff = None
        self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
        if self.apply_query_key_layer_scaling:
            coeff = self.layer_number
            self.norm_factor *= coeff

        self.scale_mask_softmax = FusedScaleMaskSoftmax(
            self.fp16, self.bf16,
            self.attn_mask_type,
            args.masked_softmax_fusion,
            attention_mask_func,
            self.attention_softmax_in_fp32,
            coeff)

        # 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.
        self.attention_dropout = torch.nn.Dropout(args.attention_dropout)
        
    def forward(self, query_layer, key_layer,
                value_layer, attention_mask):

        # ===================================
        # Raw attention scores. [b, np, s, s]
        # ===================================

        # [b, np, sq, sk]
        output_size = (query_layer.size(1),
                       query_layer.size(2),
                       query_layer.size(0),
                       key_layer.size(0))

        # [sq, b, np, hn] -> [sq, b * np, hn]
        query_layer = query_layer.view(output_size[2],
                                       output_size[0] * output_size[1], -1)
        # [sk, b, np, hn] -> [sk, b * np, hn]
        key_layer = key_layer.view(output_size[3],
                                   output_size[0] * output_size[1], -1)

        # preallocting result tensor: [b * np, sq, sk]
        matmul_result = torch.empty(
            output_size[0]*output_size[1],
            output_size[2],
            output_size[3],
            dtype=query_layer.dtype,
            device=torch.cuda.current_device())

        # Raw attention scores. [b * np, sq, sk]
        matmul_result = torch.baddbmm(
            matmul_result,
            query_layer.transpose(0, 1),   # [b * np, sq, hn]
            key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
            beta=0.0, alpha=(1.0/self.norm_factor))

        # change view to [b, np, sq, sk]
        attention_scores = matmul_result.view(*output_size)

        # ===========================
        # Attention probs and dropout
        # ===========================

        # attention scores and attention mask [b, np, sq, sk]
        attention_probs = self.scale_mask_softmax(attention_scores,
                                                  attention_mask)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.

        if not self.model_parallel_memory_opt:
            with mpu.get_cuda_rng_tracker().fork():
                attention_probs = self.attention_dropout(attention_probs)
        else:
            attention_probs = self.attention_dropout(attention_probs)

        # =========================
        # Context layer. [sq, b, hp]
        # =========================

        # value_layer -> context layer.
        # [sk, b, np, hn] --> [b, np, sq, hn]

        # context layer shape: [b, np, sq, hn]
        output_size = (value_layer.size(1),
                       value_layer.size(2),
                       query_layer.size(0),
                       value_layer.size(3))

        # change view [sk, b * np, hn]
        value_layer = value_layer.view(value_layer.size(0),
                                       output_size[0] * output_size[1], -1)

        # change view [b * np, sq, sk]
        attention_probs = attention_probs.view(output_size[0] * output_size[1],
                                               output_size[2], -1)

        # matmul: [b * np, sq, hn]
        context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))

        # change view [b, np, sq, hn]
        context_layer = context_layer.view(*output_size)

        # [b, np, sq, hn] --> [sq, b, np, hn]
        context_layer = context_layer.permute(2, 0, 1, 3).contiguous()

        # [sq, b, np, hn] --> [sq, b, hp]
        new_context_layer_shape = context_layer.size()[:-2] + \
            (self.hidden_size_per_partition,)
        context_layer = context_layer.view(*new_context_layer_shape)

        return context_layer


305
class ParallelAttention(MegatronModule):
306
307
308
309
310
    """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
311

312
    def __init__(self, init_method,
313
314
315
316
                 output_layer_init_method, layer_number,
                 attention_type=AttnType.self_attn,
                 attn_mask_type=AttnMaskType.padding):
        super(ParallelAttention, self).__init__()
Mohammad's avatar
Mohammad committed
317
        args = get_args()
318
        self.layer_number = max(1, layer_number)
319
320
        self.attention_type = attention_type
        self.attn_mask_type = attn_mask_type
321
        self.params_dtype = args.params_dtype
322
323
        self.checkpoint_attention = args.checkpoint_attention
        #assert args.activations_checkpoint_method is None
324
325

        projection_size = args.kv_channels * args.num_attention_heads
326
327

        # Per attention head and per partition values.
328
        world_size = mpu.get_tensor_model_parallel_world_size()
329
        self.hidden_size_per_attention_head = mpu.divide(
330
            projection_size, args.num_attention_heads)
331
        self.num_attention_heads_per_partition = mpu.divide(
Mohammad's avatar
Mohammad committed
332
            args.num_attention_heads, world_size)
333
334

        # Strided linear layer.
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
        if attention_type == AttnType.self_attn:
            self.query_key_value = mpu.ColumnParallelLinear(
                args.hidden_size,
                3 * projection_size,
                gather_output=False,
                init_method=init_method)
        else:
            assert attention_type == AttnType.cross_attn
            self.query = mpu.ColumnParallelLinear(
                args.hidden_size,
                projection_size,
                gather_output=False,
                init_method=init_method)

            self.key_value = mpu.ColumnParallelLinear(
                args.hidden_size,
                2 * projection_size,
                gather_output=False,
                init_method=init_method)
354

355
356
        self.core_attention = CoreAttention(self.layer_number,
                                            self.attn_mask_type)
357
358
359

        # Output.
        self.dense = mpu.RowParallelLinear(
360
            projection_size,
Mohammad's avatar
Mohammad committed
361
            args.hidden_size,
362
            input_is_parallel=True,
363
364
            init_method=output_layer_init_method,
            skip_bias_add=True)
Vijay Korthikanti's avatar
Vijay Korthikanti committed
365

366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
    def _checkpointed_attention_forward(self, query_layer, key_layer,
                                        value_layer, attention_mask):
        """Forward method with activation checkpointing."""
        def custom_forward(*inputs):
            query_layer = inputs[0]
            key_layer = inputs[1]
            value_layer = inputs[2]
            attention_mask = inputs[3]
            output_ = self.core_attention(query_layer, key_layer,
                                          value_layer, attention_mask)
            return output_

        hidden_states = mpu.checkpoint(
            custom_forward,
            False, query_layer, key_layer, value_layer, attention_mask)

        return hidden_states
383
384
385
386
387
388
389
390
391
392
393

    def _allocate_memory(self, inference_max_sequence_len, batch_size):
        return torch.empty(
            inference_max_sequence_len,
            batch_size,
            self.num_attention_heads_per_partition,
            self.hidden_size_per_attention_head,
            dtype=self.params_dtype,
            device=torch.cuda.current_device())

    def forward(self, hidden_states, attention_mask,
mshoeybi's avatar
mshoeybi committed
394
                encoder_output=None, inference_params=None):
395
        # hidden_states: [sq, b, h]
396

397
398
399
        # =================================================
        # Pre-allocate memory for key-values for inference.
        # =================================================
mshoeybi's avatar
mshoeybi committed
400
        if inference_params:
401
            if self.layer_number not in inference_params.key_value_memory_dict:
mshoeybi's avatar
mshoeybi committed
402
                inf_max_seq_len = inference_params.max_sequence_len
mshoeybi's avatar
mshoeybi committed
403
                inf_max_batch_size = inference_params.max_batch_size
404
                inference_key_memory = self._allocate_memory(
mshoeybi's avatar
mshoeybi committed
405
                    inf_max_seq_len, inf_max_batch_size)
406
                inference_value_memory = self._allocate_memory(
mshoeybi's avatar
mshoeybi committed
407
                    inf_max_seq_len, inf_max_batch_size)
408
409
410
411
412
                inference_params.key_value_memory_dict[self.layer_number] = (
                    inference_key_memory, inference_value_memory)
            else:
                inference_key_memory, inference_value_memory = \
                    inference_params.key_value_memory_dict[self.layer_number]
mshoeybi's avatar
mshoeybi committed
413

414
415
416
        # =====================
        # Query, Key, and Value
        # =====================
417

418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
        if self.attention_type == AttnType.self_attn:
            # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
            mixed_x_layer, _ = self.query_key_value(hidden_states)

            # [sq, b, (np * 3 * hn)] --> [sq, b, np, 3 * hn]
            new_tensor_shape = mixed_x_layer.size()[:-1] + \
                (self.num_attention_heads_per_partition,
                 3 * self.hidden_size_per_attention_head)
            mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)

            # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
            (query_layer,
             key_layer,
             value_layer) = mpu.split_tensor_along_last_dim(mixed_x_layer, 3)
        else:
            # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)]
            mixed_kv_layer, _ = self.key_value(encoder_output)

            # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn]
            new_tensor_shape = mixed_kv_layer.size()[:-1] + \
                (self.num_attention_heads_per_partition,
                 2 * self.hidden_size_per_attention_head)
            mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape)

            # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn]
            (key_layer,
             value_layer) = mpu.split_tensor_along_last_dim(mixed_kv_layer, 2)

            # Attention head [sq, b, h] --> [sq, b, hp]
            query_layer, _ = self.query(hidden_states)
            # [sq, b, hp] --> [sq, b, np, hn]
            new_tensor_shape = query_layer.size()[:-1] + \
                (self.num_attention_heads_per_partition,
                 self.hidden_size_per_attention_head)
            query_layer = query_layer.view(*new_tensor_shape)
453

mshoeybi's avatar
mshoeybi committed
454
455
456
        # ==================================
        # Adjust key and value for inference
        # ==================================
457

mshoeybi's avatar
mshoeybi committed
458
        if inference_params:
mshoeybi's avatar
mshoeybi committed
459
460
            batch_start = inference_params.batch_size_offset
            batch_end = batch_start + key_layer.size(1)
461
            assert batch_end <= inference_key_memory.size(1)
mshoeybi's avatar
mshoeybi committed
462
463
            sequence_start = inference_params.sequence_len_offset
            sequence_end = sequence_start + key_layer.size(0)
464
            assert sequence_end <= inference_key_memory.size(0)
465
            # Copy key and values.
466
467
468
469
470
            inference_key_memory[sequence_start:sequence_end,
                                 batch_start:batch_end, ...] = key_layer
            inference_value_memory[sequence_start:sequence_end,
                                   batch_start:batch_end, ...] = value_layer
            key_layer = inference_key_memory[
mshoeybi's avatar
mshoeybi committed
471
                :sequence_end, batch_start:batch_end, ...]
472
            value_layer = inference_value_memory[
mshoeybi's avatar
mshoeybi committed
473
                :sequence_end, batch_start:batch_end, ...]
474

475
476
477
        # ==================================
        # core attention computation
        # ==================================
478

479
480
481
        if self.checkpoint_attention:
            context_layer = self._checkpointed_attention_forward(
                query_layer, key_layer, value_layer, attention_mask)
Vijay Korthikanti's avatar
Vijay Korthikanti committed
482
        else:
483
484
            context_layer = self.core_attention(
                query_layer, key_layer, value_layer, attention_mask)
485
486

        # =================
487
        # Output. [sq, b, h]
488
489
490
        # =================

        output, bias = self.dense(context_layer)
491

492
493
494
        return output, bias


495
def bias_dropout_add(x, bias, residual, prob, training):
496
497
498
499
500
501
502
503
504
505
506
507
508
    # type: (Tensor, Tensor, Tensor, float, bool) -> Tensor
    out = torch.nn.functional.dropout(x + bias, p=prob, training=training)
    out = residual + out
    return out


def get_bias_dropout_add(training):
    def _bias_dropout_add(x, bias, residual, prob):
        return bias_dropout_add(x, bias, residual, prob, training)
    return _bias_dropout_add


@torch.jit.script
509
510
511
512
def bias_dropout_add_fused_train(x: torch.Tensor,
                                 bias: torch.Tensor,
                                 residual: torch.Tensor,
                                 prob: float) -> torch.Tensor:
513
514
515
516
    return bias_dropout_add(x, bias, residual, prob, True)


@torch.jit.script
517
518
519
520
def bias_dropout_add_fused_inference(x: torch.Tensor,
                                     bias: torch.Tensor,
                                     residual: torch.Tensor,
                                     prob: float) -> torch.Tensor:
521
    return bias_dropout_add(x, bias, residual, prob, False)
522
523
524
525
526


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

527
    Transformer layer takes input with size [b, s, h] and returns an
528
529
    output of the same size.
    """
Neel Kant's avatar
Neel Kant committed
530

531
532
    def __init__(self, init_method, output_layer_init_method,
                 layer_number, layer_type=LayerType.encoder,
533
534
                 self_attn_mask_type=AttnMaskType.padding,
                 drop_path_rate=0.):
Mohammad's avatar
Mohammad committed
535
        args = get_args()
536
537

        super(ParallelTransformerLayer, self).__init__()
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
538
        self.layer_number = layer_number
539
        self.layer_type = layer_type
540
541

        self.apply_residual_connection_post_layernorm \
Mohammad's avatar
Mohammad committed
542
            = args.apply_residual_connection_post_layernorm
543

Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
544
545
546
        self.bf16 = args.bf16
        self.fp32_residual_connection = args.fp32_residual_connection

547
548
        # Layernorm on the input data.
        self.input_layernorm = LayerNorm(
Mohammad's avatar
Mohammad committed
549
            args.hidden_size,
Sangkug Lym's avatar
Sangkug Lym committed
550
            eps=args.layernorm_epsilon,
551
552
            no_persist_layer_norm=args.no_persist_layer_norm,
            sequence_parallel=args.model_parallel_memory_opt)
553
554

        # Self attention.
555
556
557
558
559
560
        self.self_attention = ParallelAttention(
            init_method,
            output_layer_init_method,
            layer_number,
            attention_type=AttnType.self_attn,
            attn_mask_type=self_attn_mask_type)
561
562
        self.hidden_dropout = args.hidden_dropout
        self.bias_dropout_fusion = args.bias_dropout_fusion
Vijay Korthikanti's avatar
Vijay Korthikanti committed
563
        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None
564

565
        # Layernorm on the attention output
566
        self.post_attention_layernorm = LayerNorm(
Mohammad's avatar
Mohammad committed
567
            args.hidden_size,
Sangkug Lym's avatar
Sangkug Lym committed
568
            eps=args.layernorm_epsilon,
569
570
            no_persist_layer_norm=args.no_persist_layer_norm,
            sequence_parallel=args.model_parallel_memory_opt)
571

572
573
574
575
576
577
578
579
580
        if self.layer_type == LayerType.decoder:
            self.inter_attention = ParallelAttention(
                init_method,
                output_layer_init_method,
                layer_number,
                attention_type=AttnType.cross_attn)
            # Layernorm on the attention output.
            self.post_inter_attention_layernorm = LayerNorm(
                args.hidden_size,
Sangkug Lym's avatar
Sangkug Lym committed
581
                eps=args.layernorm_epsilon,
582
583
                no_persist_layer_norm=args.no_persist_layer_norm,
                sequence_parallel=args.model_parallel_memory_opt)
584

585
        # MLP
rprenger's avatar
rprenger committed
586
587
588
589
        if args.num_experts is not None:
            self.mlp = SwitchMLP(init_method, output_layer_init_method)
        else:
            self.mlp = ParallelMLP(init_method, output_layer_init_method)
590

591
    def forward(self, hidden_states, attention_mask,
mshoeybi's avatar
mshoeybi committed
592
593
                encoder_output=None, enc_dec_attn_mask=None,
                inference_params=None):
594
595
        # hidden_states: [b, s, h]

596
        # Layer norm at the beginning of the transformer layer.
597
598
        layernorm_output = self.input_layernorm(hidden_states)
        # Self attention.
599
        attention_output, attention_bias = \
600
601
602
            self.self_attention(
                layernorm_output,
                attention_mask,
mshoeybi's avatar
mshoeybi committed
603
                inference_params=inference_params)
604

605
606
        # Residual connection.
        if self.apply_residual_connection_post_layernorm:
607
608
609
610
            residual = layernorm_output
        else:
            residual = hidden_states

Vijay Korthikanti's avatar
Vijay Korthikanti committed
611
        if self.drop_path is None:
612
613
614
615
616
617
618
619
620
            # jit scripting for a nn.module (with dropout) is not
            # trigerring the fusion kernel. For now, we use two
            # different nn.functional routines to account for varying
            # dropout semantics during training and inference phases.
            if self.bias_dropout_fusion:
                if self.training:
                    bias_dropout_add_func = bias_dropout_add_fused_train
                else:
                    bias_dropout_add_func = bias_dropout_add_fused_inference
621
            else:
622
                bias_dropout_add_func = get_bias_dropout_add(self.training)
623

624
625
626
627
628
629
630
631
632
633
634
635
            # re-enable torch grad to enable fused optimization.
            with torch.enable_grad():
                layernorm_input = bias_dropout_add_func(
                    attention_output,
                    attention_bias.expand_as(residual),
                    residual,
                    self.hidden_dropout)
        else:
            out = torch.nn.functional.dropout(attention_output + attention_bias,
                                              p=self.hidden_dropout,
                                              training=self.training)
            layernorm_input = residual + self.drop_path(out)
636

637
638
639
        # Layer norm post the self attention.
        layernorm_output = self.post_attention_layernorm(layernorm_input)

640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
        if self.layer_type == LayerType.decoder:
            attention_output, attention_bias = \
                self.inter_attention(layernorm_output,
                                     enc_dec_attn_mask,
                                     encoder_output=encoder_output)
            # residual connection
            if self.apply_residual_connection_post_layernorm:
                residual = layernorm_output
            else:
                residual = layernorm_input

            # re-enable torch grad to enable fused optimization.
            with torch.enable_grad():
                layernorm_input = bias_dropout_add_func(
                    attention_output,
                    attention_bias.expand_as(residual),
                    residual,
                    self.hidden_dropout)

            # Layer norm post the decoder attention
            layernorm_output = self.post_inter_attention_layernorm(layernorm_input)

662
        # MLP.
663
        mlp_output, mlp_bias = self.mlp(layernorm_output)
664

665
666
        # Second residual connection.
        if self.apply_residual_connection_post_layernorm:
667
            residual = layernorm_output
668
        else:
669
670
            residual = layernorm_input

Vijay Korthikanti's avatar
Vijay Korthikanti committed
671
        if self.drop_path is None:
672
673
674
675
676
677
678
679
680
681
682
683
            # re-enable torch grad to enable fused optimization.
            with torch.enable_grad():
                output = bias_dropout_add_func(
                    mlp_output,
                    mlp_bias.expand_as(residual),
                    residual,
                    self.hidden_dropout)
        else:
            out = torch.nn.functional.dropout(mlp_output + mlp_bias,
                                              p=self.hidden_dropout,
                                              training=self.training)
            output = residual + self.drop_path(out)
684
685
686
687

        return output


688
689
690
class NoopTransformerLayer(MegatronModule):
    """A single 'no-op' transformer layer.

Lawrence McAfee's avatar
Lawrence McAfee committed
691
    The sole purpose of this layer is for when a standalone embedding layer
692
    is used (i.e., args.standalone_embedding_stage == True). In this case,
Lawrence McAfee's avatar
Lawrence McAfee committed
693
694
695
696
697
698
699
700
701
    zero transformer layers are assigned when pipeline rank == 0. Additionally,
    when virtual pipeline rank >= 1, zero total model parameters are created
    (virtual rank 0 contains the input embedding). This results in the model's
    input and output tensors being the same, which causes an error when
    performing certain memory optimiations on the output tensor (e.g.,
    deallocating it). Thus, this layer disconnects the input from the output
    via a clone. Since ranks containing a no-op layer are generally under-
    utilized (both compute and memory), there's no worry of any performance
    degredation.
702
703
704
705
706
707
708
709
710
711
712
713
    """

    def __init__(self, layer_number):
        super().__init__()
        self.layer_number = layer_number

    def forward(self, hidden_states, attention_mask,
                encoder_output=None, enc_dec_attn_mask=None,
                inference_params=None):
        return hidden_states.clone()


714
715
716
class ParallelTransformer(MegatronModule):
    """Transformer class."""

717
    def __init__(self, init_method, output_layer_init_method,
718
                 layer_type=LayerType.encoder,
719
                 self_attn_mask_type=AttnMaskType.padding,
720
721
                 pre_process=True, post_process=True,
                 drop_path_rate=0.0):
722
        super(ParallelTransformer, self).__init__()
Mohammad's avatar
Mohammad committed
723
        args = get_args()
724

725
726
        self.layer_type = layer_type
        self.model_type = args.model_type
Mohammad Shoeybi's avatar
Mohammad Shoeybi committed
727
        self.bf16 = args.bf16
728
        self.fp32_residual_connection = args.fp32_residual_connection
729
730
731
        self.pre_process = pre_process
        self.post_process = post_process
        self.input_tensor = None
732
        self.drop_path_rate = drop_path_rate
733

734
        # Store activation checkpoiting flag.
735
736
        self.activations_checkpoint_method = args.activations_checkpoint_method
        self.activations_checkpoint_num_layers = args.activations_checkpoint_num_layers
737
738
        self.distribute_checkpointed_activations = \
            args.distribute_checkpointed_activations and not args.model_parallel_memory_opt
739

740
741
        self.model_parallel_memory_opt = args.model_parallel_memory_opt

742
        # Number of layers.
743
744
        self.num_layers = mpu.get_num_layers(
            args, args.model_type == ModelType.encoder_and_decoder)
Mohammad's avatar
Mohammad committed
745

Vijay Korthikanti's avatar
Vijay Korthikanti committed
746
        self.drop_path_rates = [rate.item() for rate in torch.linspace(0, self.drop_path_rate, args.num_layers)]
747

Mohammad's avatar
Mohammad committed
748
749
        # Transformer layers.
        def build_layer(layer_number):
750
            return ParallelTransformerLayer(
751
752
753
                init_method,
                output_layer_init_method,
                layer_number,
754
                layer_type=layer_type,
755
                self_attn_mask_type=self_attn_mask_type,
Vijay Korthikanti's avatar
Vijay Korthikanti committed
756
                drop_path_rate=self.drop_path_rates[layer_number - 1])
757
758
        if args.virtual_pipeline_model_parallel_size is not None:
            assert args.num_layers % args.virtual_pipeline_model_parallel_size == 0, \
759
760
                'num_layers_per_stage must be divisible by ' \
                'virtual_pipeline_model_parallel_size'
Vijay Korthikanti's avatar
Vijay Korthikanti committed
761
            assert args.model_type != ModelType.encoder_and_decoder
762
763
            # Number of layers in each model chunk is the number of layers in the stage,
            # divided by the number of model chunks in a stage.
764
            self.num_layers = self.num_layers // args.virtual_pipeline_model_parallel_size
765
766
767
768
769
770
771
772
            # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of
            # layers to stages like (each list is a model chunk):
            # Stage 0: [0]  [2]  [4]  [6]
            # Stage 1: [1]  [3]  [5]  [7]
            # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
            # layers to stages like (each list is a model chunk):
            # Stage 0: [0, 1]  [4, 5]
            # Stage 1: [2, 3]  [6, 7]
773
            offset = mpu.get_virtual_pipeline_model_parallel_rank() * (
774
                args.num_layers // args.virtual_pipeline_model_parallel_size) + \
775
776
                (mpu.get_pipeline_model_parallel_rank() * self.num_layers)
        else:
777
            # Each stage gets a contiguous set of layers.
Vijay Korthikanti's avatar
Vijay Korthikanti committed
778
779
            if args.model_type == ModelType.encoder_and_decoder and \
                    mpu.get_pipeline_model_parallel_world_size() > 1:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
780
781
782
783
784
785
786
787
                pipeline_rank = mpu.get_pipeline_model_parallel_rank()
                if layer_type == LayerType.encoder:
                    offset = pipeline_rank * self.num_layers
                else:
                    num_ranks_in_enc = args.pipeline_model_parallel_split_rank
                    offset = (pipeline_rank - num_ranks_in_enc) * self.num_layers
            else:
                offset = mpu.get_pipeline_model_parallel_rank() * self.num_layers
788

789
        if self.num_layers == 0:
Lawrence McAfee's avatar
Lawrence McAfee committed
790
            # When a standalone embedding stage is used (e.g.,
791
            # args.standalone_embedding_stage == True), virtual pipeline ranks
792
            # on pipeline rank 0 will have zero transformer layers assigned to
Lawrence McAfee's avatar
Lawrence McAfee committed
793
794
795
796
797
            # them. This results in the model's input and output tensors to be
            # the same, which will cause failure for certain output tensor
            # optimizations (e.g., pipeline output deallocation). To remedy
            # this, we assign a 'no-op' layer on these ranks, which will
            # disconnect the input tensor from the output tensor.
798
799
800
801
802
            self.num_layers = 1
            self.layers = torch.nn.ModuleList([ NoopTransformerLayer(1) ])
        else:
            self.layers = torch.nn.ModuleList(
                [build_layer(i + 1 + offset) for i in range(self.num_layers)])
803

804
        if self.post_process:
805
806
807
            # Final layer norm before output.
            self.final_layernorm = LayerNorm(
                args.hidden_size,
Sangkug Lym's avatar
Sangkug Lym committed
808
                eps=args.layernorm_epsilon,
809
810
                no_persist_layer_norm=args.no_persist_layer_norm,
                sequence_parallel=args.model_parallel_memory_opt)
811

Mohammad's avatar
Mohammad committed
812
    def _get_layer(self, layer_number):
813
        return self.layers[layer_number]
Mohammad's avatar
Mohammad committed
814

815
816
    def _checkpointed_forward(self, hidden_states, attention_mask,
                              encoder_output, enc_dec_attn_mask):
817
818
819
820
        """Forward method with activation checkpointing."""
        def custom(start, end):
            def custom_forward(*inputs):
                x_ = inputs[0]
821
822
823
                attention_mask = inputs[1]
                encoder_output = inputs[2]
                enc_dec_attn_mask = inputs[3]
Mohammad's avatar
Mohammad committed
824
825
                for index in range(start, end):
                    layer = self._get_layer(index)
826
                    x_ = layer(x_, attention_mask, encoder_output, enc_dec_attn_mask)
827
828
829
                return x_
            return custom_forward

830
831
832
833
834
835
836
837
        if self.activations_checkpoint_method == 'uniform':
            # Uniformly divide the total number of Transformer layers and checkpoint
            # the input activation of each divided chunk.
            # A method to further reduce memory usage reducing checkpoints.
            l = 0
            while l < self.num_layers:
                hidden_states = mpu.checkpoint(
                    custom(l, l + self.activations_checkpoint_num_layers),
838
                    self.distribute_checkpointed_activations,
839
840
841
842
843
844
845
846
847
848
                    hidden_states, attention_mask, encoder_output, enc_dec_attn_mask)
                l += self.activations_checkpoint_num_layers
        elif self.activations_checkpoint_method == 'block':
            # Checkpoint the input activation of only a set number of individual
            # Transformer layers and skip the rest.
            # A method fully use the device memory removing redundant re-computation.
            for l in range(self.num_layers):
                if l < self.activations_checkpoint_num_layers:
                    hidden_states = mpu.checkpoint(
                        custom(l, l + 1),
849
                        self.distribute_checkpointed_activations,
850
851
852
853
854
855
                        hidden_states, attention_mask, encoder_output, enc_dec_attn_mask)
                else:
                    hidden_states = custom(l, l + 1)(
                        hidden_states, attention_mask, encoder_output, enc_dec_attn_mask)
        else:
            raise ValueError("Invalid activation checkpoint method.")
856
857
858

        return hidden_states

859
    def set_input_tensor(self, input_tensor):
860
861
862
863
864
865
866
        """Set input tensor to be used instead of forward()'s input.

        When doing pipeline parallelism the input from the previous
        stage comes from communication, not from the input, so the
        model's forward_step_func won't have it. This function is thus
        used by internal code to bypass the input provided by the
        forward_step_func"""
867
868
        self.input_tensor = input_tensor

869
    def forward(self, hidden_states, attention_mask,
mshoeybi's avatar
mshoeybi committed
870
871
                encoder_output=None, enc_dec_attn_mask=None,
                inference_params=None):
872

873
        # Checks.
mshoeybi's avatar
mshoeybi committed
874
        if inference_params:
875
            assert self.activations_checkpoint_method is None, \
876
                'inference does not work with activation checkpointing'
877

878
        if self.pre_process:
879
            # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
mshoeybi's avatar
mshoeybi committed
880
            # If the input flag for fp32 residual connection is set, convert for float.
881
882
            if self.fp32_residual_connection:
                hidden_states = hidden_states.transpose(0, 1).contiguous().float()
mshoeybi's avatar
mshoeybi committed
883
            # Otherwise, leave it as is.
884
885
            else:
                hidden_states = hidden_states.transpose(0, 1).contiguous()
886
887

            if self.model_parallel_memory_opt:
888
                hidden_states = mpu.scatter_to_sequence_parallel_region(hidden_states)
889

890
        else:
891
            # See set_input_tensor()
892
            hidden_states = self.input_tensor
893

894
895
        # Viewless tensor.
        # - We only need to create a viewless tensor in the case of micro batch
896
897
898
899
        #   size (mbs) == 1, since in this case, 'hidden_states.transpose()'
        #   above creates a view tensor, and '.contiguous()' is a pass-through.
        #   For mbs >= 2, '.contiguous()' creates a new tensor, eliminating
        #   the need to make it viewless.
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
        #
        #   However, we don't explicitly check mbs == 1 here because
        #   make_viewless_tensor() has negligible overhead when its input
        #   is already viewless.
        # 
        # - For the 'else' case above, calling make_viewless_tensor() here is
        #   likely redundant, since p2p_communication.py (likely originator)
        #   already creates viewless tensors. That said, make_viewless_tensor()
        #   is called here to be future-proof and corner-case-proof.
        hidden_states = mpu.make_viewless_tensor(
            hidden_states,
            requires_grad = True,
            keep_graph = True,
        )

        # Transpose encoder output.
916
917
        if encoder_output is not None and \
                not self.model_parallel_memory_opt:
918
            encoder_output = encoder_output.transpose(0, 1).contiguous()
919
920
921
            if self.model_parallel_memory_opt:
                encoder_output = mpu.scatter_to_sequence_parallel_region(encoder_output)

Vijay Korthikanti's avatar
Vijay Korthikanti committed
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
        if self.model_parallel_memory_opt:
            with mpu.get_cuda_rng_tracker().fork():
                # Forward pass.
                if self.activations_checkpoint_method is not None:
                    hidden_states = self._checkpointed_forward(hidden_states,
                                                               attention_mask,
                                                               encoder_output,
                                                               enc_dec_attn_mask)
                else:
                    for index in range(self.num_layers):
                        layer = self._get_layer(index)
                        hidden_states = layer(
                            hidden_states,
                            attention_mask,
                            encoder_output=encoder_output,
                            enc_dec_attn_mask=enc_dec_attn_mask,
                            inference_params=inference_params)
939
        else:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
            # Forward pass.
            if self.activations_checkpoint_method is not None:
                hidden_states = self._checkpointed_forward(hidden_states,
                                                           attention_mask,
                                                           encoder_output,
                                                           enc_dec_attn_mask)
            else:
                for index in range(self.num_layers):
                    layer = self._get_layer(index)
                    hidden_states = layer(
                        hidden_states,
                        attention_mask,
                        encoder_output=encoder_output,
                        enc_dec_attn_mask=enc_dec_attn_mask,
                        inference_params=inference_params)
mshoeybi's avatar
mshoeybi committed
955

956
        # Final layer norm.
957
        if self.post_process:
958
            # Reverting data format change [s b h] --> [b s h].
959
960
            hidden_states = self.final_layernorm(hidden_states)

Vijay Korthikanti's avatar
Vijay Korthikanti committed
961
962
            if self.layer_type == LayerType.encoder and \
                    self.model_type == ModelType.encoder_and_decoder and \
963
964
965
                    self.model_parallel_memory_opt:
                output = hidden_states
            else:
Vijay Korthikanti's avatar
Vijay Korthikanti committed
966

967
968
                if self.model_parallel_memory_opt:
                    hidden_states = mpu.gather_from_sequence_parallel_region(hidden_states)
969

970
                output = hidden_states.transpose(0, 1).contiguous()
971
972
        else:
            output = hidden_states
973

974
        return output