transformer.py 29.9 KB
Newer Older
1
# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
Przemek Tredak's avatar
Przemek Tredak committed
2
3
4
5
6
#
# See LICENSE for license information.

"""Transformer."""
import os
7
import warnings
Przemek Tredak's avatar
Przemek Tredak committed
8
from contextlib import nullcontext
9
from typing import Any, Callable, List, Optional, Tuple, Union
Przemek Tredak's avatar
Przemek Tredak committed
10
11
12

import torch

13
import transformer_engine_extensions as tex
14
from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm
15
from transformer_engine.pytorch.attention import MultiheadAttention
Przemek Tredak's avatar
Przemek Tredak committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
from transformer_engine.pytorch.jit import (
    set_jit_fusion_options,
    warmup_jit_bias_dropout_add_all_dtypes,
    get_bias_dropout_add,
    bias_dropout_add_fused_train,
    bias_dropout_add_fused_inference,
)
from transformer_engine.pytorch.utils import (
    cast_if_needed,
    get_default_init_method,
)
from transformer_engine.pytorch.constants import (
    AttnMaskTypes,
    LayerTypes,
    dist_group_type,
)
32
33
from transformer_engine.pytorch.distributed import get_distributed_world_size

Przemek Tredak's avatar
Przemek Tredak committed
34

35
warnings.filterwarnings("module", category=DeprecationWarning, module="transformer")
cyanguwa's avatar
cyanguwa committed
36
37


38
__all__ = ["TransformerLayer"]
cyanguwa's avatar
cyanguwa committed
39

Przemek Tredak's avatar
Przemek Tredak committed
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

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

    def __init__(self, drop_prob: float = 0.0) -> None:
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        """DropPath FWD"""

        if self.drop_prob == 0.0 or not self.training:
            return hidden_state
        keep_prob = 1 - self.drop_prob
        # work with diff dim tensors, not just 2D ConvNets
        shape = (hidden_state.shape[0],) + (1,) * (hidden_state.ndim - 1)
        random_tensor = keep_prob + torch.rand(
            shape, dtype=hidden_state.dtype, device=hidden_state.device
        )
        random_tensor.floor_()  # binarize
        output = hidden_state.div(keep_prob) * random_tensor
        return output


