configuration_evo2.py

"""Configuration for Evo2 (StripedHyena2)."""

from __future__ import annotations

import json
from typing import List

from transformers import PretrainedConfig


class Evo2Config(PretrainedConfig):
    """Evo2 config.

    Defaults match evo2_7b_base (32-layer, 4096-hidden). Per-variant overrides
    are written into the config.json of each repo by convert_checkpoint.py.

    Block dispatch is driven by four index lists:
      attn_layer_idxs : transformer (MHA + RoPE) blocks
      hcl_layer_idxs  : Hyena Cascade Long  (IIR via log_poles + residues)
      hcm_layer_idxs  : Hyena Cascade Medium (FIR, fir_inner_filter_length=128)
      hcs_layer_idxs  : Hyena Cascade Short  (FIR, fir_inner_filter_length=7)
    Their disjoint union must equal range(num_layers).
    """

    model_type = "evo2"

    def __init__(
        self,
        # Architecture core
        vocab_size: int = 512,
        hidden_size: int = 4096,
        num_filters: int = 4096,
        inner_mlp_size: int = 11008,
        num_layers: int = 32,
        num_attention_heads: int = 32,
        # Block dispatch
        attn_layer_idxs: List[int] | None = None,
        hcl_layer_idxs: List[int] | None = None,
        hcm_layer_idxs: List[int] | None = None,
        hcs_layer_idxs: List[int] | None = None,
        # Filter geometry
        hcm_filter_length: int = 128,
        hcs_filter_length: int = 7,
        hcl_filter_groups: int = 4096,
        hcm_filter_groups: int = 256,
        hcs_filter_groups: int = 256,
        hyena_filter_groups: int = 1,
        short_filter_length: int = 3,
        short_filter_bias: bool = False,
        state_size: int = 16,
        # Channel-split conventions
        column_split: bool = True,
        column_split_hyena: bool = False,
        interleave: bool = True,
        hyena_flip_x1x2: bool = False,
        # Norms / activations
        eps: float = 1e-6,
        final_norm: bool = True,
        mlp_activation: str = "gelu",
        evo2_style_activations: bool = True,
        # Linear biases
        mha_out_proj_bias: bool = True,
        hyena_out_proj_bias: bool = True,
        qkv_proj_bias: bool = False,
        # MLP geometry
        inner_size_multiple_of: int = 16,
        make_vocab_size_divisible_by: int = 8,
        # Embeddings
        tie_embeddings: bool = True,
        # Rotary / sequence
        max_seqlen: int = 8192,
        max_batch_size: int = 1,
        rotary_emb_base: float = 10000.0,
        use_interpolated_rotary_pos_emb: bool = False,
        rotary_emb_scaling_factor: float = 1.0,
        # Inference engine
        prefill_style: str = "fft",
        inference_mode: bool = True,
        # Projection precision (TE FP8 - required for 1b / 20b / 40b)
        use_fp8_input_projections: bool = False,
        # TE scaling recipe for those projections: "delayed" matches the original
        # Evo 2 (vortex); "current" also works when training across GPUs.
        fp8_recipe: str = "delayed",
        # Opt-in Triton kernels for the Hyena short / medium / long blocks
        # (vortex PR #77). Inference only; not bit-exact with the default path.
        use_hcs_kernel: bool = False,
        use_hcm_kernel: bool = False,
        use_hcl_kernel: bool = False,
        # Backend toggles
        use_cache: bool = True,
        # Vortex used this flag at runtime. Our HF port ignores it entirely
        # (attention dispatch is driven by ``config._attn_implementation``,
        # the standard HF mechanism). Default is False here so the config
        # reflects what actually runs: SDPA by default unless the user passes
        # ``attn_implementation="flash_attention_2"`` to from_pretrained.
        use_flash_attn: bool = False,
        use_flash_rmsnorm: bool = False,
        use_flash_depthwise: bool = False,
        use_flashfft: bool = False,
        # Per-block dtypes (cast at module init)
        attn_block_dtype: str = "bfloat16",
        hyena_block_dtype: str = "bfloat16",
        mlp_dtype: str = "bfloat16",
        # Multi-tensor parallel knobs (kept for ParallelGatedMLP._compute_inner_size)
        model_parallel_size: int = 1,
        pipe_parallel_size: int = 1,
        # GQA
        proj_groups: int = 1,
        smeared_gqa: bool = False,
        **kwargs,
    ):
        if attn_layer_idxs is None:
            attn_layer_idxs = [3, 10, 17, 24, 31]
        if hcl_layer_idxs is None:
            hcl_layer_idxs = [2, 6, 9, 13, 16, 20, 23, 27, 30]
        if hcm_layer_idxs is None:
            hcm_layer_idxs = [1, 5, 8, 12, 15, 19, 22, 26, 29]
        if hcs_layer_idxs is None:
            hcs_layer_idxs = [0, 4, 7, 11, 14, 18, 21, 25, 28]

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_filters = num_filters
        self.inner_mlp_size = inner_mlp_size
        self.num_layers = num_layers
        self.num_attention_heads = num_attention_heads

        self.attn_layer_idxs = attn_layer_idxs
        self.hcl_layer_idxs = hcl_layer_idxs
        self.hcm_layer_idxs = hcm_layer_idxs
        self.hcs_layer_idxs = hcs_layer_idxs

        self.hcm_filter_length = hcm_filter_length
        self.hcs_filter_length = hcs_filter_length
        self.hcl_filter_groups = hcl_filter_groups
        self.hcm_filter_groups = hcm_filter_groups
        self.hcs_filter_groups = hcs_filter_groups
        self.hyena_filter_groups = hyena_filter_groups
        self.short_filter_length = short_filter_length
        self.short_filter_bias = short_filter_bias
        self.state_size = state_size

        self.column_split = column_split
        self.column_split_hyena = column_split_hyena
        self.interleave = interleave
        self.hyena_flip_x1x2 = hyena_flip_x1x2

        self.eps = eps
        self.final_norm = final_norm
        self.mlp_activation = mlp_activation
        self.evo2_style_activations = evo2_style_activations

        self.mha_out_proj_bias = mha_out_proj_bias
        self.hyena_out_proj_bias = hyena_out_proj_bias
        self.qkv_proj_bias = qkv_proj_bias

        self.inner_size_multiple_of = inner_size_multiple_of
        self.make_vocab_size_divisible_by = make_vocab_size_divisible_by

        self.tie_embeddings = tie_embeddings

        self.max_seqlen = max_seqlen
        self.max_batch_size = max_batch_size
        self.rotary_emb_base = rotary_emb_base
        self.use_interpolated_rotary_pos_emb = use_interpolated_rotary_pos_emb
        self.rotary_emb_scaling_factor = rotary_emb_scaling_factor

        self.prefill_style = prefill_style
        self.inference_mode = inference_mode

        self.use_fp8_input_projections = use_fp8_input_projections
        if fp8_recipe not in ("delayed", "current"):
            raise ValueError(f"fp8_recipe must be 'delayed' or 'current', got {fp8_recipe!r}")
        self.fp8_recipe = fp8_recipe
        self.use_hcs_kernel = use_hcs_kernel
        self.use_hcm_kernel = use_hcm_kernel
        self.use_hcl_kernel = use_hcl_kernel

        self.use_cache = use_cache
        self.use_flash_attn = use_flash_attn
        self.use_flash_rmsnorm = use_flash_rmsnorm
        self.use_flash_depthwise = use_flash_depthwise
        self.use_flashfft = use_flashfft

        self.attn_block_dtype = attn_block_dtype
        self.hyena_block_dtype = hyena_block_dtype
        self.mlp_dtype = mlp_dtype

        self.model_parallel_size = model_parallel_size
        self.pipe_parallel_size = pipe_parallel_size

        self.proj_groups = proj_groups
        self.smeared_gqa = smeared_gqa

        super().__init__(**kwargs)

    # HF generation helpers expect `num_hidden_layers`.
    @property
    def num_hidden_layers(self) -> int:
        return self.num_layers

    # Internal blocks were originally written against a dotdict and call
    # `config.get(key, default)` extensively; provide a dict-like getter.
    def get(self, key, default=None):
        return getattr(self, key, default)

    @classmethod
    def from_original_config(cls, config_path: str, **kwargs) -> "Evo2Config":
        with open(config_path, "r") as f:
            config = json.load(f)
        return cls(**config, **kwargs)