"""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, ): # `gradient_checkpointing` is a training-time switch that the training pipeline # stores in its config dump. It is not an architecture property; enable it with # `model.gradient_checkpointing_enable()` instead. 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"]