235 lines
8.9 KiB
Python
235 lines
8.9 KiB
Python
"""
|
||
CyberSLM Model Configuration
|
||
=============================
|
||
Defines the complete, validated configuration for the CyberSLM decoder-only
|
||
transformer. All hyperparameters are frozen at construction time and validated
|
||
for mathematical consistency before any model component is instantiated.
|
||
|
||
Architecture summary
|
||
--------------------
|
||
- Hidden dim : 384
|
||
- Decoder layers : 12
|
||
- Attention heads : 6
|
||
- Head dim : 64 (hidden_dim / num_heads = 384 / 6 = 64)
|
||
- FFN inner dim : 1024 (SwiGLU gate + value = 2 × 1024 → projects back to 384)
|
||
- Vocab size : 32 000
|
||
- Max context : 4 096
|
||
- Approx params : 33.53 M
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass, field
|
||
from typing import Optional
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class CyberSLMConfig:
|
||
"""
|
||
Immutable configuration for the CyberSLM decoder-only transformer.
|
||
|
||
All fields are set once at construction; the frozen dataclass guarantees
|
||
no accidental mutation during training. Call :meth:`validate` immediately
|
||
after construction or use the convenience constructor
|
||
:func:`default_config`.
|
||
|
||
Attributes
|
||
----------
|
||
vocab_size : int
|
||
Number of tokens in the SentencePiece BPE vocabulary.
|
||
max_seq_len : int
|
||
Maximum token sequence length (context window).
|
||
hidden_dim : int
|
||
Embedding dimension ``d_model``.
|
||
num_layers : int
|
||
Number of stacked transformer decoder blocks.
|
||
num_heads : int
|
||
Number of attention heads. Must evenly divide ``hidden_dim``.
|
||
head_dim : int
|
||
Dimension of each attention head. Must equal ``hidden_dim // num_heads``.
|
||
ffn_hidden_dim : int
|
||
Inner dimension of the SwiGLU feed-forward network.
|
||
The gate and value projections each map hidden_dim → ffn_hidden_dim,
|
||
and the output projection maps ffn_hidden_dim → hidden_dim.
|
||
rope_base : int
|
||
Base for Rotary Position Embedding frequency computation (θ = 10 000).
|
||
norm_eps : float
|
||
Epsilon added inside RMSNorm to prevent division by zero.
|
||
tie_weights : bool
|
||
When True the output projection shares weights with the token embedding.
|
||
bias : bool
|
||
When True linear layers include a bias term (False = modern practice).
|
||
dropout : float
|
||
Residual / feed-forward dropout probability (0.0 = disabled).
|
||
attn_dropout : float
|
||
Attention weight dropout probability (0.0 = disabled).
|
||
pad_token_id : Optional[int]
|
||
Token ID used for padding; None if the dataset never pads.
|
||
bos_token_id : Optional[int]
|
||
Beginning-of-sequence token ID.
|
||
eos_token_id : Optional[int]
|
||
End-of-sequence token ID.
|
||
"""
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Vocabulary & sequence #
|
||
# ------------------------------------------------------------------ #
|
||
vocab_size: int = 32_000
|
||
max_seq_len: int = 4_096
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Transformer dimensions #
|
||
# ------------------------------------------------------------------ #
|
||
hidden_dim: int = 384
|
||
num_layers: int = 12
|
||
num_heads: int = 6
|
||
head_dim: int = 64 # must equal hidden_dim // num_heads
|
||
ffn_hidden_dim: int = 1_024
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Positional encoding #
|
||
# ------------------------------------------------------------------ #
|
||
rope_base: int = 10_000
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Normalization #
|
||
# ------------------------------------------------------------------ #
|
||
norm_eps: float = 1e-6
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Architecture flags #
|
||
# ------------------------------------------------------------------ #
|
||
tie_weights: bool = True
|
||
bias: bool = False
|
||
dropout: float = 0.0
|
||
attn_dropout: float = 0.0
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Special token IDs (set by tokenizer integration layer) #
|
||
# ------------------------------------------------------------------ #
|
||
pad_token_id: Optional[int] = None
|
||
bos_token_id: Optional[int] = 2 # real SentencePiece BOS id
|
||
eos_token_id: Optional[int] = 3 # real SentencePiece EOS id
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Validation #
|
||
# ------------------------------------------------------------------ #
|
||
def validate(self) -> "CyberSLMConfig":
|
||
"""
|
||
Assert mathematical consistency of every hyperparameter.
|
||
|
||
Raises
|
||
------
|
||
ValueError
|
||
If any hyperparameter violates an architectural constraint.
|
||
|
||
Returns
|
||
-------
|
||
CyberSLMConfig
|
||
Self, to allow chaining: ``cfg = CyberSLMConfig().validate()``.
|
||
"""
|
||
errors: list[str] = []
|
||
|
||
# Positivity checks
|
||
for name, value in [
|
||
("vocab_size", self.vocab_size),
|
||
("max_seq_len", self.max_seq_len),
|
||
("hidden_dim", self.hidden_dim),
|
||
("num_layers", self.num_layers),
|
||
("num_heads", self.num_heads),
|
||
("head_dim", self.head_dim),
|
||
("ffn_hidden_dim", self.ffn_hidden_dim),
|
||
("rope_base", self.rope_base),
|
||
]:
|
||
if value <= 0:
|
||
errors.append(f"{name} must be positive, got {value}")
|
||
|
||
# Attention head consistency
|
||
if self.hidden_dim % self.num_heads != 0:
|
||
errors.append(
|
||
f"hidden_dim ({self.hidden_dim}) must be divisible by "
|
||
f"num_heads ({self.num_heads})"
|
||
)
|
||
expected_head_dim = self.hidden_dim // self.num_heads
|
||
if self.head_dim != expected_head_dim:
|
||
errors.append(
|
||
f"head_dim ({self.head_dim}) must equal "
|
||
f"hidden_dim // num_heads = {expected_head_dim}"
|
||
)
|
||
|
||
# RoPE requires even head_dim (pairs of sin/cos)
|
||
if self.head_dim % 2 != 0:
|
||
errors.append(
|
||
f"head_dim ({self.head_dim}) must be even for RoPE"
|
||
)
|
||
|
||
# Dropout bounds
|
||
for name, value in [("dropout", self.dropout), ("attn_dropout", self.attn_dropout)]:
|
||
if not (0.0 <= value < 1.0):
|
||
errors.append(f"{name} must be in [0, 1), got {value}")
|
||
|
||
# norm_eps positivity
|
||
if self.norm_eps <= 0.0:
|
||
errors.append(f"norm_eps must be positive, got {self.norm_eps}")
|
||
|
||
if errors:
|
||
raise ValueError(
|
||
"CyberSLMConfig validation failed:\n"
|
||
+ "\n".join(f" • {e}" for e in errors)
|
||
)
|
||
|
||
return self
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Derived properties #
|
||
# ------------------------------------------------------------------ #
|
||
@property
|
||
def total_attention_dim(self) -> int:
|
||
"""``num_heads × head_dim`` — equals ``hidden_dim`` by construction."""
|
||
return self.num_heads * self.head_dim
|
||
|
||
@property
|
||
def rope_half_dim(self) -> int:
|
||
"""Number of frequency pairs in RoPE (``head_dim // 2``)."""
|
||
return self.head_dim // 2
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Display #
|
||
# ------------------------------------------------------------------ #
|
||
def __str__(self) -> str:
|
||
lines = [
|
||
"CyberSLMConfig",
|
||
"=" * 40,
|
||
f" vocab_size : {self.vocab_size:,}",
|
||
f" max_seq_len : {self.max_seq_len:,}",
|
||
f" hidden_dim : {self.hidden_dim}",
|
||
f" num_layers : {self.num_layers}",
|
||
f" num_heads : {self.num_heads}",
|
||
f" head_dim : {self.head_dim}",
|
||
f" ffn_hidden_dim : {self.ffn_hidden_dim}",
|
||
f" rope_base : {self.rope_base}",
|
||
f" norm_eps : {self.norm_eps}",
|
||
f" tie_weights : {self.tie_weights}",
|
||
f" bias : {self.bias}",
|
||
f" dropout : {self.dropout}",
|
||
f" attn_dropout : {self.attn_dropout}",
|
||
]
|
||
return "\n".join(lines)
|
||
|
||
|
||
def default_config() -> CyberSLMConfig:
|
||
"""
|
||
Return the validated default CyberSLM configuration.
|
||
|
||
This is the single source of truth for all training runs.
|
||
All hyperparameters match the finalized architecture specification.
|
||
|
||
Returns
|
||
-------
|
||
CyberSLMConfig
|
||
A validated, immutable configuration object.
|
||
"""
|
||
cfg = CyberSLMConfig()
|
||
cfg.validate()
|
||
return cfg
|