class TransformerLayer(torch.nn.Module):
67
    r"""
Przemek Tredak's avatar
Przemek Tredak committed
68
69
70
    TransformerLayer is made up of an attention block and a feedforward network (MLP).
    This standard layer is based on the paper "Attention Is All You Need".

71
72
73
    .. warning::

        Arguments :attr:`attention_softmax_in_fp32` and :attr:`apply_query_key_layer_scaling`
74
        are deprecated and will be fully removed in the next release (v1.0.0).
75

76
    .. note::
77

78
79
        Argument :attr:`attention_mask` will be ignored in the `forward` call when
        :attr:`self_attn_mask_type` is set to `"causal"`.
80

Przemek Tredak's avatar
Przemek Tredak committed
81
82
83
84
85
86
87
88
    Parameters
    ----------
    hidden_size : int
                 size of each input sample.
    ffn_hidden_size : int
                     intermediate size to which input samples are projected.
    num_attention_heads : int
                         number of attention heads in the transformer layer.
89
90
91
92
93
94
95
96
    num_gqa_groups : int, default = `None`
                         number of GQA groups in the transformer layer.
                         Grouped Query Attention is described in
                         `this paper <https://arxiv.org/pdf/2305.13245.pdf>`_.
                         This only affects the keys and values, not the querys.
                         GQA-1 is equivalent to Multi-Query Attention
                         (`MQA <https://arxiv.org/pdf/1911.02150.pdf>`_), while GQA-H
                         is equivalent to MHA, i.e. `num_gqa_groups = num_attention_heads`.
Przemek Tredak's avatar
Przemek Tredak committed
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
    layernorm_epsilon : float, default = 1e-5
                       a value added to the denominator of layer normalization
                       for numerical stability.
    hidden_dropout: float, default = 0.1
                   dropout probability for the dropout op after FC2 layer.
    attention_dropout: float, default = 0.1
                      dropout probability for the dropout op during multi-head attention.
    init_method : Callable, default = `None`
                 used for initializing weights of QKV and FC1 weights in the following way:
                 `init_method(weight)`. When set to `None`, defaults to
                 `torch.nn.init.normal_(mean=0.0, std=0.023)`.
    output_layer_init_method : Callable, default = `None`
                              used for initializing weights of PROJ and FC2 in the following way:
                              `output_layer_init_method(weight)`. When set to `None`, defaults to
                              `torch.nn.init.normal_(mean=0.0, std=0.023)`.
    apply_residual_connection_post_layernorm : bool, default = `False`
                                              if set to `True`, residual connections are taken
                                              from the output of layer norm (default is taken
                                              from input of layer norm)
    layer_number: int, default = `None`
                 layer number of the current `TransformerLayer` when multiple such modules are
                 concatenated to form a transformer block.
    output_layernorm: bool, default = `False`
                     if set to `True`, layer normalization is applied on the output side,
                     after the final dropout-add. default behavior is to apply layer
                     normalization on the input side, before the QKV transformation.
    layer_type: {'encoder', 'decoder'}, default = `encoder`
               if set to `decoder`, an additional cross-attn block is added after self-attn.
               This can be used for structures like `T5` Transformer in conjunction with the
               `encoder` option.
    kv_channels: int, default = `None`
                number of key-value channels. defaults to
                :attr:`hidden_size` / :attr:`num_attention_heads` if `None`.
130
    self_attn_mask_type: {'causal', 'padding', 'no_mask', 'arbitrary'}, default = `causal`
131
132
133
134
135
                        type of attention mask passed into softmax operation. Overridden by
                        :attr:`self_attn_mask_type` in the `forward` method. The forward
                        arg is useful for dynamically changing mask types, e.g. a different
                        mask for training and inference. The init arg is useful for cases
                        involving compilation/tracing, e.g. ONNX export.
136
137
138
139
140
141
142
    zero_centered_gamma : bool, default = 'False'
                         if set to 'True', gamma parameter in LayerNorm is initialized to 0 and
                         the LayerNorm formula changes to

                         .. math::
                            y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \varepsilon}} *
                            (1 + \gamma) + \beta
143
144
    normalization : { 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm'
                   type of normalization applied.
145
146
147
148
149
150
    qkv_weight_interleaved : bool, default = `True`
                            if set to `False`, the QKV weight is interpreted as a concatenation of
                            query, key, and value weights along the `0th` dimension. The default
                            interpretation is that the individual `q`, `k`, and `v` weights for each
                            attention head are interleaved. This parameter is set to `False` when
                            using :attr:`fuse_qkv_params=False`.
ngoyal2707's avatar
ngoyal2707 committed
151
152
    bias : bool, default = `True`
          if set to `False`, the transformer layer will not learn any additive biases.
153
154
155
    activation : str, default = 'gelu'
          Type of activation used in MLP block.
          Options are: 'gelu', 'relu', 'reglu', 'geglu' and 'swiglu'.
156
157
158
159
    device : Union[torch.device, str], default = "cuda"
          The device on which the parameters of the model will allocated. It is the user's
          responsibility to ensure all parameters are moved to the GPU before running the
          forward pass.
ngoyal2707's avatar
ngoyal2707 committed
160

Przemek Tredak's avatar
Przemek Tredak committed
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
    Parallelism parameters
    ----------------------
    set_parallel_mode : bool, default = `False`
                      if set to `True`, QKV and FC1 layers are used as Column Parallel
                      whereas PROJ and FC2 is used as Row Parallel as described
                      `here <https://arxiv.org/pdf/1909.08053.pdf>`_.
    sequence_parallel : bool, default = `False`
                       if set to `True`, uses sequence parallelism.
    tp_group : ProcessGroup, default = `None`
              tensor parallel process group.
    tp_size : int, default = 1
             used as TP (tensor parallel) world size when TP groups are not formed during
             initialization. In this case, users must call the
             `set_tensor_parallel_group(tp_group)` method on the initialized module before the
             forward pass to supply the tensor parallel group needed for tensor and sequence
             parallel collectives.

    Optimization parameters
    -----------------------
    fuse_wgrad_accumulation : bool, default = 'False'
                             if set to `True`, enables fusing of creation and accumulation of
182
183
184
185
                             the weight gradient. When enabled, it is assumed that the weights
                             have an additional `main_grad` attribute (used instead of the
                             regular `grad`) which is a pre-allocated buffer of the correct
                             size to accumulate gradients in.
186
    params_dtype : torch.dtype, default = `torch.get_default_dtype()`
Przemek Tredak's avatar
Przemek Tredak committed
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
                  it controls the type used to allocate the initial parameters. Useful when
                  the model is trained with lower precision and the original FP32 parameters
                  would not fit in GPU memory.
    seq_length: int
               sequence length of input samples. Needed for JIT Warmup, a technique where jit
               fused functions are warmed up before training to ensure same kernels are used for
               forward propogation and activation recompute phase.
    micro_batch_size: int
                     batch size per training step. Needed for JIT Warmup, a technique where jit
                     fused functions are warmed up before training to ensure same kernels are
                     used for forward propogation and activation recompute phase.
    drop_path_rate: float, default = 0.0
                   when > 0.0, applies stochastic depth per sample in
                   the main path of the residual block.
    fuse_qkv_params: bool, default = 'False'
                    if set to `True`, `TransformerLayer` module exposes a single fused
                    parameter for query-key-value. This enables optimizations such as QKV
                    fusion without concatentations/splits and also enables the argument
                    `fuse_wgrad_accumulation`.
    """

    def __init__(
        self,
        hidden_size: int,
        ffn_hidden_size: int,
        num_attention_heads: int,
213
        num_gqa_groups: Optional[int] = None,
Przemek Tredak's avatar
Przemek Tredak committed
214
215
216
217
218
219
220
        layernorm_epsilon: float = 1e-5,
        hidden_dropout: float = 0.1,
        attention_dropout: float = 0.1,
        init_method: Optional[Callable] = None,
        output_layer_init_method: Optional[Callable] = None,
        layer_number: Optional[int] = None,
        kv_channels: Optional[int] = None,
221
        self_attn_mask_type: str = "causal",
Przemek Tredak's avatar
Przemek Tredak committed
222
223
        tp_group: Optional[dist_group_type] = None,
        tp_size: int = 1,
224
        params_dtype: Optional[torch.dtype] = None,
Przemek Tredak's avatar
Przemek Tredak committed
225
226
        get_rng_state_tracker: Optional[Callable] = None,
        fuse_wgrad_accumulation: bool = False,
227
228
        apply_query_key_layer_scaling: bool = False, # pylint: disable=unused-argument
        attention_softmax_in_fp32: bool = True, # pylint: disable=unused-argument
Przemek Tredak's avatar
Przemek Tredak committed
229
230
231
232
233
234
235
236
237
        seq_length: Optional[int] = None,
        micro_batch_size: Optional[int] = None,
        sequence_parallel: bool = False,
        apply_residual_connection_post_layernorm: bool = False,
        output_layernorm: bool = False,
        layer_type: str = "encoder",
        drop_path_rate: float = 0.0,
        set_parallel_mode: bool = False,
        fuse_qkv_params: bool = False,
238
        zero_centered_gamma: bool = False,
239
        qkv_weight_interleaved: bool = True,
240
        ub_tp_comm_overlap: bool = False,
ngoyal2707's avatar
ngoyal2707 committed
241
        bias: bool = True,
242
243
        activation: str = 'gelu',
        normalization: str = "LayerNorm",
244
        device: Union[torch.device, str] = "cuda",
Przemek Tredak's avatar
Przemek Tredak committed
245
246
247
    ) -> None:
        super().__init__()

