File size: 3,113 Bytes
e56d7ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
# Copyright 2022 MosaicML Examples authors
# SPDX-License-Identifier: Apache-2.0
"""GPT Blocks used for the GPT Model."""
from typing import Optional, Tuple
import torch
import torch.nn as nn
from .attention import MultiheadAttention
from .low_precision_layernorm import LPLayerNorm
class GPTMLP(nn.Module):
def __init__(self,
d_model: int,
mlp_ratio: int,
device: Optional[str] = None):
super().__init__()
self.mlp_up = nn.Linear(d_model, mlp_ratio * d_model, device=device)
self.mlp_act = nn.GELU(approximate='none')
self.mlp_down = nn.Linear(mlp_ratio * d_model, d_model, device=device)
self.mlp_down._is_residual = True # type: ignore
def forward(self, x):
return self.mlp_down(self.mlp_act(self.mlp_up(x)))
class GPTBlock(nn.Module):
def __init__(self,
attn_impl: str,
d_model: int,
n_heads: int,
mlp_ratio: int,
attn_clip_qkv: Optional[float] = None,
attn_qk_ln: bool = False,
softmax_scale: Optional[float] = None,
attn_pdrop: float = 0.0,
alibi: bool = False,
resid_pdrop: float = 0.0,
low_precision_layernorm: bool = False,
device: Optional[str] = None,
**kwargs):
del kwargs # unused, just to capture any extra args from the config
super().__init__()
layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
self.ln_1 = layernorm_class(d_model, device=device)
self.attn = MultiheadAttention(
attn_impl=attn_impl,
attn_clip_qkv=attn_clip_qkv,
attn_qk_ln=attn_qk_ln,
softmax_scale=softmax_scale,
attn_pdrop=attn_pdrop,
d_model=d_model,
n_heads=n_heads,
device=device,
)
self.ln_2 = layernorm_class(d_model, device=device)
self.mlp = GPTMLP(
d_model=d_model,
mlp_ratio=mlp_ratio,
device=device,
)
self.resid_attn_dropout = nn.Dropout(resid_pdrop)
self.resid_mlp_dropout = nn.Dropout(resid_pdrop)
def forward(
self,
x: torch.Tensor,
past_key_value: Optional[Tuple[torch.Tensor]] = None,
attn_bias: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.ByteTensor] = None,
is_causal: bool = True,
) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:
a = self.ln_1(x)
b, _, past_key_value = self.attn(a,
past_key_value=past_key_value,
attn_bias=attn_bias,
attention_mask=attention_mask,
is_causal=is_causal)
x = x + self.resid_attn_dropout(b)
m = self.ln_2(x)
n = self.mlp(m)
x = x + self.resid_mlp_dropout(n)
return x, past_key_value
|