File size: 6,643 Bytes
9df61c9 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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"]
|