Create ffn.py
Browse files
ffn.py
ADDED
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
"""GPT Blocks used for the GPT Model."""
|
2 |
+
from typing import Any, Optional
|
3 |
+
import torch
|
4 |
+
import torch.nn as nn
|
5 |
+
from .fc import FC_CLASS_REGISTRY
|
6 |
+
try:
|
7 |
+
import transformer_engine.pytorch as te
|
8 |
+
except:
|
9 |
+
te = None
|
10 |
+
|
11 |
+
class MPTMLP(nn.Module):
|
12 |
+
|
13 |
+
def __init__(self, d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None, bias: bool=True):
|
14 |
+
super().__init__()
|
15 |
+
fc_kwargs: dict[str, Any] = {'bias': bias}
|
16 |
+
if fc_type != 'te':
|
17 |
+
fc_kwargs['device'] = device
|
18 |
+
self.up_proj = FC_CLASS_REGISTRY[fc_type](d_model, expansion_ratio * d_model, **fc_kwargs)
|
19 |
+
self.act = nn.GELU(approximate='none')
|
20 |
+
self.down_proj = FC_CLASS_REGISTRY[fc_type](expansion_ratio * d_model, d_model, **fc_kwargs)
|
21 |
+
self.down_proj._is_residual = True
|
22 |
+
|
23 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
24 |
+
return self.down_proj(self.act(self.up_proj(x)))
|
25 |
+
FFN_CLASS_REGISTRY = {'mptmlp': MPTMLP}
|
26 |
+
if te is not None:
|
27 |
+
te.LayerNormMLP._has_norm = True
|
28 |
+
FFN_CLASS_REGISTRY['te_ln_mlp'] = te.LayerNormMLP
|
29 |
+
|
30 |
+
def build_ffn(d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None, bias: bool=True, **kwargs: Any) -> nn.Module:
|
31 |
+
ffn_type = kwargs.pop('ffn_type')
|
32 |
+
if ffn_type == 'mptmlp':
|
33 |
+
if len(kwargs) > 0:
|
34 |
+
raise ValueError(f'MPTMLP got an unexpected keyword argument: {kwargs}')
|
35 |
+
return MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, fc_type=fc_type, device=device, bias=bias)
|
36 |
+
elif ffn_type == 'te_ln_mlp':
|
37 |
+
assert te is not None
|
38 |
+
return te.LayerNormMLP(hidden_size=d_model, ffn_hidden_size=d_model * expansion_ratio, bias=bias, **kwargs)
|
39 |
+
raise ValueError(f'ffn_type={ffn_type!r} not recognized.')
|