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# Note: MP and PP utilities are removed for ease of use and editing.5 6import torch7import torch.nn as nn8import torch.nn.functional as F9from torch.utils.checkpoint import checkpoint10 11from .utils import print_rank_0, column_split12from .cache import InferenceParams, RecurrentInferenceParams13from .engine import HyenaInferenceEngine14from .layers import (15    RMSNorm,16    ParallelGatedMLP,17    VocabParallelEmbedding,18)19 20try:21    from flash_attn.modules.mha import MHA22except ImportError:23    "flash_attn not installed"24 25 26class AttentionBlock(nn.Module):27    def __init__(self, config, layer_idx) -> None:28        super().__init__()29        self.config = config30        self.pre_norm, self.post_norm = RMSNorm(config), RMSNorm(config)31        self.layer_idx = layer_idx32        self.proj_groups = config.get("proj_groups", 1)33        dtype = config.get("attn_block_dtype", torch.bfloat16)34        mlp_dtype = config.get("mlp_dtype", torch.bfloat16)35        self.num_attention_heads = config.num_attention_heads36        self.hidden_size_per_attention_head = config.hidden_size // config.num_attention_heads37 38        self.counter = 039        self.inner_mha_cls = MHA(40            embed_dim=config.hidden_size,41            num_heads=config.num_attention_heads,42            num_heads_kv=config.num_attention_heads // self.proj_groups,43            rotary_emb_dim=config.hidden_size // config.num_attention_heads,44            qkv_proj_bias=config.get("qkv_proj_bias", True),45            rotary_emb_base=config.get("rotary_emb_base", 10000),46            causal=True,47            layer_idx=layer_idx,48            out_proj_bias=config.get("mha_out_proj_bias", True),49            use_flash_attn=self.config.use_flash_attn,50        ).to(dtype=dtype)51 52        if self.config.get("smeared_gqa", False):53            self.inner_mha_cls.num_heads_kv = self.inner_mha_cls.num_heads54        self.inner_mha_cls.rotary_emb.register_buffer(55            "inv_freq", self.inner_mha_cls.rotary_emb.inv_freq56        )57 58        self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)59 60    def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):61        if (62            type(padding_mask) == torch.Tensor63        ):  # workaround for masking bug in FA. This works because Wqkv does not have bias64            # and attention scores will be also automatically zeroed.65            u = u * padding_mask[..., None]66 67        u = (68            self.inner_mha_cls(69                self.pre_norm(u),70                inference_params=inference_params,71            )72            + u73        )74        if type(padding_mask) == torch.Tensor:  # guard against bias75            u = u * padding_mask[..., None]76        u = self.mlp(self.post_norm(u)) + u77        return u, None78 79 80class ParallelHyenaFilter(nn.Module):81    def __init__(self, config, layer_idx) -> None:82        super().__init__()83        self.config = config84        self.layer_idx = layer_idx85        self.hyena_filter_groups = config.get("hyena_filter_groups", self.config.hidden_size)86 87        self.use_flashfft = config.get("use_flashfft", False)88        self.state_size = config.state_size89        self.hidden_size = config.hidden_size90        self.num_filters = config.num_filters91        self.inference_mode = config.get("inference_mode", True)92        self.counter = 093        self.column_split_hyena = config.get("column_split_hyena", True)94 95        assert self.hidden_size % self.num_filters == 0 and self.num_filters <= self.hidden_size96 97        self.D = nn.Parameter(torch.zeros(self.hidden_size))98 99        # attention heads are not used except to split post short_filter100        # projections in the same way as the checkpoint101        self.num_attention_heads = config.num_attention_heads102        self.hidden_size_per_attention_head = self.hidden_size // self.num_attention_heads103 104        # after preprocessing here we can save the new checkpoint105        self.short_filter_length = config.short_filter_length106        self.short_filter_weight = nn.Parameter(107            torch.randn(3 * config.hidden_size, 1, config.short_filter_length)108        )109        self.short_filter_bias = (110            nn.Parameter(torch.randn(3 * config.hidden_size)) if config.short_filter_bias else None111        )112 113        self.engine = HyenaInferenceEngine(layer_idx=layer_idx)114        self.use_flash_depthwise = config.get("use_flash_depthwise", False)115        self.data_dtype = None116 117        if self.use_flash_depthwise:118            self.fir_fn = FlashDepthwiseConv1d(119                channels=3 * self.hidden_size,120                kernel_size=self.short_filter_length,121                padding=self.short_filter_length - 1,122                weights=self.short_filter_weight,123                bias=self.short_filter_bias,124                device=None,125                dtype=self.config.get("depthwise_dtype", torch.bfloat16),126            )127        else:128            self.fir_fn = F.conv1d129 130        self.fftconv_fn = None131        self.long_fir_threshold = config.get("long_fir_threshold", None)132        if self.long_fir_threshold is not None:133            assert (134                self.use_flashfft is False135            ), "long_fir_threshold not compatible with fused flashfft"136 137        self.num_systems = self.hidden_size // self.hyena_filter_groups138        self.poles = nn.Parameter(torch.randn(self.num_systems, self.state_size, 1, 2))139        self.residues = nn.Parameter(torch.randn(self.num_systems, self.state_size, 1, 2))140        self.h = None141 142    def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):143        if (144            inference_params is not None145            and self.layer_idx in inference_params.fir_state_dict.keys()146        ):147            return self.sequential_forward(u, inference_params)148 149        else:150            return self.parallel_forward(u, inference_params, padding_mask)151 152    def parallel_forward(self, u, inference_params=None, padding_mask=None):153        L = u.shape[1]154        z_pre, fir_state = self.engine.parallel_fir(155            self.fir_fn,156            u,157            self.short_filter_weight,158            self.short_filter_bias,159            L,160            fir_length=self.short_filter_length,161            inference_params=inference_params,162            padding_mask=padding_mask,163        )164        if inference_params:165            inference_params.fir_state_dict[self.layer_idx] = fir_state166 167        if self.h is None:168            h, filter_dtype, poles, residues = self.compute_filter(L, u.device)169        else:170            h = self.h171            filter_dtype = self.h.dtype172 173        if self.hyena_filter_groups > 1:174            h = h.repeat_interleave(self.hidden_size // self.hyena_filter_groups, 1)175 176        # if inference_params is not None, we plan to perform generation:177        # prefilling for the IIR portion of the filter is handled by the engine.178        dims = (179            self.hidden_size,180            self.num_attention_heads,181            self.hidden_size_per_attention_head,182            self.state_size,183            self.hyena_filter_groups,184        )185        y = self.engine.parallel_iir(186            z_pre,187            h,188            self.D,189            L,190            t=self.t,191            poles=self.poles,192            dims=dims,193            inference_params=inference_params,194            layer_idx=self.layer_idx,195            prefill_style=self.config.get("prefill_style", "fft"),196            use_flashfft=self.use_flashfft,197            fftconv_fn=self.fftconv_fn,198            column_split_hyena=self.column_split_hyena,199            long_fir_threshold=self.long_fir_threshold,200            padding_mask=padding_mask,201        )202 203        return y, inference_params204 205    def sequential_forward(self, u, inference_params):206        if self.data_dtype is None:207            self.data_dtype = u.dtype208        if len(u.shape) > 2:209            u = u[:, -1]210 211        fir_state, iir_state = (212            inference_params.fir_state_dict[self.layer_idx],213            inference_params.state_dict[self.layer_idx],214        )215 216        z_pre, fir_state = self.engine.step_fir(217            u, fir_state, weight=self.short_filter_weight, bias=self.short_filter_bias218        )219        x2, x1, v = (220            column_split(z_pre, self.num_attention_heads, self.hidden_size_per_attention_head)221            if self.column_split_hyena222            else z_pre.split([self.hidden_size, self.hidden_size, self.hidden_size], dim=1)223        )224 225        y, iir_state = self.engine.step_iir(226            x2,227            x1,228            v,229            self.D,230            self.residues,231            self.poles,232            iir_state,233            iir_groups=self.hyena_filter_groups,234        )235 236        inference_params.fir_state_dict[self.layer_idx] = fir_state237        inference_params.state_dict[self.layer_idx] = iir_state238        y = y.to(dtype=self.data_dtype)239        return y[:, None], inference_params240 241    def update_time(self, L, device):242        """243        Set [0, 1, ..., L-1] where L is the length of the current batch of inputs.244        If L is greater than the length of the previous batch, then the time vector is245        reinitialized. Otherwise, the time vector is truncated from cache.246        """247        if not hasattr(self, "t"):248            self.t = torch.arange(L, device=device)[None, None]249        elif self.t.shape[-1] < L:250            self.t = torch.arange(L, device=device)[None, None]251        else:252            self.t = self.t[..., :L]253 254    def compute_filter(self, L, device):255        self.update_time(L, device)256        filter_dtype = torch.float32257        residues, log_poles = (258            torch.view_as_complex(self.residues.to(filter_dtype)),259            torch.view_as_complex(self.poles.to(filter_dtype)).log(),260        )261        h = (residues * (log_poles * self.t).exp()).real.sum(1)[None]262        return h, filter_dtype, log_poles, residues263 264 265class ParallelGatedConvBlock(nn.Module):266    def __init__(self, config, layer_idx) -> None:267        super().__init__()268        self.config = config269        self.layer_idx = layer_idx270        dtype = config.get("hyena_block_dtype", torch.float32)271        mlp_dtype = config.get("mlp_dtype", torch.bfloat16)272        self.pre_norm, self.post_norm = RMSNorm(config).to(dtype=dtype), RMSNorm(config).to(273            dtype=dtype274        )275        self.filter = ParallelHyenaFilter(config, layer_idx).to(dtype=dtype)276        self.projections = nn.Linear(config.hidden_size, 3 * config.hidden_size)277        self.out_filter_dense = nn.Linear(config.hidden_size, config.hidden_size).to(dtype)278        self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)279 280    def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):281        z = self.projections(self.pre_norm(u))282        if type(padding_mask) == torch.Tensor:  # guard against bias283            z = z * padding_mask[..., None]284 285        z, inference_params = self.filter(286            z, inference_params=inference_params, padding_mask=padding_mask287        )288 289        u = self.out_filter_dense(z) + u290        if type(padding_mask) == torch.Tensor:  # guard against bias291            u = u * padding_mask[..., None]292        u = self.mlp(self.post_norm(u)) + u293        return u, inference_params294 295 296def get_block(config, layer_idx, flash_fft=None):297    if layer_idx in config.attn_layer_idxs:298        return AttentionBlock(config, layer_idx)299    elif layer_idx in config.hyena_layer_idxs:300        block = ParallelGatedConvBlock(config, layer_idx)301        if config.get("use_flashfft", "False"):302            block.filter.fftconv_fn = flash_fft303        return block304    else:305        raise NotImplementedError306 307 308class StripedHyena(nn.Module):309    def __init__(self, config):310        super().__init__()311        self.config = config312        self.embedding_layer = VocabParallelEmbedding(config)313        self.norm = RMSNorm(config) if config.get("final_norm", True) else None314        self.unembed = self.emb if config.tie_embeddings else VocabParallelEmbedding(config)315        self.gradient_checkpointing = False316        317        if config.get("use_flashfft", "False"):318            raise NotImplementedError("Please use standalone SH code for other custom kernels")319        else:320            self.flash_fft = None321 322        self.blocks = nn.ModuleList(323            get_block(config, layer_idx, flash_fft=self.flash_fft)324            for layer_idx in range(config.num_layers)325        )326 327    def forward(self, x, inference_params_dict=None, padding_mask=None):328        L = x.shape[1]329        x = self.embedding_layer.embed(x)330        if inference_params_dict is not None:331            x, inference_params_dict_out = self.stateful_forward(332                x,333                inference_params_dict=inference_params_dict,334            )335        else:336            x, inference_params_dict_out = self.stateless_forward(x, padding_mask=padding_mask)337        x = self.norm(x)338        x = self.unembed.unembed(x)339        return x, inference_params_dict_out340 341    def stateful_forward(self, x, inference_params_dict=None):342        for block_idx, block in enumerate(self.blocks):343            block_name = "mha" if block_idx in self.config.attn_layer_idxs else "hyena"344            inference_params = inference_params_dict[block_name]345            x, _ = block(x, inference_params=inference_params)346 347        return x, inference_params_dict348 349    def stateless_forward(self, x, padding_mask=None):350        if type(padding_mask) == torch.Tensor:351            x = x * padding_mask[..., None]352 353        for block_idx, block in enumerate(self.blocks):354            if self.gradient_checkpointing and self.training:355                def create_custom_forward(module):356                    def custom_forward(*inputs):357                        # None for past_key_value358                        return module(*inputs, inference_params=None, padding_mask=padding_mask)359 360                    return custom_forward361 362                x, _ = checkpoint(create_custom_forward(block), x, use_reentrant=False)363            else:364                x, _ = block(x, inference_params=None, padding_mask=padding_mask)365        return x, None366 367    def initialize_inference_params(self):368        print_rank_0("Initializing inference params...")369        inference_params_dict = {370            "mha": InferenceParams(371                max_seqlen=self.config.get("max_seqlen", 8192),372                max_batch_size=self.config.get("max_batch_size", 1),373                seqlen_offset=0,374            ),375            "hyena": RecurrentInferenceParams(376                fir_filter_length=self.config.short_filter_length,377                state_dim=self.config.state_size,378                seqlen_offset=0,379            ),380        }381        return inference_params_dict382 383    def precompute_filters(self, L, device):384        for block_idx, block in enumerate(self.blocks):385            if type(block) == ParallelGatedConvBlock:386                if type(block.filter) == ParallelHyenaFilter:387                    L = block.filter.long_fir_threshold or L388                    print_rank_0(f"Precomputing filters, L={L}...")389 390                    filter_dtype = torch.float16 if L >= 2048 else torch.float32391 392                    block.filter._set_time(L, device)393                    residues, poles = (394                        torch.view_as_complex(block.filter.residues.to(torch.float16)),395                        torch.view_as_complex(block.filter.poles.to(torch.float16)),396                    )397 398                    block.filter.h = (residues * poles**block.filter.t).real.sum(1)[None]399                    block.filter.h = block.filter.h.to(dtype=filter_dtype)400 401    def load_poles_residues(self, path):402        "Load different poles and residues for each layer."403        for block_idx, block in enumerate(self.blocks):404            if type(block) == ParallelGatedConvBlock:405                if type(block.filter) == ParallelHyenaFilter:406                    print(f"Loading poles and residues for block {block_idx}")407                    poles = torch.load(path + f"/approx_poles_{block_idx+1}.pt", map_location="cpu")408                    poles = torch.view_as_real(poles)409                    residues = torch.load(410                        path + f"/approx_residues_{block_idx+1}.pt", map_location="cpu"411                    )412                    residues = torch.view_as_real(residues)413                    poles = poles.permute(1, 0, 2).unsqueeze(-2)414                    residues = residues.permute(1, 0, 2).unsqueeze(-2)415 416                    block.filter.poles = nn.Parameter(poles)417                    block.filter.residues = nn.Parameter(residues)418 419    def to_bfloat16_except_poles_residues(self):420        """Convert all parameters to bfloat16 except for the poles and residues.421 422        Particularly important for longer prompts.423        """424        for k, p in self.named_parameters():425            if "poles" not in k and "residues" not in k:426                p.data = p.data.to(torch.bfloat16)427