togethercomputer/StripedHyena-Nous-7B

1# Copyright (c) Together2# This software is distributed under the terms of the Apache License, Version 2.03# Author: Michael Poli4 5import torch6import torch.nn as nn7import torch.nn.functional as F8 9try:10    import conv1d_cpp11except:12    pass13from .utils import column_split14 15 16def canonicalize_modal_system(poles, residues):17    """Canonicalize a modal system.18 19    Args:20        poles (Tensor): The poles of the system.21        residues (Tensor): The residues of the system.22 23    Returns:24        Tuple[Tensor, Tensor]: The canonicalized poles and residues.25    """26    raise NotImplementedError27 28 29IIR_PREFILL_MODES = [30    "recurrence",31    "modal-fft",32    "hybrid-modal-recurrence",33    "modal-scan",34    "canonical-fft",35    "iir-fir-caching",36]37 38 39class HyenaInferenceEngine:40    def __init__(41        self, fir_fn=None, fftconv_fn=None, iir_prefill_style="modal-fft", layer_idx=None42    ) -> None:43        self.fir_fn = fir_fn44        self.fftconv_fn = fftconv_fn45        assert (46            iir_prefill_style in IIR_PREFILL_MODES47        ), f"iir_prefill_style must be one of {IIR_PREFILL_MODES}"48        self.iir_prefill_style = iir_prefill_style49        self.layer_idx = layer_idx50        self.low_mem_mode = False51 52    def parallel_fir(53        self,54        fir_fn,55        u,56        weight,57        bias,58        L,59        fir_length=3,60        inference_params=None,61        prefill_mode=None,62        padding_mask=None,63    ):64        """Compute the output state of the long convolutional filter."""65        # prepare input layout, dimensions and dispatch to fir kernel66        if fir_fn != torch.nn.functional.conv1d:67            z_pre = fir_fn(u)[:, :L]  # B, L, D68            z_pre = z_pre.permute(0, 2, 1)69        else:70            u = u.permute(0, 2, 1)  # B, D, L71            z_pre = fir_fn(72                u,73                weight,74                bias,75                stride=1,76                padding=fir_length - 1,77                groups=u.shape[1],78            )[..., :L]79 80        # handle padding post fir, the only place with biases81        if type(padding_mask) == torch.Tensor:82            z_pre = z_pre * padding_mask[:, None]83 84        if inference_params is not None:85            # handle seqlen last and dim last cases for `u`86            if fir_fn != torch.nn.functional.conv1d:87                fir_state = u[:, -fir_length + 1 :].permute(0, 2, 1)88            else:89                fir_state = u[..., -fir_length + 1 :]90        else:91            fir_state = None92 93        return z_pre, fir_state94 95    def parallel_iir(96        self,97        z_pre,98        h,99        D,100        L,101        poles,102        t,103        dims,104        layer_idx,105        inference_params=None,106        prefill_style="fft",107        fftconv_fn=None,108        padding_mask=None,109        use_flashfft=False,110        column_split_hyena=False,111        long_fir_threshold=None,112    ):113        """Compute the output state of the short convolutional filter."""114        fft_size = 2 * L115        hidden_size, num_attention_heads, hidden_size_per_attention_head, _, _ = dims116        # Compatibility with training infra that column splits the projections117        if column_split_hyena:118            z = z_pre.reshape(119                z_pre.shape[0],120                num_attention_heads,121                3 * hidden_size_per_attention_head,122                z_pre.shape[2],123            )124            x2, x1, v = (125                z[:, :, :hidden_size_per_attention_head],126                z[127                    :,128                    :,129                    hidden_size_per_attention_head : 2 * hidden_size_per_attention_head,130                ],131                z[:, :, 2 * hidden_size_per_attention_head :],132            )133            x2, x1, v = (134                x2.reshape(x2.shape[0], -1, x2.shape[-1]),135                x1.reshape(x1.shape[0], -1, x1.shape[-1]),136                v.reshape(v.shape[0], -1, v.shape[-1]),137            )138        else:139            x2, x1, v = z_pre.split([hidden_size, hidden_size, hidden_size], dim=1)140 141        x1v = x1 * v142 143        if use_flashfft and (L % 2) == 0:  # only works with even L144            y = fftconv_fn(145                x1v.to(dtype=torch.bfloat16).contiguous(),146                h.to(dtype=torch.float32),147            )148            X_s = None149 150        elif long_fir_threshold is None:151            H = torch.fft.rfft(h.to(dtype=torch.float32), n=fft_size) / fft_size152            X_s = torch.fft.fft(x1v.to(dtype=torch.float32), n=fft_size)153            X = X_s[..., : H.shape[-1]]154            if len(z_pre.shape) > 3:155                H = H.unsqueeze(1)156            y = torch.fft.irfft(X * H, n=fft_size, norm="forward")[..., :L]157        else:158            assert h.shape[0] == 1, "batch size must be 1 for long_fir_threshold"159            h = h[0][:, None]  # rearrange to d, 1, l for depthwise conv1d160            h = h[..., :long_fir_threshold]161            y = F.conv1d(162                x1v,163                h.to(dtype=x1v.dtype),164                stride=1,165                groups=x1v.shape[1],166                padding=h.shape[-1] - 1,167            )[..., :L]168 169        y = y.to(dtype=x1v.dtype)170        y = (y + x1v * D.unsqueeze(-1)) * x2171        if inference_params is not None:172            if prefill_style == "fft":173                self.prefill_via_modal_fft(174                    inference_params=inference_params,175                    x1v=x1v,176                    X_s=X_s,177                    L=L,178                    t=t,179                    poles=poles,180                    dims=dims,181                    layer_idx=layer_idx,182                    use_flashfft=use_flashfft,183                )184 185            elif prefill_style == "recurrence":186                self.prefill_via_direct_recurrence(187                    inference_params=inference_params,188                    x1v=x1v,189                    L=L,190                    poles=poles,191                )192 193            else:194                raise NotImplementedError195            if self.low_mem_mode:196                del z_pre, x2, x1, v, x1v, h197                torch.cuda.empty_cache()198 199        return y.permute(0, 2, 1)200 201    def step_fir(self, u, fir_state, weight, bias=None):202        """Step the FIR filter.203 204        Note:205        `fir_state` contains the last `short_filter_length - 1` elements of `u`: `u_(L-2), u_{L-1), ...`206        We assume dimensions of `short_filter_weight` to be `[d, 1, short_filter_len]` (SISO / multi SISO layout).207        """208        h0, h = weight[..., 0, -1], weight[..., 0, :-1]209        h0, h = h0[None], h[None]210        y = h0 * u + torch.sum(fir_state * h, dim=-1) + bias211 212        # update213        fir_state = torch.roll(fir_state, -1, dims=2)214        fir_state[..., -1] = u215        return y, fir_state216 217    def step_iir(self, x2, x1, v, D, residues, poles, iir_state, iir_groups=1):218        x1v = x1 * v219 220        residues, poles = (221            torch.view_as_complex(residues.to(torch.float32)),222            torch.view_as_complex(poles.to(torch.float32)),223        )224        # squeeze the dummy seqlen dimension225        # D, state_dim, 1 -> 1, D, state_dim226        residues, poles = residues[..., 0][None], poles[..., 0][None]227        iir_state = poles * iir_state + x1v[..., None]228 229        res_state = torch.sum(residues * iir_state, dim=-1).real230 231        if iir_groups > 1:232            raise NotImplementedError233        y = x2 * (res_state + D * x1v)234 235        return y, iir_state236 237    def prefill_via_fir_caching(self, u, inference_params, L, *args, **kwargs):238        """Turns the IIR filter into a FIR and uses a cache for decoding."""239        raise NotImplementedError(":)")240 241    def prefill_via_direct_recurrence(self, inference_params, x1v, L, poles, *args, **kwargs):242        """243        Compute the IIR state via explicit SSM recurrence (modal form)244        """245        x1v_ = x1v[..., None, None]  # b, d, l, sdim, reim246        x1v_ = x1v_.repeat(1, 1, 1, 1, 2)  # b, d, l, sdim, reim247 248        state = x1v_[:, :, 0]249        poles = poles[:, :, 0].to(dtype=torch.float32)250 251        for i in range(L):252            state = poles * state + x1v_[:, :, i]253        inference_params.state_dict[self.layer_idx] = torch.view_as_complex(254            state.to(dtype=torch.float32)255        )256 257    def prefill_via_hybrid_recurrence(258        self, inference_params, u, log_poles, x1v_f_a, L, *args, **kwargs259    ):260        """261        Compute the IIR state via hybrid recurrence-convolution over blocks262        """263        raise NotImplementedError(":)")264 265    def prefill_via_scan(self, u, inference_params=None, *args, **kwargs):266        raise NotImplementedError267 268    def prefill_via_canonical_fft(self, u, inference_params=None, *args, **kwargs):269        """270        Compute the IIR state via a single FFT with the denominator of the SSM in companion form.271 272        This is the most memory efficient "parallelized" prefilling method for Hyena.273 274        From: https://arxiv.org/abs/2310.18780275        """276        raise NotImplementedError(":)")277 278    def prefill_via_modal_fft(279        self,280        inference_params,281        x1v,282        L,283        poles,284        t,285        dims,286        layer_idx,287        X_s=None,288        use_flashfft=False,289        state_dtype=torch.complex64,290        *args,291        **kwargs,292    ):293        """294        Compute the IIR state via a single FFT, using the poles of the SSM in modal form.295        """296        # When the model has a long convolution derived from a SSM in modal form and prefill_style is "fft",297        # we split the filter into poles and residues and reuse FFT computation on the input.298        # This optimization is currently not supported when using flashfftconv.299        hidden_size, _, _, state_size, hyena_filter_groups = dims300 301        if use_flashfft:302            # using real states303            poles = poles.squeeze().reshape(poles.shape[0], -1)[..., None]304 305            state_s = poles**t306            if hyena_filter_groups > 1:307                raise NotImplementedError308 309            x1v = x1v[:, :, None].repeat(1, 1, 2 * state_size, 1)310            x1v = x1v.reshape(x1v.shape[0], -1, x1v.shape[-1])311            state_s = state_s[None]312 313            state = self.fftconv_fn(314                x1v.contiguous(),315                state_s.to(dtype=torch.float32),316            )317            state = state[..., L - 1].reshape(x1v.shape[0], hidden_size, state_size, 2)318            state = torch.view_as_complex(state.contiguous())319            inference_params.state_dict[self.layer_idx] = state.to(dtype=state_dtype)320        else:321            assert X_s is not None322            bs = x1v.shape[0]323            fft_size = 2 * L324            poles = torch.view_as_complex(poles.to(torch.float32))325            state_s = poles**t326            state_S = torch.fft.fft(state_s, n=fft_size).repeat(327                bs, 1, 1, 1328            )  # B, D, state_dim, 2 * L329            if hyena_filter_groups > 1:330                state_S = state_S.repeat_interleave(hidden_size // hyena_filter_groups, 1)331            state = torch.fft.ifft(X_s[..., None, :] * state_S, n=fft_size)332            inference_params.state_dict[layer_idx] = state[..., L - 1].to(dtype=state_dtype)333 334    def _compute_state(self, log_poles, u, t, L, *args, **kwargs):335        """336        Compute the IIR state given an input `u` and log_poles of the modal system.337        """338        bs = u.shape[0]339        fft_size = 2 * L340        U = torch.fft.rfft(u.to(torch.float32), n=fft_size)341        fft_size = 2 * L342        x = (log_poles * t).exp()343        # [batch, hidden_size, state_dim, 2 * seqlen]344        X = torch.fft.fft(x, n=fft_size).repeat(bs, 1, 1, 1)345        state = torch.fft.ifft(U[..., None, :] * X, n=fft_size)[..., :L]346        return state347