File size: 13,714 Bytes
2719787 | 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 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 | """Configuration for the owned dual-attention decoder LM."""
from dataclasses import dataclass
from .attention_config import SUPPORTED_SEQUENCE_BOUNDARY_POLICIES
SUPPORTED_PE_TYPES = {"sinusoidal", "learned", "relative", "rope", "none"}
SUPPORTED_SYMBOL_RETRIEVAL = {"symbolic", "positional", "relative", "relsymbolic"}
SUPPORTED_RA_TYPES = {"ra", "rca", "disrca"}
SUPPORTED_RA_ACTIVATIONS = {"softmax", "identity", "relu", "tanh", "sigmoid", "gelu"}
SUPPORTED_FFN_ACTIVATIONS = {"gelu", "relu", "swiglu", "identity"}
SUPPORTED_FFN_HIDDEN_DIM_MODES = {"dff_factor", "swiglu_parameter_matched"}
SUPPORTED_INIT_SCHEMES = {"xavier_uniform", "normal_0_02_scaled_projection"}
@dataclass(frozen=True)
class DatLMConfig:
vocab_size: int
max_seq_len: int
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
sequence_boundary_policy: str = "eos_document"
segment_boundary_token_id: 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
pad_token_id: int = 0
bos_token_id: int = 1
eos_token_id: int = 2
mlm_head_enabled: bool = False
def __post_init__(self) -> None:
if self.vocab_size <= 0:
raise ValueError(f"vocab_size must be positive, got {self.vocab_size}")
if self.max_seq_len <= 0:
raise ValueError(f"max_seq_len must be positive, got {self.max_seq_len}")
if self.pe_type not in SUPPORTED_PE_TYPES:
raise ValueError(f"Unsupported pe_type: {self.pe_type}")
if self.hidden_dim <= 0:
raise ValueError(f"hidden_dim must be positive, got {self.hidden_dim}")
if self.n_heads_sa <= 0:
raise ValueError(f"n_heads_sa must be positive for DAT, got {self.n_heads_sa}")
if self.n_heads_ra <= 0:
raise ValueError(f"n_heads_ra must be positive for DAT, got {self.n_heads_ra}")
total_heads = self.total_n_heads
if self.hidden_dim % total_heads != 0:
raise ValueError(
f"hidden_dim ({self.hidden_dim}) must be divisible by total DAT heads "
f"({total_heads} = {self.n_heads_sa} SA + {self.n_heads_ra} RA)"
)
if self.n_layers <= 0:
raise ValueError(f"n_layers must be positive, got {self.n_layers}")
if not 0.0 <= self.dropout < 1.0:
raise ValueError(f"dropout must be in [0.0, 1.0), got {self.dropout}")
if self.dff_factor <= 0:
raise ValueError(f"dff_factor must be positive, got {self.dff_factor}")
if self.ffn_hidden_dim_mode not in SUPPORTED_FFN_HIDDEN_DIM_MODES:
raise ValueError(f"Unsupported ffn_hidden_dim_mode for DAT: {self.ffn_hidden_dim_mode}")
if self.ffn_activation not in SUPPORTED_FFN_ACTIVATIONS:
raise ValueError(f"Unsupported ffn_activation for DAT: {self.ffn_activation}")
if self.ffn_hidden_dim_mode == "swiglu_parameter_matched" and self.ffn_activation != "swiglu":
raise ValueError(
"ffn_hidden_dim_mode='swiglu_parameter_matched' requires "
f"ffn_activation='swiglu', got {self.ffn_activation}"
)
if self.rope_theta <= 0.0:
raise ValueError(f"rope_theta must be positive, got {self.rope_theta}")
if self.max_rel_pos is not None and self.max_rel_pos <= 0:
raise ValueError(f"max_rel_pos must be positive when provided, got {self.max_rel_pos}")
if self.sequence_boundary_policy not in SUPPORTED_SEQUENCE_BOUNDARY_POLICIES:
raise ValueError(
"Unsupported sequence_boundary_policy for DAT: "
f"{self.sequence_boundary_policy}"
)
if self.sequence_boundary_policy == "segment_document":
if self.segment_boundary_token_id is None:
raise ValueError(
"segment_boundary_token_id is required when "
"sequence_boundary_policy='segment_document'"
)
if self.init_range <= 0.0:
raise ValueError(f"init_range must be positive, got {self.init_range}")
if self.init_scheme not in SUPPORTED_INIT_SCHEMES:
raise ValueError(f"Unsupported init_scheme for DAT: {self.init_scheme}")
if self.norm_type not in {"layernorm", "rmsnorm"}:
raise ValueError(
f"norm_type must be 'layernorm' or 'rmsnorm', got {self.norm_type}"
)
if self.pe_type == "rope" and self.head_dim % 2 != 0:
raise ValueError(
"RoPE requires even DAT head_dim, got "
f"{self.head_dim} from hidden_dim={self.hidden_dim}, total_heads={total_heads}"
)
if self.pe_type == "sinusoidal" and self.hidden_dim % 2 != 0:
raise ValueError(f"Sinusoidal encoding requires even hidden_dim, got {self.hidden_dim}")
if self.share_attn_params and self.n_heads_sa != self.n_heads_ra:
raise ValueError(
"share_attn_params=True requires n_heads_sa == n_heads_ra, "
f"got {self.n_heads_sa} and {self.n_heads_ra}"
)
if self.symbol_dim is not None and self.symbol_dim <= 0:
raise ValueError(f"symbol_dim must be positive when provided, got {self.symbol_dim}")
if self.symbolic_attn_n_heads is not None and self.symbolic_attn_n_heads <= 0:
raise ValueError(
"symbolic_attn_n_heads must be positive when provided, "
f"got {self.symbolic_attn_n_heads}"
)
if self.symbol_retrieval not in SUPPORTED_SYMBOL_RETRIEVAL:
raise ValueError(f"Unsupported symbol_retrieval for DAT: {self.symbol_retrieval}")
if self.symbol_retrieval == "symbolic":
symbolic_heads = self.resolved_symbolic_attn_n_heads
if self.hidden_dim % symbolic_heads != 0:
raise ValueError(
f"hidden_dim ({self.hidden_dim}) must be divisible by symbolic_attn_n_heads "
f"({symbolic_heads}) for symbolic retrieval"
)
if self.resolved_symbol_dim % symbolic_heads != 0:
raise ValueError(
f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by symbolic_attn_n_heads "
f"({symbolic_heads}) for symbolic retrieval"
)
if (
self.symbol_retrieval == "positional"
and self.positional_symbols_sinusoidal
and self.resolved_symbol_dim % 2 != 0
):
raise ValueError(
"Sinusoidal positional symbols require even symbol_dim, "
f"got {self.resolved_symbol_dim}"
)
if self.relative_symbols_rope:
if self.symbol_retrieval != "relative":
raise ValueError(
"relative_symbols_rope=True requires symbol_retrieval='relative', "
f"got {self.symbol_retrieval!r}"
)
if self.resolved_symbol_dim % 2 != 0:
raise ValueError(
"RoPE relative symbols require even symbol_dim, "
f"got {self.resolved_symbol_dim}"
)
if self.positional_symbols_sinusoidal and self.symbol_retrieval != "positional":
raise ValueError(
"positional_symbols_sinusoidal=True requires symbol_retrieval='positional', "
f"got {self.symbol_retrieval!r}"
)
if self.resolved_n_symbols <= 0:
raise ValueError(f"resolved_n_symbols must be positive, got {self.resolved_n_symbols}")
if self.symbol_retrieval == "relsymbolic":
if self.relsymbolic_rel_n_heads <= 0:
raise ValueError(
f"relsymbolic_rel_n_heads must be positive, got {self.relsymbolic_rel_n_heads}"
)
if self.hidden_dim % self.relsymbolic_rel_n_heads != 0:
raise ValueError(
f"hidden_dim ({self.hidden_dim}) must be divisible by "
f"relsymbolic_rel_n_heads ({self.relsymbolic_rel_n_heads})"
)
if self.relsymbolic_symbolic_attn_n_heads <= 0:
raise ValueError(
"relsymbolic_symbolic_attn_n_heads must be positive, "
f"got {self.relsymbolic_symbolic_attn_n_heads}"
)
if self.hidden_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
raise ValueError(
f"hidden_dim ({self.hidden_dim}) must be divisible by "
f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
)
if self.resolved_symbol_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
raise ValueError(
f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by "
f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
)
if self.relsymbolic_neighborhood_size <= 0:
raise ValueError(
"relsymbolic_neighborhood_size must be positive, "
f"got {self.relsymbolic_neighborhood_size}"
)
if not 0.0 <= self.relsymbolic_dropout < 1.0:
raise ValueError(
f"relsymbolic_dropout must be in [0.0, 1.0), got {self.relsymbolic_dropout}"
)
if self.relsymbolic_rel_scale is not None and self.relsymbolic_rel_scale <= 0.0:
raise ValueError(
"relsymbolic_rel_scale must be positive when provided, "
f"got {self.relsymbolic_rel_scale}"
)
if (
self.relsymbolic_symbolic_attn_scale is not None
and self.relsymbolic_symbolic_attn_scale <= 0.0
):
raise ValueError(
"relsymbolic_symbolic_attn_scale must be positive when provided, "
f"got {self.relsymbolic_symbolic_attn_scale}"
)
if self.ra_type not in SUPPORTED_RA_TYPES:
raise ValueError(f"Unsupported ra_type for DAT: {self.ra_type}")
if self.ra_n_relations is not None and self.ra_n_relations <= 0:
raise ValueError(
f"ra_n_relations must be positive when provided, got {self.ra_n_relations}"
)
if self.ra_type != "ra" and self.ra_n_relations is not None:
raise ValueError(f"ra_n_relations applies only to ra_type='ra', got {self.ra_type}")
if self.ra_type != "ra" and self.ra_symmetric_rels:
raise ValueError(f"ra_symmetric_rels applies only to ra_type='ra', got {self.ra_type}")
n_relations = self.resolved_ra_n_relations
if self.ra_type == "ra" and (self.head_dim * self.n_heads_ra) % n_relations != 0:
raise ValueError(
f"head_dim * n_heads_ra ({self.head_dim * self.n_heads_ra}) must be "
f"divisible by ra_n_relations ({n_relations})"
)
if self.ra_rel_activation not in SUPPORTED_RA_ACTIVATIONS:
raise ValueError(
f"Unsupported ra_rel_activation for DAT: {self.ra_rel_activation}"
)
@property
def total_n_heads(self) -> int:
return self.n_heads_sa + self.n_heads_ra
@property
def head_dim(self) -> int:
return self.hidden_dim // self.total_n_heads
@property
def resolved_symbol_dim(self) -> int:
return self.hidden_dim if self.symbol_dim is None else self.symbol_dim
@property
def resolved_symbolic_attn_n_heads(self) -> int:
return self.total_n_heads if self.symbolic_attn_n_heads is None else self.symbolic_attn_n_heads
@property
def resolved_ffn_hidden_dim(self) -> int:
if self.ffn_hidden_dim_mode == "swiglu_parameter_matched":
return int(8 / 3 * self.hidden_dim)
return self.hidden_dim * self.dff_factor
@property
def resolved_n_symbols(self) -> int:
return self.max_seq_len if self.n_symbols is None else self.n_symbols
@property
def resolved_ra_n_relations(self) -> int:
return self.n_heads_ra if self.ra_n_relations is None else self.ra_n_relations
|