| """Configuration for Comet-1 (OpenModel-1), a hybrid attention / causal-convolution language model.""" |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| class CometConfig(PretrainedConfig): |
| r""" |
| Configuration class for [`CometModel`] and [`CometForCausalLM`]. |
| |
| The field names mirror the ``model_config`` block that the OpenModel-1 training |
| pipeline (``train.py``) writes next to every checkpoint, so a training-time |
| ``config.json`` loads unchanged. |
| |
| Args: |
| vocab_size (`int`, defaults to 50277): |
| Size of the token embedding table (``len(tokenizer)`` of the GPT-NeoX-20B BPE tokenizer). |
| hidden_size (`int`, defaults to 4096): |
| Width of the residual stream. |
| num_hidden_layers (`int`, defaults to 36): |
| Number of blocks. |
| num_attention_heads (`int`, defaults to 32): |
| Number of query heads in the "global" (attention) blocks. |
| num_key_value_heads (`int`, defaults to 8): |
| Number of key/value heads (grouped-query attention). |
| intermediate_size (`int`, defaults to 12288): |
| Inner width of the SwiGLU feed-forward. |
| max_position_embeddings (`int`, defaults to 8192): |
| Number of positions in the precomputed RoPE table. Inputs longer than this are rejected. |
| rope_theta (`float`, defaults to 1000000.0): |
| RoPE base frequency. |
| memory_width (`int`, defaults to 2048): |
| Channel count of the depthwise convolution lane in "recall" blocks. |
| memory_kernel_size (`int`, defaults to 7): |
| Kernel size of the depthwise causal convolution. |
| memory_dilations (`list[int]`, defaults to `[1, 2, 4, 8]`): |
| Dilation table. Block ``i`` uses ``memory_dilations[i % len(memory_dilations)]``. |
| layer_pattern (`list[str]`, defaults to `["global", "global", "global", "recall"]`): |
| Repeating block schedule. Block ``i`` is of kind ``layer_pattern[i % len(layer_pattern)]``; |
| ``"global"`` is grouped-query attention, anything else is the causal recall mixer. |
| layer_scale_init (`float`, defaults to 0.5): |
| Initial value of the learned per-channel residual scales. |
| qk_norm (`bool`, defaults to `True`): |
| Apply RMSNorm to queries and keys (per head) before RoPE. |
| rms_norm_eps (`float`, defaults to 1e-6): |
| Epsilon of every RMSNorm. |
| fused_projections (`bool`, defaults to `True`): |
| Parameter layout. ``True`` stores Q/K/V as one ``qkv_proj`` matrix and the two SwiGLU input |
| projections as one ``gate_up`` matrix (the layout of the released checkpoint). |
| dropout (`float`, defaults to 0.0): |
| Dropout on both residual branches (training only). |
| initializer_range (`float`, defaults to 0.02): |
| Standard deviation used when initialising new weights. |
| tie_word_embeddings (`bool`, defaults to `True`): |
| Share the input embedding matrix with the output projection. |
| use_cache (`bool`, defaults to `True`): |
| Return a key/value + convolution-state cache from the forward pass. |
| variant (`str`, *optional*): |
| Free-form size label written by the training pipeline (for example ``"7.0b"``). |
| """ |
|
|
| model_type = "comet" |
|
|
| def __init__( |
| self, |
| vocab_size=50277, |
| hidden_size=4096, |
| num_hidden_layers=36, |
| num_attention_heads=32, |
| num_key_value_heads=8, |
| intermediate_size=12288, |
| max_position_embeddings=8192, |
| rope_theta=1_000_000.0, |
| memory_width=2048, |
| memory_kernel_size=7, |
| memory_dilations=(1, 2, 4, 8), |
| layer_pattern=("global", "global", "global", "recall"), |
| layer_scale_init=0.5, |
| qk_norm=True, |
| rms_norm_eps=1e-6, |
| fused_projections=True, |
| dropout=0.0, |
| initializer_range=0.02, |
| tie_word_embeddings=True, |
| use_cache=True, |
| variant=None, |
| bos_token_id=0, |
| eos_token_id=0, |
| pad_token_id=1, |
| **kwargs, |
| ): |
| |
| |
| |
| kwargs.pop("gradient_checkpointing", None) |
|
|
| 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.num_key_value_heads = num_key_value_heads |
| self.intermediate_size = intermediate_size |
| self.max_position_embeddings = max_position_embeddings |
| self.rope_theta = float(rope_theta) |
| self.memory_width = memory_width |
| self.memory_kernel_size = memory_kernel_size |
| self.memory_dilations = list(memory_dilations) |
| self.layer_pattern = list(layer_pattern) |
| self.layer_scale_init = layer_scale_init |
| self.qk_norm = qk_norm |
| self.rms_norm_eps = rms_norm_eps |
| self.fused_projections = fused_projections |
| self.dropout = dropout |
| self.initializer_range = initializer_range |
| self.use_cache = use_cache |
| self.variant = variant |
|
|
| if hidden_size % num_attention_heads != 0: |
| raise ValueError("hidden_size must be divisible by num_attention_heads") |
| if (hidden_size // num_attention_heads) % 2 != 0: |
| raise ValueError("the attention head dimension must be even for RoPE") |
| if num_attention_heads % num_key_value_heads != 0: |
| raise ValueError("num_attention_heads must be divisible by num_key_value_heads") |
| if memory_kernel_size < 1: |
| raise ValueError("memory_kernel_size must be positive") |
| if len(self.layer_pattern) == 0 or len(self.memory_dilations) == 0: |
| raise ValueError("layer_pattern and memory_dilations must not be empty") |
|
|
| super().__init__( |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| pad_token_id=pad_token_id, |
| tie_word_embeddings=tie_word_embeddings, |
| **kwargs, |
| ) |
|
|
| def block_kind(self, index): |
| """Kind of block ``index``: ``"global"`` (attention) or ``"recall"`` (causal convolution).""" |
| return self.layer_pattern[index % len(self.layer_pattern)] |
|
|
| def block_dilation(self, index): |
| """Convolution dilation that block ``index`` uses if it is a recall block.""" |
| return self.memory_dilations[index % len(self.memory_dilations)] |
|
|
|
|
| __all__ = ["CometConfig"] |
|
|