248
249
        warnings.warn(
            "Arguments `attention_softmax_in_fp32` and `apply_query_key_layer_scaling`"
250
            "are deprecated and will be fully removed in the next release (v1.0.0).",
251
252
253
            category=DeprecationWarning,
        )

254
255
256
257
258
        if ub_tp_comm_overlap:
            assert (
                tex.userbuf_comm_available()
            ), "Userbuffer communication backend not available."

259
        self.self_attn_mask_type = self_attn_mask_type
260
        params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype
261
262
263
264
265
        ub_tp_comm_overlap = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_OVERLAP", "1")))
        ub_bulk_wgrad = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_BULK_WGRAD", "1")))
        ub_bulk_dgrad = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_BULK_DGRAD", "1")))
        ub_split_ag = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_SPLIT_AG", "1")))
        ub_split_rs = ub_tp_comm_overlap and bool(int(os.getenv("NVTE_UB_SPLIT_RS", "1")))
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
        ub_atomic_gemm_rs = (ub_tp_comm_overlap
                             and bool(int(os.getenv("NVTE_UB_ATOMIC_GEMM_RS", "0"))))
        assert (
            not (ub_split_rs and ub_atomic_gemm_rs)
        ), "Only one type of RS overlap NVTE_UB_SPLIT_RS/NVTE_UB_ATOMIC_GEMM_RS should be enabled."
        ub_atomic_gemm_ag = (ub_tp_comm_overlap
                             and bool(int(os.getenv("NVTE_UB_ATOMIC_GEMM_AG", "0"))))
        assert (
            not (ub_split_ag and ub_atomic_gemm_ag)
        ), "Only one type of AG overlap NVTE_UB_SPLIT_AG/NVTE_UB_ATOMIC_GEMM_AG should be enabled."

        if ub_atomic_gemm_rs or ub_atomic_gemm_ag:
            warnings.warn(
                "Atomic gemm uses a beta API from cublas and is not tested for all use cases."
            )

