|
from transformers import PretrainedConfig |
|
import json |
|
|
|
|
|
class StripedHyenaConfig(PretrainedConfig): |
|
model_type = "stripedhyena" |
|
|
|
def __init__( |
|
self, |
|
vocab_size=32000, |
|
hidden_size=4096, |
|
num_filters=4096, |
|
inner_mlp_size=14336, |
|
attn_layer_idxs=[], |
|
hyena_layer_idxs=[], |
|
num_layers=32, |
|
tie_embeddings=False, |
|
short_filter_length=3, |
|
num_attention_heads=32, |
|
proj_groups=4, |
|
hyena_filter_groups=1, |
|
split_k0=True, |
|
column_split_hyena=True, |
|
column_split=False, |
|
model_parallel_size=1, |
|
pipe_parallel_size=1, |
|
short_filter_bias=True, |
|
mha_out_proj_bias=False, |
|
qkv_proj_bias=False, |
|
final_norm=True, |
|
use_cache=True, |
|
use_flash_attention_2=True, |
|
use_flash_rmsnorm=True, |
|
use_flash_depthwise=False, |
|
use_flashfft=False, |
|
inference_mode=False, |
|
prefill_style="fft", |
|
max_seqlen=32768, |
|
eps=1e-5, |
|
state_size=2, |
|
rotary_emb_base=500000, |
|
smeared_gqa=False, |
|
make_vocab_size_divisible_by=8, |
|
log_intermediate_values=False, |
|
**kwargs, |
|
): |
|
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.tie_embeddings = tie_embeddings |
|
self.short_filter_length = short_filter_length |
|
self.num_attention_heads = num_attention_heads |
|
self.proj_groups = proj_groups |
|
self.hyena_filter_groups = hyena_filter_groups |
|
self.split_k0 = split_k0 |
|
self.column_split_hyena = column_split_hyena |
|
self.column_split = column_split |
|
self.model_parallel_size = model_parallel_size |
|
self.pipe_parallel_size = pipe_parallel_size |
|
self.short_filter_bias = short_filter_bias |
|
self.mha_out_proj_bias = mha_out_proj_bias |
|
self.qkv_proj_bias = qkv_proj_bias |
|
self.final_norm = final_norm |
|
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.inference_mode = inference_mode |
|
self.prefill_style = prefill_style |
|
self.max_seqlen = max_seqlen |
|
self.eps = eps |
|
self.state_size = state_size |
|
self.rotary_emb_base = rotary_emb_base |
|
self.smeared_gqa = smeared_gqa |
|
self.make_vocab_size_divisible_by = make_vocab_size_divisible_by |
|
self.log_intermediate_values = log_intermediate_values |
|
super().__init__(**kwargs) |
|
|
|
def to_dict(self): |
|
return {attr: getattr(self, attr) for attr in self.__dict__} |
|
|
|
@classmethod |
|
def from_original_config(cls, config_path, **kwargs): |
|
with open(config_path, "r") as f: |
|
config = json.load(f) |
|
|
|
return cls(**config, **kwargs) |
|
|
