msj19/gated_deltaproduct

1# -*- coding: utf-8 -*-2 3# from https://github.com/HazyResearch/zoology/blob/main/zoology/mixers/convolution.py4 5import math6import warnings7from typing import Optional8 9import torch10import torch.nn as nn11import torch.nn.functional as F12from einops import rearrange13 14from ..modules.activations import ACT2FN15from ..utils import checkpoint16 17try:18    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update19except ImportError:20    causal_conv1d_fn = None21    causal_conv1d_update = None22 23 24def fft_conv(u, k, dropout_mask, gelu=True, k_rev=None):25    seqlen = u.shape[-1]26    fft_size = 2 * seqlen27    k_f = torch.fft.rfft(k, n=fft_size) / fft_size28    if k_rev is not None:29        k_rev_f = torch.fft.rfft(k_rev, n=fft_size) / fft_size30        k_f = k_f + k_rev_f.conj()31    u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)32 33    if len(u.shape) > 3:34        k_f = k_f.unsqueeze(1)35    y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen]36 37    out = y + u38    if gelu:39        out = F.gelu(out)40    if dropout_mask is not None:41        return (out * rearrange(dropout_mask, "b H -> b H 1")).to(dtype=u.dtype)42    else:43        return out.to(dtype=u.dtype)44 45 46@checkpoint47def proj_then_conv1d(48    x: torch.Tensor,49    proj_weight: torch.Tensor,50    conv1d_weight: torch.Tensor,51    conv1d_bias: Optional[torch.Tensor] = None,52    cache: Optional[torch.Tensor] = None53) -> torch.Tensor:54    # We do matmul and transpose BLH -> HBL at the same time55    x = rearrange(proj_weight @ rearrange(x, "b l d -> d (b l)"), "d (b l) -> b d l", l=x.shape[-2])56 57    if causal_conv1d_fn is None:58        raise ImportError("`causal_conv1d_fn` is not available. Please install `causal-conv1d` first.")59    if cache is None:60        x = causal_conv1d_fn(61            x=x,62            weight=rearrange(conv1d_weight, "d 1 w -> d w"),63            bias=conv1d_bias,64            activation="silu",65        ).transpose(1, 2)66    else:67        assert x.shape[-1] == 1, "Only support decoding with 1 token at a time for now"68        x = x.squeeze(-1)69        x = causal_conv1d_update(70            x=x,71            weight=rearrange(conv1d_weight, "d 1 w -> d w"),72            bias=conv1d_bias,73            cache=cache,74            activation="silu",75        )76    return x77 78 79class ShortConvolution(nn.Conv1d):80    """81    Simple wrapper around `nn.Conv1d` that accepts dimension last.82    """83 84    def __init__(85        self,86        hidden_size: int,87        kernel_size: int,88        bias: bool = False,89        activation: Optional[str] = 'silu',90        use_fast_conv1d: Optional[bool] = True91    ):92        super().__init__(93            in_channels=hidden_size,94            out_channels=hidden_size,95            kernel_size=kernel_size,96            groups=hidden_size,97            bias=bias,98            padding=kernel_size - 199        )100 101        self.hidden_size = hidden_size102        self.activation = None103        if activation is not None:104            assert activation in ['silu', 'swish'], f"Activation `{activation}` not supported yet."105            self.activation = activation106 107        if causal_conv1d_fn is None:108            if use_fast_conv1d:109                raise RuntimeError(110                    "Please either install `causal-conv1d>=1.4.0` to enable fast causal short convolution CUDA kernel "111                    "or set `use_fast_conv1d` to False"112                )113            else:114                warnings.warn(115                    "The naive Pytorch verison is very slow in practice, "116                    "please run `pip install causal-conv1d>=1.4.0` to install fast causal short convolution CUDA kernel"117                )118        self.use_fast_conv1d = use_fast_conv1d119 120    def extra_repr(self):121        s = ('{in_channels}, {out_channels}, kernel_size={kernel_size}'122             ', stride={stride}')123        if self.padding != (0,) * len(self.padding):124            s += ', padding={padding}'125        if self.dilation != (1,) * len(self.dilation):126            s += ', dilation={dilation}'127        if self.output_padding != (0,) * len(self.output_padding):128            s += ', output_padding={output_padding}'129        if self.groups != 1:130            s += ', groups={groups}'131        if self.bias is None:132            s += ', bias=False'133        if self.padding_mode != 'zeros':134            s += ', padding_mode={padding_mode}'135        if self.activation is not None:136            s += ', activation={activation}'137        if not self.use_fast_conv1d:138            s += ', use_fast_conv1d={use_fast_conv1d}'139        return s.format(**self.__dict__)140 141    def forward(142        self,143        x: torch.Tensor,144        mask: Optional[torch.Tensor] = None,145        cache: Optional[torch.Tensor] = None146    ) -> torch.Tensor:147        """148        Args:149            x (`torch.Tensor`):150                Tensor of shape `[batch_size, seq_len, hidden_size]`151            mask (`Optional[torch.Tensor]`):152                Attention mask dealing with padded positions.153            cache (`Optional[torch.Tensor]`):154                Previous cache tensor of shape `[batch_size, hidden_size, kernel_size]`,155        Returns:156            Tensor of shape `[batch_size, seq_len, hidden_size]`. The `cache` (if provided) is updated inplace.157        """158 159        if mask is not None:160            x = x.mul_(mask.unsqueeze(-1))161        if cache is not None and x.shape[1] == 1:162            return self.step(x, cache)163        x = rearrange(x, "b l d -> b d l")164        # Update state (B D W)165        if cache is not None:166            cache.copy_(F.pad(x, (self.kernel_size[0] - x.shape[-1], 0)))167        if self.use_fast_conv1d:168            x = causal_conv1d_fn(169                x=x,170                weight=rearrange(self.weight, "d 1 w -> d w"),171                bias=self.bias,172                activation=self.activation,173            )174        else:175            x = self._conv_forward(x, self.weight, self.bias)[..., :x.shape[-1]]176            if self.activation is not None:177                x = ACT2FN[self.activation](x)178        return rearrange(x, "b d l -> b l d")179 180    def step(181        self,182        x: torch.Tensor,183        cache: torch.Tensor184    ):185        assert x.shape[1] == 1, "Only support decoding with 1 token at a time for now"186 187        x = x.squeeze(1)188        if self.use_fast_conv1d:189            x = causal_conv1d_update(190                x=x,191                conv_state=cache,192                weight=rearrange(self.weight, "d 1 w -> d w"),193                bias=self.bias,194                activation=self.activation,195            )196        else:197            dtype = x.dtype198            cache.copy_(torch.roll(cache, shifts=-1, dims=-1))199            cache[:, :, -1] = x200            x = torch.sum(cache * rearrange(self.weight, "d 1 w -> d w"), dim=-1)201            if self.bias is not None:202                x = x + self.bias203            if self.activation is not None:204                x = ACT2FN[self.activation](x).to(dtype=dtype)205        return x.unsqueeze(1)206 207    @property208    def state_size(self) -> int:209        return self.hidden_size * self.kernel_size210 211 212class LongConvolution(nn.Module):213    """214    LongConvolution applies a convolution operation on the input tensor using a fixed215    filter of length l_max.216    The filter is learned during training and is applied using FFT convolution.217    Args:218        hidden_size (int): The number of expected features in the input and output.219        l_max (int): The maximum sequence length.220    Returns:221        y: (b, l, d) tensor222    """223 224    def __init__(225        self,226        hidden_size: int,227        l_max: int,228        **kwargs,229    ):230        """231        Initializes the LongConvolution module.232        Args:233            hidden_size (int): The number of expected features in the input and output.234            l_max (int): The maximum sequence length.235        """236        super().__init__()237        self.hidden_size = hidden_size238        self.filter = nn.Parameter(torch.randn(self.hidden_size, l_max), requires_grad=True)239 240    def forward(self, x: torch.Tensor, *args, **kwargs):241        """242        Applies the LongConvolution operation on the input tensor.243        Args:244            x: (b, l, d) tensor245        Returns:246            y: (b, l, d) tensor247        """248        x = x.transpose(1, 2)249        y = fft_conv(x, self.filter, dropout_mask=None, gelu=False)250        y = y.transpose(1, 2)251        return y.to(dtype=x.dtype)252 253 254class PositionalEmbedding(nn.Module):255    def __init__(self, emb_dim: int, seq_len: int, **kwargs):256        """Complex exponential positional embeddings for implicit long convolution filters."""257        super().__init__()258 259        self.seq_len = seq_len260        # The time embedding fed to the filteres is normalized so that t_f = 1261        t = torch.linspace(0, 1, self.seq_len)[None, :, None]  # 1, L, 1262 263        if emb_dim > 1:264            bands = (emb_dim - 1) // 2265        # To compute the right embeddings we use the "proper" linspace266        t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]267        w = 2 * math.pi * t_rescaled / seq_len  # 1, L, 1268 269        f = torch.linspace(1e-4, bands - 1, bands)[None, None]270        z = torch.exp(-1j * f * w)271        z = torch.cat([t, z.real, z.imag], dim=-1)272        self.z = nn.Parameter(z, requires_grad=False)273 274    def forward(self, L):275        return self.z[:, :L]276 277 278class ImplicitLongConvolution(nn.Module):279    """280    Long convolution with implicit filter parameterized by an MLP.281 282    Args:283        hidden_size (int):284            The number of expected features in the input and output.285        l_max (int):286            The maximum sequence length.287        d_emb (Optional[int]):288            The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine).289            Defaults to 3.290        d_hidden (Optional[int]):291            The number of features in the hidden layer of the MLP. Defaults to 16.292 293    Attributes:294        pos_emb (`PositionalEmbedding`): The positional embedding layer.295        mlp (`nn.Sequential`): The MLP that parameterizes the implicit filter.296 297    """298 299    def __init__(300        self,301        hidden_size: int,302        l_max: int,303        d_emb: int = 3,304        d_hidden: int = 16,305        **kwargs,306    ):307        """308        Long convolution with implicit filter parameterized by an MLP.309 310 311        """312        super().__init__()313        self.hidden_size = hidden_size314        self.d_emb = d_emb315 316        assert (317            d_emb % 2 != 0 and d_emb >= 3318        ), "d_emb must be odd and greater or equal to 3 (time, sine and cosine)"319        self.pos_emb = PositionalEmbedding(d_emb, l_max)320 321        # final linear layer322        self.mlp = nn.Sequential(323            nn.Linear(d_emb, d_hidden),324            torch.nn.ReLU(),325            nn.Linear(d_hidden, hidden_size),326        )327 328    def filter(self, seq_len: int, *args, **kwargs):329        k = self.mlp(self.pos_emb(seq_len))330 331        return k.transpose(1, 2)332 333    def forward(self, x: torch.Tensor, *args, **kwargs):334        """335        Args:336            x: (b, l, d) tensor337        Returns:338            y: (b, l, d) tensor339        """340        x = x.transpose(1, 2)341        k = self.filter(x.shape[-1])342        y = fft_conv(x, k, dropout_mask=None, gelu=False)343 344        y = y.transpose(1, 2)345        return y.to(dtype=x.dtype)346