File size: 7,768 Bytes
5ecd8e9 | 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 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | """HuggingFace PretrainedConfig for the dual-attention (DAT) decoder LM.
This is the source-of-truth copy. scripts/convert_dat_to_hf.py copies it into a
generated HF repository (alongside modeling_dat.py and the flattened model
source) so the model can be loaded with
``AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True)``.
The config simply carries every field of models.attention.dat.config.DatLMConfig
so that modeling_dat.py can rebuild the exact DatLMConfig at load time.
"""
from transformers import PretrainedConfig
# Every field needed to rebuild DatLMConfig. New exported fields are appended
# after existing constructor parameters to preserve positional compatibility.
# modeling_dat.py rebuilds DatLMConfig as
# DatLMConfig(**{f: getattr(config, f) for f in DAT_LM_FIELDS}).
DAT_LM_FIELDS = (
"vocab_size",
"max_seq_len",
"pe_type",
"hidden_dim",
"n_heads_sa",
"n_heads_ra",
"n_layers",
"dropout",
"dff_factor",
"ffn_hidden_dim_mode",
"ffn_activation",
"rope_theta",
"max_rel_pos",
"init_range",
"init_scheme",
"norm_type",
"norm_first",
"use_bias_qkv",
"use_bias_out",
"use_bias_ffn",
"tie_lm_head",
"symbol_dim",
"n_symbols",
"symbolic_attn_n_heads",
"symbol_retrieval",
"symbolic_use_bias",
"shared_symbol_retriever",
"share_attn_params",
"positional_symbols_sinusoidal",
"relative_symbols_rope",
"relsymbolic_rel_n_heads",
"relsymbolic_symbolic_attn_n_heads",
"relsymbolic_neighborhood_size",
"relsymbolic_include_self",
"relsymbolic_normalize_rels",
"relsymbolic_trainable_symbols",
"relsymbolic_dropout",
"relsymbolic_rel_scale",
"relsymbolic_symbolic_attn_scale",
"relsymbolic_use_bias",
"ra_type",
"ra_n_relations",
"ra_rel_activation",
"ra_symmetric_rels",
"sequence_boundary_policy",
"segment_boundary_token_id",
"pad_token_id",
"bos_token_id",
"eos_token_id",
"mlm_head_enabled",
)
class DatConfig(PretrainedConfig):
model_type = "dat"
# Expose the model width under HF's conventional name so generic eval
# harnesses (e.g. babylm-eval's finetuning head, which reads
# config.hidden_size) can size their classification head. The model's own
# field is hidden_dim; this maps the alias onto it for both get and set.
attribute_map = {"hidden_size": "hidden_dim"}
def __init__(
self,
vocab_size: int = 16384,
max_seq_len: int = 513,
pe_type: str = "rope",
hidden_dim: int = 256,
n_heads_sa: int = 2,
n_heads_ra: int = 2,
n_layers: int = 4,
dropout: float = 0.0,
dff_factor: int = 4,
ffn_hidden_dim_mode: str = "dff_factor",
ffn_activation: str = "gelu",
rope_theta: float = 10000.0,
max_rel_pos: int | None = None,
init_range: float = 0.15,
init_scheme: str = "xavier_uniform",
norm_type: str = "rmsnorm",
norm_first: bool = True,
use_bias_qkv: bool = False,
use_bias_out: bool = True,
use_bias_ffn: bool = True,
tie_lm_head: bool = True,
symbol_dim: int | None = None,
n_symbols: int | None = None,
symbolic_attn_n_heads: int | None = None,
symbol_retrieval: str = "symbolic",
symbolic_use_bias: bool = False,
shared_symbol_retriever: bool = True,
share_attn_params: bool = False,
positional_symbols_sinusoidal: bool = False,
relative_symbols_rope: bool = False,
relsymbolic_rel_n_heads: int = 4,
relsymbolic_symbolic_attn_n_heads: int = 4,
relsymbolic_neighborhood_size: int = 2,
relsymbolic_include_self: bool = False,
relsymbolic_normalize_rels: bool = True,
relsymbolic_trainable_symbols: bool = True,
relsymbolic_dropout: float = 0.0,
relsymbolic_rel_scale: float | None = None,
relsymbolic_symbolic_attn_scale: float | None = None,
relsymbolic_use_bias: bool = False,
ra_type: str = "ra",
ra_n_relations: int | None = None,
ra_rel_activation: str = "identity",
ra_symmetric_rels: bool = False,
sequence_boundary_policy: str = "eos_document",
segment_boundary_token_id: int | None = None,
pad_token_id: int = 0,
bos_token_id: int = 1,
eos_token_id: int = 2,
mlm_head_enabled: bool = False,
**kwargs,
) -> None:
self.vocab_size = vocab_size
self.max_seq_len = max_seq_len
self.pe_type = pe_type
self.hidden_dim = hidden_dim
self.n_heads_sa = n_heads_sa
self.n_heads_ra = n_heads_ra
self.n_layers = n_layers
self.dropout = dropout
self.dff_factor = dff_factor
self.ffn_hidden_dim_mode = ffn_hidden_dim_mode
self.ffn_activation = ffn_activation
self.rope_theta = rope_theta
self.max_rel_pos = max_rel_pos
self.init_range = init_range
self.init_scheme = init_scheme
self.norm_type = norm_type
self.norm_first = norm_first
self.use_bias_qkv = use_bias_qkv
self.use_bias_out = use_bias_out
self.use_bias_ffn = use_bias_ffn
self.tie_lm_head = tie_lm_head
self.symbol_dim = symbol_dim
self.n_symbols = n_symbols
self.symbolic_attn_n_heads = symbolic_attn_n_heads
self.symbol_retrieval = symbol_retrieval
self.symbolic_use_bias = symbolic_use_bias
self.shared_symbol_retriever = shared_symbol_retriever
self.share_attn_params = share_attn_params
self.positional_symbols_sinusoidal = positional_symbols_sinusoidal
self.relative_symbols_rope = relative_symbols_rope
self.relsymbolic_rel_n_heads = relsymbolic_rel_n_heads
self.relsymbolic_symbolic_attn_n_heads = relsymbolic_symbolic_attn_n_heads
self.relsymbolic_neighborhood_size = relsymbolic_neighborhood_size
self.relsymbolic_include_self = relsymbolic_include_self
self.relsymbolic_normalize_rels = relsymbolic_normalize_rels
self.relsymbolic_trainable_symbols = relsymbolic_trainable_symbols
self.relsymbolic_dropout = relsymbolic_dropout
self.relsymbolic_rel_scale = relsymbolic_rel_scale
self.relsymbolic_symbolic_attn_scale = relsymbolic_symbolic_attn_scale
self.relsymbolic_use_bias = relsymbolic_use_bias
self.ra_type = ra_type
self.ra_n_relations = ra_n_relations
self.ra_rel_activation = ra_rel_activation
self.ra_symmetric_rels = ra_symmetric_rels
self.sequence_boundary_policy = sequence_boundary_policy
self.segment_boundary_token_id = segment_boundary_token_id
self.mlm_head_enabled = mlm_head_enabled
# HF base also exposes max_position_embeddings for generation utilities.
self.max_position_embeddings = max_seq_len
# tie_lm_head is the source of truth for weight tying; keep HF's
# tie_word_embeddings in lockstep so from_pretrained never re-ties an
# untied head (or fails to tie a tied one). Drop any value coming in via
# kwargs (e.g. a serialized config.json) so the two cannot disagree.
# NOTE: tie_lm_head controls only lm_head <-> token_embeddings tying.
# The mlm_head.linear_out <-> lm_head tying is handled separately by
# the get_expanded_tied_weights_keys override in modeling_dat.py.
kwargs.pop("tie_word_embeddings", None)
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=tie_lm_head,
**kwargs,
)
|