| """Configuration for Cherry Alpha.""" | |
| from __future__ import annotations | |
| from transformers import PretrainedConfig | |
| class CherryAlphaConfig(PretrainedConfig): | |
| model_type = "cherry_alpha" | |
| def __init__( | |
| self, | |
| vocab_size: int = 2048, | |
| hidden_size: int = 256, | |
| intermediate_size: int = 1488, | |
| num_hidden_layers: int = 3, | |
| num_attention_heads: int = 8, | |
| num_key_value_heads: int = 1, | |
| head_dim: int = 32, | |
| ngram_buckets: int = 156250, | |
| loop_schedule: tuple[int, ...] | list[int] = (0, 1, 1, 2), | |
| refresh_kernel_size: int = 9, | |
| max_position_embeddings: int = 4096, | |
| rope_theta: float = 100000.0, | |
| rms_norm_eps: float = 1e-5, | |
| initializer_range: float = 0.02, | |
| gradient_checkpointing: bool = True, | |
| tie_word_embeddings: bool = True, | |
| use_cache: bool = False, | |
| bos_token_id: int = 1, | |
| eos_token_id: int = 2, | |
| pad_token_id: int = 0, | |
| unk_token_id: int = 3, | |
| **kwargs, | |
| ): | |
| if hidden_size != num_attention_heads * head_dim: | |
| raise ValueError("hidden_size must equal num_attention_heads * head_dim") | |
| if num_attention_heads % num_key_value_heads: | |
| raise ValueError("num_attention_heads must be divisible by num_key_value_heads") | |
| if num_hidden_layers != 3: | |
| raise ValueError("Cherry Alpha requires exactly three physical layers") | |
| loop_schedule = tuple(int(index) for index in loop_schedule) | |
| if loop_schedule != (0, 1, 1, 2): | |
| raise ValueError("Cherry Alpha requires L1 -> L2 -> L2 -> L3") | |
| if ngram_buckets <= 0: | |
| raise ValueError("ngram_buckets must be positive") | |
| if refresh_kernel_size <= 0 or refresh_kernel_size % 2 == 0: | |
| raise ValueError("refresh_kernel_size must be a positive odd integer") | |
| self.vocab_size = vocab_size | |
| self.hidden_size = hidden_size | |
| self.intermediate_size = intermediate_size | |
| self.num_hidden_layers = num_hidden_layers | |
| self.num_attention_heads = num_attention_heads | |
| self.num_key_value_heads = num_key_value_heads | |
| self.head_dim = head_dim | |
| self.ngram_buckets = ngram_buckets | |
| self.loop_schedule = list(loop_schedule) | |
| self.refresh_kernel_size = refresh_kernel_size | |
| self.max_position_embeddings = max_position_embeddings | |
| self.rope_theta = rope_theta | |
| self.rms_norm_eps = rms_norm_eps | |
| self.initializer_range = initializer_range | |
| self.gradient_checkpointing = gradient_checkpointing | |
| self.use_cache = False | |
| super().__init__( | |
| tie_word_embeddings=tie_word_embeddings, | |
| bos_token_id=bos_token_id, | |
| eos_token_id=eos_token_id, | |
| pad_token_id=pad_token_id, | |
| unk_token_id=unk_token_id, | |
| **kwargs, | |
| ) | |
| def parameter_counts(self) -> dict[str, int]: | |
| hidden = self.hidden_size | |
| query_width = self.num_attention_heads * self.head_dim | |
| kv_width = self.num_key_value_heads * self.head_dim | |
| token_embedding = self.vocab_size * hidden | |
| ngram_embedding = self.ngram_buckets * hidden + 1 | |
| attention = ( | |
| hidden * query_width | |
| + 2 * hidden * kv_width | |
| + query_width * hidden | |
| + 2 * self.head_dim | |
| ) | |
| refresh = ( | |
| 3 * hidden * hidden | |
| + hidden * self.refresh_kernel_size | |
| + 3 * hidden | |
| + 1 | |
| ) | |
| mlp = 3 * hidden * self.intermediate_size | |
| block_norms = 2 * hidden | |
| physical_layer = attention + refresh + mlp + block_norms | |
| neural_core = ( | |
| token_embedding | |
| + self.num_hidden_layers * physical_layer | |
| + hidden | |
| ) | |
| total = neural_core + ngram_embedding | |
| if not self.tie_word_embeddings: | |
| total += token_embedding | |
| neural_core += token_embedding | |
| return { | |
| "total": total, | |
| "neural_core": neural_core, | |
| "ngram": ngram_embedding, | |
| "token_embedding": token_embedding, | |
| "physical_layer": physical_layer, | |
| "physical_layers": self.num_hidden_layers, | |
| "effective_layer_passes": len(self.loop_schedule), | |
| } | |
| CherryAlphaConfig.register_for_auto_class("AutoConfig") | |
| __all__ = ["CherryAlphaConfig"] | |