File size: 5,748 Bytes
0eca2aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | # Apache 2.0 - port of togethercomputer's StripedHyenaConfig.
"""Configuration for Evo1 (StripedHyena 7B family)."""
from __future__ import annotations
import json
from transformers import PretrainedConfig
class Evo1Config(PretrainedConfig):
"""Configuration for the Evo1 family.
Defaults match the evo-1-8k-base / evo-1-131k-base / evo-1.5-8k-base
checkpoints (32 layers, 4096 hidden, 32 heads, attn at idx [8, 16, 24]
and Hyena everywhere else, byte-level vocab_size=512). The 131k variant
overrides ``use_interpolated_rotary_pos_emb`` and ``rotary_emb_scaling_factor``
plus a longer ``max_seqlen``.
"""
model_type = "evo1"
def __init__(
self,
# Architecture
vocab_size: int = 512,
hidden_size: int = 4096,
num_filters: int = 4096,
inner_mlp_size: int = 10928,
attn_layer_idxs=None,
hyena_layer_idxs=None,
num_layers: int = 32,
num_attention_heads: int = 32,
proj_groups: int = 1,
hyena_filter_groups: int = 1,
short_filter_length: int = 3,
short_filter_bias: bool = True,
state_size: int = 8,
column_split: bool = False,
column_split_hyena: bool = True,
split_k0: bool = True,
smeared_gqa: bool = False,
# Norms
eps: float = 1e-6,
final_norm: bool = True,
# Linear biases
mha_out_proj_bias: bool = True,
qkv_proj_bias: bool = True,
# Embeddings
tie_embeddings: bool = True,
make_vocab_size_divisible_by: int = 8,
# Activations
mlp_activation: str = "gelu",
# Sequence length / RoPE
max_seqlen: int = 8192,
rotary_emb_base: float = 10000,
use_interpolated_rotary_pos_emb: bool = False,
rotary_emb_scaling_factor: float = 1.0,
# Inference engine
prefill_style: str = "fft",
inference_mode: bool = False,
# Backend toggles
use_cache: bool = True,
use_flash_attention_2: bool = True,
use_flash_rmsnorm: bool = False,
use_flash_depthwise: bool = False,
use_flashfft: bool = False,
use_flash_attn: bool = False,
# Misc
log_intermediate_values: bool = False,
model_parallel_size: int = 1,
pipe_parallel_size: int = 1,
**kwargs,
):
if attn_layer_idxs is None:
attn_layer_idxs = [8, 16, 24]
if hyena_layer_idxs is None:
hyena_layer_idxs = [i for i in range(num_layers) if i not in attn_layer_idxs]
# Architecture
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_filters = num_filters
self.inner_mlp_size = inner_mlp_size
self.attn_layer_idxs = attn_layer_idxs
self.hyena_layer_idxs = hyena_layer_idxs
self.num_layers = num_layers
self.num_attention_heads = num_attention_heads
self.proj_groups = proj_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.split_k0 = split_k0
self.smeared_gqa = smeared_gqa
# Norms
self.eps = eps
self.final_norm = final_norm
# Biases
self.mha_out_proj_bias = mha_out_proj_bias
self.qkv_proj_bias = qkv_proj_bias
# Embeddings
self.tie_embeddings = tie_embeddings
self.make_vocab_size_divisible_by = make_vocab_size_divisible_by
# Activations
self.mlp_activation = mlp_activation
# Length / RoPE
self.max_seqlen = max_seqlen
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
# Engine
self.prefill_style = prefill_style
self.inference_mode = inference_mode
# Backend toggles
self.use_cache = use_cache
self.use_flash_attention_2 = use_flash_attention_2
self.use_flash_rmsnorm = use_flash_rmsnorm
self.use_flash_depthwise = use_flash_depthwise
self.use_flashfft = use_flashfft
self.use_flash_attn = use_flash_attn
# Misc
self.log_intermediate_values = log_intermediate_values
self.model_parallel_size = model_parallel_size
self.pipe_parallel_size = pipe_parallel_size
super().__init__(**kwargs)
# ------------------------------------------------------------------
# Backwards-compatible attribute access.
#
# The internal blocks (RMSNorm, ParallelGatedMLP, ...) call
# ``config.get(key, default)`` because they were originally written
# against a `dotdict`. PretrainedConfig has a different `.get`, so we
# provide a dict-like one that delegates to attribute access.
# ------------------------------------------------------------------
@property
def num_hidden_layers(self) -> int:
# HF generation utilities (DynamicCache, etc.) expect this name; we
# keep ``num_layers`` as the source of truth to match the upstream
# StripedHyena config.
return self.num_layers
def get(self, key, default=None):
# Dict-style access used by internal blocks (RMSNorm, MHA, ...).
return getattr(self, key, default)
@classmethod
def from_original_config(cls, config_path: str, **kwargs) -> "Evo1Config":
with open(config_path, "r") as f:
config = json.load(f)
return cls(**config, **kwargs)
|