Przemek Tredak's avatar
Przemek Tredak committed
282
283
284
285
286
287
288
        bias_dropout_fusion = bool(int(os.getenv("NVTE_BIAS_DROPOUT_FUSION", "1")))
        self.layer_number = layer_number
        self.output_layernorm = output_layernorm
        self.layer_type = layer_type
        self.apply_residual_connection_post_layernorm = (
            apply_residual_connection_post_layernorm
        )
289

Przemek Tredak's avatar
Przemek Tredak committed
290
291
292
293
294
295
296
        assert layer_type in LayerTypes, f"layer_type {layer_type} not supported"

        if not fuse_qkv_params:
            assert (
                not fuse_wgrad_accumulation
            ), "Gradient accumulation fusion requires single QKV parameter."

297
298
299
        if not fuse_qkv_params:
            qkv_weight_interleaved = False

Przemek Tredak's avatar
Przemek Tredak committed
300
301
302
303
304
305
306
307
308
        self.kv_channels = (
            kv_channels if kv_channels else (hidden_size // num_attention_heads)
        )

        if init_method is None:
            init_method = get_default_init_method()
        if output_layer_init_method is None:
            output_layer_init_method = get_default_init_method()

309
310
311
        self.tp_size = tp_size if tp_group is None else get_distributed_world_size(tp_group)
        self.sequence_parallel = (self.tp_size > 1) and sequence_parallel
        self.seq_length = seq_length
Przemek Tredak's avatar
Przemek Tredak committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326

        self.get_rng_state_tracker = get_rng_state_tracker

        attention_args = (
            hidden_size,
            num_attention_heads,
            self.kv_channels,
            attention_dropout,
            layernorm_epsilon,
            init_method,
            output_layer_init_method,
        )
        common_attention_kwargs = {
            "layer_number": layer_number,
            "tp_group": tp_group,
327
            "tp_size": self.tp_size,
328
            "num_gqa_groups": num_gqa_groups,
Przemek Tredak's avatar
Przemek Tredak committed
329
330
331
332
333
334
335
            "fuse_wgrad_accumulation": fuse_wgrad_accumulation,
            "get_rng_state_tracker": get_rng_state_tracker,
            "sequence_parallel": self.sequence_parallel,
            "params_dtype": params_dtype,
            "return_layernorm_output": apply_residual_connection_post_layernorm,
            "set_parallel_mode": set_parallel_mode,
            "fuse_qkv_params": fuse_qkv_params,
cyanguwa's avatar
cyanguwa committed
336
            "zero_centered_gamma": zero_centered_gamma,
337
            "qkv_weight_interleaved" : qkv_weight_interleaved,
338
339
340
341
            "ub_bulk_wgrad" : ub_bulk_wgrad,
            "ub_bulk_dgrad" : ub_bulk_dgrad,
            "ub_split_ag" : ub_split_ag,
            "ub_split_rs" : ub_split_rs,
342
343
            "ub_atomic_gemm_rs" : ub_atomic_gemm_rs,
            "ub_atomic_gemm_ag" : ub_atomic_gemm_ag,
Przemek Tredak's avatar
Przemek Tredak committed
344
345
        }

346
        self.self_attention = MultiheadAttention(
Przemek Tredak's avatar
Przemek Tredak committed
347
348
349
350
            *attention_args,
            **common_attention_kwargs,
            input_layernorm=not output_layernorm,
            attention_type="self",
ngoyal2707's avatar
ngoyal2707 committed
351
            bias=bias,
352
            return_bias=True,
353
            normalization=normalization,
354
            device=device,
Przemek Tredak's avatar
Przemek Tredak committed
355
356
357
        )

        if layer_type == "decoder":
358
            self.inter_attention = MultiheadAttention(
Przemek Tredak's avatar
Przemek Tredak committed
359
360
361
362
363
                *attention_args,
                **common_attention_kwargs,
                attn_mask_type="padding",
                input_layernorm=True,
                attention_type="cross",
ngoyal2707's avatar
ngoyal2707 committed
364
                bias=bias,
365
                return_bias=True,
366
                normalization=normalization,
367
                device=device,
Przemek Tredak's avatar
Przemek Tredak committed
368
369
            )

370
        # LayerNorm -> activation(Linear + Bias) -> Linear
Przemek Tredak's avatar
Przemek Tredak committed
371
372
        # parallel_mode not supported for LayerNormMLP,
        # FC1 is CPL and FC2 is RPL
373
374
        # In the case of GLU activation, FC1 handles both
        # Linear layers before the activation
Przemek Tredak's avatar
Przemek Tredak committed
375
376
377
378
379
380
        self.layernorm_mlp = LayerNormMLP(
            hidden_size,
            ffn_hidden_size,
            eps=layernorm_epsilon,
            fuse_wgrad_accumulation=fuse_wgrad_accumulation,
            tp_group=tp_group,
381
            tp_size=self.tp_size,
Przemek Tredak's avatar
Przemek Tredak committed
382
383
384
            get_rng_state_tracker=get_rng_state_tracker,
            init_method=init_method,
            output_layer_init_method=output_layer_init_method,
ngoyal2707's avatar
ngoyal2707 committed
385
            bias=bias,
Przemek Tredak's avatar
Przemek Tredak committed
386
387
388
389
390
391
392
            return_bias=True,
            sequence_parallel=self.sequence_parallel,
            params_dtype=params_dtype,
            return_layernorm_output=apply_residual_connection_post_layernorm,
            seq_length=seq_length,
            micro_batch_size=micro_batch_size,
            set_parallel_mode=set_parallel_mode,
393
            zero_centered_gamma=zero_centered_gamma,
394
395
396
397
            ub_bulk_wgrad=ub_bulk_wgrad,
            ub_bulk_dgrad=ub_bulk_dgrad,
            ub_split_rs=ub_split_rs,
            ub_split_ag=ub_split_ag,
398
399
            ub_atomic_gemm_rs=ub_atomic_gemm_rs,
            ub_atomic_gemm_ag=ub_atomic_gemm_ag,
400
            activation=activation,
401
            normalization=normalization,
402
            device=device,
Przemek Tredak's avatar
Przemek Tredak committed
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
        )

        self.hidden_dropout = hidden_dropout
        self.bias_dropout_fusion = bias_dropout_fusion
        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None

        # Set bias+dropout+add fusion grad_enable execution handler.
        TORCH_MAJOR = int(torch.__version__.split(".")[0])
        TORCH_MINOR = int(torch.__version__.split(".")[1])
        use_nvfuser = TORCH_MAJOR > 1 or (TORCH_MAJOR == 1 and TORCH_MINOR >= 10)
        self.bias_dropout_add_exec_handler = (
            nullcontext if use_nvfuser else torch.enable_grad
        )

        if self.bias_dropout_fusion:
            set_jit_fusion_options()
            if seq_length and micro_batch_size:
                if self.sequence_parallel:
421
                    seq_length = seq_length // self.tp_size
Przemek Tredak's avatar
Przemek Tredak committed
422
423
424
425
                warmup_jit_bias_dropout_add_all_dtypes(
                    hidden_size, seq_length, micro_batch_size
                )

426
427
428
429
        norm_module = {
                "LayerNorm": LayerNorm,
                "RMSNorm": RMSNorm,
        }
Przemek Tredak's avatar
Przemek Tredak committed
430
        if self.output_layernorm:
431
            self.layernorm = norm_module[normalization](
Przemek Tredak's avatar
Przemek Tredak committed
432
433
434
435
                hidden_size,
                eps=layernorm_epsilon,
                sequence_parallel=self.sequence_parallel,
                params_dtype=params_dtype,
436
437
                zero_centered_gamma=zero_centered_gamma,
                device=device,
Przemek Tredak's avatar
Przemek Tredak committed
438
439
440
441
442
443
444
445
446
447
448
            )

    def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> None:
        """Set TP group"""
        # Deep iterate but skip self to avoid infinite recursion.
        for index, child in enumerate(self.modules()):
            if index == 0:
                continue
            if hasattr(child, "set_tensor_parallel_group"):
                child.set_tensor_parallel_group(tp_group)

449
450
451
    def set_context_parallel_running(
        self,
        cp_group: Union[dist_group_type, None],
452
        cp_global_ranks: List[int],
453
454
455
456
457
458
459
460
461
462
        cp_stream: torch.cuda.Stream,
    ) -> None:
        """Set CP group and CP dual-stream running"""
        # Deep iterate but skip self to avoid infinite recursion.
        for index, child in enumerate(self.modules()):
            if index == 0:
                continue
            if hasattr(child, "set_context_parallel_running"):
                child.set_context_parallel_running(cp_group, cp_global_ranks, cp_stream)

Przemek Tredak's avatar
Przemek Tredak committed
463
464
465
    def forward(
        self,
        hidden_states: torch.Tensor,
cyanguwa's avatar
cyanguwa committed
466
        attention_mask: Optional[torch.Tensor] = None,
467
        self_attn_mask_type: Optional[str] = None,
Przemek Tredak's avatar
Przemek Tredak committed
468
469
470
        encoder_output: Optional[torch.Tensor] = None,
        enc_dec_attn_mask: Optional[torch.Tensor] = None,
        is_first_microbatch: Optional[bool] = None,
cyanguwa's avatar
cyanguwa committed
471
        checkpoint_core_attention: bool = False,
Przemek Tredak's avatar
Przemek Tredak committed
472
        inference_params: Optional[Any] = None,
473
        rotary_pos_emb: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None,
474
475
476
        core_attention_bias_type: str = "no_bias",
        core_attention_bias: Optional[torch.Tensor] = None,
        fast_zero_fill: bool = True,
Przemek Tredak's avatar
Przemek Tredak committed
477
478
479
480
    ) -> torch.Tensor:
        """
        Transformer Layer: attention block and a feedforward network (MLP)

481
482
        .. note::

483
484
            Argument :attr:`attention_mask` is only used when :attr:`self_attn_mask_type`
            includes `"padding"` or `"arbitrary"`.
485

Przemek Tredak's avatar
Przemek Tredak committed
486
487
488
489
        Parameters
        ----------
        hidden_states : torch.Tensor
             Input tensor.
490
491
492
493
        attention_mask : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = `None`
                        Boolean tensor used to mask out self-attention softmax input.
                        Can be a tuple of 2 masks for cross attention with padding masks.
        self_attn_mask_type: {'causal', 'padding', 'no_mask', 'arbitrary'}, default = `causal`
494
                            type of attention mask passed into softmax operation.
cyanguwa's avatar
cyanguwa committed
495
        encoder_output : Optional[torch.Tensor], default = `None`
Przemek Tredak's avatar
Przemek Tredak committed
496
497
             Output of the encoder block to be fed into the decoder block if using
             `layer_type="decoder"`.
cyanguwa's avatar
cyanguwa committed
498
        enc_dec_attn_mask : Optional[torch.Tensor], default = `None`
Przemek Tredak's avatar
Przemek Tredak committed
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
             Boolean tensor used to mask out inter-attention softmax input if using
             `layer_type="decoder"`.
        is_first_microbatch : {True, False, None}, default = None
                             During training using either gradient accumulation or
                             pipeline parallelism a minibatch of data is further split
                             into microbatches. Between the microbatches of the same minibatch
                             the model weights are not updated. Setting this parameter indicates
                             whether the current microbatch is the first in a minibatch or not.
                             When set, this parameter enables additional optimizations:

                             * during FP8 training, it allows caching of the FP8 versions of
                               the weights
                             * it also allows skipping gradient accumulation during the
                               first microbatch (since it is the first gradient being
                               produced)
cyanguwa's avatar
cyanguwa committed
514
        checkpoint_core_attention: bool, default = `False`
Przemek Tredak's avatar
Przemek Tredak committed
515
516
517
518
                                  If true, forward activations for core attention are recomputed
                                  during the backward pass in order to save memory that would
                                  otherwise be occupied to store the forward activations until
                                  backprop.
519
520
521
        rotary_pos_emb: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = `None`
                       Embeddings for query and key tensors for applying rotary position
                       embedding. By default no input embedding is applied.
522
523
524
525
526
527
        core_attention_bias_type: str, default = `no_bias`
                    Bias type, {`no_bias`, `pre_scale_bias`, 'post_scale_bias`}
        core_attention_bias: Optional[torch.Tensor], default = `None`
                    Bias tensor for Q * K.T
        fast_zero_fill: bool, default = `True`
                    Whether to set output tensors to 0 or not before use.
Przemek Tredak's avatar
Przemek Tredak committed
528
529
        """

530
        if self_attn_mask_type is None:
531
532
533
534
535
536
            self_attn_mask_type = self.self_attn_mask_type

        assert (
            self_attn_mask_type in AttnMaskTypes
        ), f"self_attn_mask_type {self_attn_mask_type} not supported"

537
538
        hidden_states = hidden_states.contiguous()

539
540
541
542
543
        if self.sequence_parallel and self.seq_length is not None:
            assert (
                hidden_states.shape[0] == self.seq_length // self.tp_size
            ), "Sequence dimension must be split across TP group when using sequence parallel."

544
        if self_attn_mask_type != "causal" and attention_mask is not None:
545
546
547
548
            assert (
                attention_mask.dtype == torch.bool
            ), "Attention mask must be a boolean tensor"

Przemek Tredak's avatar
Przemek Tredak committed
549
550
551
552
553
554
555
556
557
        # For AMP
        if torch.is_autocast_enabled():
            hidden_states = cast_if_needed(
                hidden_states, torch.get_autocast_gpu_dtype()
            )

        # Self attention.
        self_attention_outputs = self.self_attention(
            hidden_states,
558
559
            attention_mask=attention_mask,
            attn_mask_type=self_attn_mask_type,
Przemek Tredak's avatar
Przemek Tredak committed
560
561
562
            inference_params=inference_params,
            is_first_microbatch=is_first_microbatch,
            checkpoint_core_attention=checkpoint_core_attention,
563
            rotary_pos_emb=rotary_pos_emb,
564
565
566
            core_attention_bias_type=core_attention_bias_type,
            core_attention_bias=core_attention_bias,
            fast_zero_fill=fast_zero_fill,
Przemek Tredak's avatar
Przemek Tredak committed
567
        )
ngoyal2707's avatar
ngoyal2707 committed
568

Przemek Tredak's avatar
Przemek Tredak committed
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
        if self.apply_residual_connection_post_layernorm and not self.output_layernorm:
            attention_output, attention_bias, residual = self_attention_outputs
        else:
            attention_output, attention_bias = self_attention_outputs
            residual = hidden_states

        # Set BDA func.
        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
        else:
            bias_dropout_add_func = get_bias_dropout_add(self.training)

        # Bias dropoout add.
ngoyal2707's avatar
ngoyal2707 committed
585
        if self.drop_path is None and attention_bias.numel() != 0:
Przemek Tredak's avatar
Przemek Tredak committed
586
587
588
589
590
            with self.bias_dropout_add_exec_handler():
                bda_output = bias_dropout_add_func(
                    attention_output, attention_bias, residual, self.hidden_dropout
                )
        else:
ngoyal2707's avatar
ngoyal2707 committed
591
592
            if attention_bias.numel() != 0:
                attention_output = attention_output + attention_bias
Przemek Tredak's avatar
Przemek Tredak committed
593
            out = torch.nn.functional.dropout(
ngoyal2707's avatar
ngoyal2707 committed
594
                attention_output,
Przemek Tredak's avatar
Przemek Tredak committed
595
596
597
                p=self.hidden_dropout,
                training=self.training,
            )
ngoyal2707's avatar
ngoyal2707 committed
598
599
600
            if self.drop_path is not None:
                out = self.drop_path(out)
            bda_output = residual + out
Przemek Tredak's avatar
Przemek Tredak committed
601
602
603
604
605

        # Cross attention.
        if self.layer_type == "decoder":
            inter_attention_outputs = self.inter_attention(
                bda_output,
606
607
                attention_mask=enc_dec_attn_mask,
                attn_mask_type=self_attn_mask_type,
Przemek Tredak's avatar
Przemek Tredak committed
608
609
610
                encoder_output=encoder_output,
                is_first_microbatch=is_first_microbatch,
                checkpoint_core_attention=checkpoint_core_attention,
611
612
613
                core_attention_bias_type=core_attention_bias_type,
                core_attention_bias=core_attention_bias,
                fast_zero_fill=fast_zero_fill,
Przemek Tredak's avatar
Przemek Tredak committed
614
615
616
617
618
619
620
            )
            if self.apply_residual_connection_post_layernorm:
                attention_output, attention_bias, residual = inter_attention_outputs
            else:
                attention_output, attention_bias = inter_attention_outputs
                residual = bda_output

ngoyal2707's avatar
ngoyal2707 committed
621
622
623
624
625
626
627
628
629
630
            if attention_bias.numel() != 0:
                with self.bias_dropout_add_exec_handler():
                    bda_output = bias_dropout_add_func(
                        attention_output, attention_bias, residual, self.hidden_dropout
                    )
            else:
                out = torch.nn.functional.dropout(
                    attention_output,
                    p=self.hidden_dropout,
                    training=self.training,
Przemek Tredak's avatar
Przemek Tredak committed
631
                )
ngoyal2707's avatar
ngoyal2707 committed
632
                bda_output = residual + out
Przemek Tredak's avatar
Przemek Tredak committed
633
634
635
636
637
638
639
640
641
642
643
        # MLP.
        mlp_outputs = self.layernorm_mlp(
            bda_output, is_first_microbatch=is_first_microbatch
        )
        if self.apply_residual_connection_post_layernorm:
            mlp_output, mlp_bias, residual = mlp_outputs
        else:
            mlp_output, mlp_bias = mlp_outputs
            residual = bda_output

        # Bias dropoout add.
ngoyal2707's avatar
ngoyal2707 committed
644
        if self.drop_path is None and mlp_bias.numel() != 0:
Przemek Tredak's avatar
Przemek Tredak committed
645
646
647
648
649
            with self.bias_dropout_add_exec_handler():
                output = bias_dropout_add_func(
                    mlp_output, mlp_bias, residual, self.hidden_dropout
                )
        else:
ngoyal2707's avatar
ngoyal2707 committed
650
651
            if mlp_bias.numel() != 0:
                mlp_output = mlp_output + mlp_bias
Przemek Tredak's avatar
Przemek Tredak committed
652
            out = torch.nn.functional.dropout(
ngoyal2707's avatar
ngoyal2707 committed
653
                mlp_output, p=self.hidden_dropout, training=self.training
Przemek Tredak's avatar
Przemek Tredak committed
654
            )
ngoyal2707's avatar
ngoyal2707 committed
655
656
657
            if self.drop_path is not None:
                out = self.drop_path(out)
            output = residual + out
Przemek Tredak's avatar
Przemek Tredak committed
658
659
660
661
662

        # For BERT like architectures.
        if self.output_layernorm:
            output = self.layernorm(output)

663
        # output: [s, b, h]
Przemek Tredak's avatar
Przemek Tredak committed
664
        return output