replit/replit-code-v1_5-3b

1"""GPT Blocks used for the GPT Model."""2from typing import Any, Optional3import torch4import torch.nn as nn5from .fc import FC_CLASS_REGISTRY6try:7    import transformer_engine.pytorch as te8except:9    te = None10 11class MPTMLP(nn.Module):12 13    def __init__(self, d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None):14        super().__init__()15        fc_kwargs = {}16        if fc_type != 'te':17            fc_kwargs['device'] = device18        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 = True22 23    def forward(self, x: torch.Tensor) -> torch.Tensor:24        return self.down_proj(self.act(self.up_proj(x)))25FFN_CLASS_REGISTRY = {'mptmlp': MPTMLP}26if te is not None:27    te.LayerNormMLP._has_norm = True28    FFN_CLASS_REGISTRY['te_ln_mlp'] = te.LayerNormMLP29 30def build_ffn(d_model: int, expansion_ratio: int, fc_type: str='torch', device: Optional[str]=None, **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)36    elif ffn_type == 'te_ln_mlp':37        assert te is not None38        return te.LayerNormMLP(hidden_size=d_model, ffn_hidden_size=d_model * expansion_ratio, **kwargs)39    raise ValueError(f'ffn_type={ffn_type!r} not recognized.')