File size: 2,128 Bytes
4634a7d | 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 | """Hugging Face configuration for ARLoop."""
from transformers import PretrainedConfig
class ARLoopConfig(PretrainedConfig):
model_type = "arloop"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(self, vocab_size=49152, hidden_size=2048, intermediate_size=5632,
physical_layers=24, loop_count=3, num_attention_heads=16,
num_key_value_heads=4, max_position_embeddings=4096,
rope_theta=50000.0, rope_condense_ratio=1, rms_norm_eps=1e-5,
attnres_block_size=24, sw_blocks=3, qk_norm=True,
embedding_scale=1.0, logit_softcap=None, use_cache=True, **kwargs):
kwargs.pop("num_hidden_layers", None)
kwargs.pop("tie_word_embeddings", None)
super().__init__(tie_word_embeddings=True, **kwargs)
for key, value in dict(vocab_size=vocab_size, hidden_size=hidden_size,
intermediate_size=intermediate_size, physical_layers=physical_layers,
loop_count=loop_count, num_attention_heads=num_attention_heads,
num_key_value_heads=num_key_value_heads,
max_position_embeddings=max_position_embeddings, rope_theta=rope_theta,
rope_condense_ratio=rope_condense_ratio, rms_norm_eps=rms_norm_eps,
attnres_block_size=attnres_block_size, sw_blocks=sw_blocks,
qk_norm=qk_norm, embedding_scale=embedding_scale,
logit_softcap=logit_softcap, use_cache=use_cache).items():
setattr(self, key, value)
self.num_hidden_layers = physical_layers * loop_count
self.head_dim = hidden_size // num_attention_heads
if hidden_size % num_attention_heads or num_attention_heads % num_key_value_heads:
raise ValueError("Invalid grouped-query attention dimensions")
if self.head_dim % 2 or hidden_size < 32:
raise ValueError("RoPE needs even head_dim; value gating needs hidden_size >= 32")
if attnres_block_size < 2 or attnres_block_size % 2 or sw_blocks < 1:
raise ValueError("Invalid residual block/window size")
|