|
|
|
|
|
from transformers import PretrainedConfig |
|
|
|
|
|
|
|
|
class NSAConfig(PretrainedConfig): |
|
|
model_type = "nsa" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
vocab_size=50257, |
|
|
hidden_size=768, |
|
|
num_hidden_layers=12, |
|
|
num_attention_heads=12, |
|
|
n_kv_groups=1, |
|
|
d_k=64, |
|
|
d_v=64, |
|
|
max_position_embeddings=2048, |
|
|
rope_theta=10000, |
|
|
nsa=None, |
|
|
**kwargs, |
|
|
): |
|
|
super().__init__(**kwargs) |
|
|
self.vocab_size = vocab_size |
|
|
self.hidden_size = hidden_size |
|
|
self.num_hidden_layers = num_hidden_layers |
|
|
self.num_attention_heads = num_attention_heads |
|
|
self.n_kv_groups = n_kv_groups |
|
|
self.d_k = d_k |
|
|
self.d_v = d_v |
|
|
self.max_position_embeddings = max_position_embeddings |
|
|
self.rope_theta = rope_theta |
|
|
self.nsa = nsa or { |
|
|
"branches": ["cmp", "sel", "win"], |
|
|
"window": 512, |
|
|
"gqa_groups": n_kv_groups, |
|
|
"block": 32, |
|
|
"stride": 16, |
|
|
"sel_block": 64, |
|
|
"sel_top_n": 16, |
|
|
} |
|
|
|
