| """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") |
|
|
