初始化项目，由ModelHub XC社区提供模型

Model: sabari2005/cyberslm-instruct
Source: Original Platform

cyberslm/model/__init__.py

@@ -0,0 +1,35 @@
+"""
+CyberSLM model package.
+
+Public API (grows as phases are added):
+    Phase 1: CyberSLMConfig, default_config, RMSNorm, RotaryPositionEmbedding, apply_rope
+    Phase 2: MultiHeadSelfAttention, SwiGLUFeedForward
+    Phase 3: DecoderBlock, CyberSLM, build_model
+    Phase 4: (training engine lives in cyberslm.training)
+"""
+
+from cyberslm.model.config import CyberSLMConfig, default_config
+from cyberslm.model.norm import RMSNorm
+from cyberslm.model.rope import RotaryPositionEmbedding, apply_rope
+from cyberslm.model.attention import MultiHeadSelfAttention
+from cyberslm.model.ffn import SwiGLUFeedForward
+from cyberslm.model.block import DecoderBlock
+from cyberslm.model.model import CyberSLM, build_model, count_parameters, model_summary
+
+__all__ = [
+    # Phase 1
+    "CyberSLMConfig",
+    "default_config",
+    "RMSNorm",
+    "RotaryPositionEmbedding",
+    "apply_rope",
+    # Phase 2
+    "MultiHeadSelfAttention",
+    "SwiGLUFeedForward",
+    # Phase 3
+    "DecoderBlock",
+    "CyberSLM",
+    "build_model",
+    "count_parameters",
+    "model_summary",
+]

cyberslm/model/attention.py

@@ -0,0 +1,302 @@
+"""
+Multi-Head Self Attention (MHSA)
+=================================
+Standard scaled dot-product multi-head self attention with:
+  - Rotary Position Embedding (RoPE) on queries and keys
+  - Causal (auto-regressive) masking
+  - No bias on projection layers
+  - Pre-norm placement handled by the enclosing DecoderBlock
+
+Mathematical definition
+-----------------------
+Given input X ∈ ℝ^{B×T×d}:
+
+  Q = X Wq,  K = X Wk,  V = X Wv       (projections, no bias)
+
+Split into H heads, each of dimension d_h = d / H:
+
+  Qₕ, Kₕ = RoPE(Qₕ), RoPE(Kₕ)         (apply rotary embeddings)
+
+Scaled dot-product attention per head:
+
+  Aₕ = softmax( (Qₕ Kₕᵀ) / √d_h + mask ) Vₕ
+
+where mask[i,j] = 0 if j ≤ i else −∞  (causal constraint).
+
+Concatenate and project:
+
+  output = concat(A₁, ..., A_H) Wo
+
+Complexity
+----------
+Time : O(T² · d)  — quadratic in sequence length (standard attention)
+Space: O(T² · H)  — attention weight matrix per head
+
+Numerical stability
+-------------------
+- Scaling by 1/√d_h keeps the pre-softmax logits in a well-conditioned
+  range, preventing vanishing gradients from very peaked softmax outputs.
+- Softmax is computed by PyTorch's numerically stable implementation
+  (subtract max before exp).
+- RoPE is applied in float32 (see rope.py).
+- Causal mask adds −∞ (not a large negative number) so masked positions
+  become exactly 0 after softmax — no gradient leakage.
+
+FlashAttention compatibility
+-----------------------------
+The forward pass is written in a way that is structurally compatible with
+a future drop-in replacement by ``torch.nn.functional.scaled_dot_product_attention``
+(PyTorch 2.0+) or the ``flash-attn`` library.  To migrate:
+  1. Replace the manual QKᵀ/softmax/V block with:
+         F.scaled_dot_product_attention(q, k, v, attn_mask=None,
+                                        dropout_p=0.0, is_causal=True)
+  2. Remove the manual mask addition (is_causal=True handles it).
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from cyberslm.model.config import CyberSLMConfig
+from cyberslm.model.rope import RotaryPositionEmbedding, apply_rope
+
+
+def _causal_bias(
+    q_len: int,
+    k_len: int,
+    past_len: int,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> Tensor:
+    """
+    Additive causal mask of shape ``(q_len, k_len)`` for a query block that
+    starts at absolute position ``past_len``.
+
+    Query row ``i`` represents absolute position ``past_len + i`` and may attend
+    to key columns ``0 .. past_len + i`` inclusive; everything after is -inf.
+    With ``past_len == 0`` this reduces to the usual upper-triangular mask.
+    """
+    q_pos = torch.arange(q_len, device=device).unsqueeze(1) + past_len  # (q,1)
+    k_pos = torch.arange(k_len, device=device).unsqueeze(0)             # (1,k)
+    return torch.where(
+        k_pos <= q_pos,
+        torch.zeros((), dtype=dtype, device=device),
+        torch.full((), float("-inf"), dtype=dtype, device=device),
+    )
+
+
+class MultiHeadSelfAttention(nn.Module):
+    """
+    Multi-Head Self Attention with RoPE and causal masking.
+
+    This module owns the four projection matrices (Wq, Wk, Wv, Wo),
+    the RoPE cache, and the causal mask buffer.
+
+    Parameters
+    ----------
+    config : CyberSLMConfig
+        Validated model configuration.
+
+    Attributes
+    ----------
+    q_proj : nn.Linear   ``(hidden_dim, hidden_dim)``, no bias
+    k_proj : nn.Linear   ``(hidden_dim, hidden_dim)``, no bias
+    v_proj : nn.Linear   ``(hidden_dim, hidden_dim)``, no bias
+    o_proj : nn.Linear   ``(hidden_dim, hidden_dim)``, no bias
+    rope   : RotaryPositionEmbedding
+
+    Shape
+    -----
+    Input  : ``(batch, seq_len, hidden_dim)``
+    Output : ``(batch, seq_len, hidden_dim)``
+    """
+
+    def __init__(
+        self,
+        config: CyberSLMConfig,
+        rope: Optional[RotaryPositionEmbedding] = None,
+    ) -> None:
+        super().__init__()
+        self.hidden_dim = config.hidden_dim
+        self.num_heads  = config.num_heads
+        self.head_dim   = config.head_dim
+        self.scale      = 1.0 / math.sqrt(self.head_dim)
+        self.attn_dropout_p = config.attn_dropout
+
+        # ------------------------------------------------------------------ #
+        # Projection layers — no bias (modern practice, saves ~4×384 params) #
+        # ------------------------------------------------------------------ #
+        self.q_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
+        self.o_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
+
+        # ------------------------------------------------------------------ #
+        # RoPE cache                                                          #
+        # ------------------------------------------------------------------ #
+        # The cos/sin tables depend only on (head_dim, max_seq_len, base), so
+        # every layer's would be byte-identical. CyberSLM builds ONE and passes
+        # it in; previously each of the 12 layers constructed its own, costing
+        # ~12 MB of duplicated buffers. Falls back to building its own so the
+        # module stays usable standalone (tests, ablations).
+        self.rope = rope if rope is not None else RotaryPositionEmbedding(
+            head_dim=config.head_dim,
+            max_seq_len=config.max_seq_len,
+            base=config.rope_base,
+        )
+
+        # Causality is enforced by scaled_dot_product_attention(is_causal=...)
+        # rather than a materialised (max_seq_len × max_seq_len) mask buffer,
+        # which previously cost ~67 MB per layer.
+
+    def forward(
+        self,
+        x: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        return_attn_weights: bool = False,
+        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor, Tensor]]]:
+        """
+        Compute multi-head self attention.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input of shape ``(batch, seq_len, hidden_dim)``.
+        attention_mask : Optional[Tensor]
+            Key-padding mask of shape ``(batch, seq_len)`` with 1 for real
+            tokens and 0 for padding. When provided, padded keys are excluded
+            from every query's attention (in addition to the causal mask).
+            ``None`` means no padding (the common training case with packed
+            sequences).
+        return_attn_weights : bool
+            If True, also return the attention weight matrix for inspection.
+            This forces the slower explicit-softmax path; leave False for
+            training so the fused kernel is used.
+
+        Returns
+        -------
+        output : Tensor
+            Shape ``(batch, seq_len, hidden_dim)``.
+        attn_weights : Optional[Tensor]
+            Shape ``(batch, num_heads, seq_len, seq_len)`` if
+            ``return_attn_weights=True``, else ``None``.
+        present : Optional[Tuple[Tensor, Tensor]]
+            The concatenated ``(k, v)`` for this layer when ``use_cache=True``,
+            to be fed back on the next decoding step. ``None`` otherwise.
+        """
+        B, T, _ = x.shape
+        # Number of tokens already in the cache == absolute position of x[0].
+        past_len = kv_cache[0].size(2) if kv_cache is not None else 0
+
+        # ------------------------------------------------------------------ #
+        # 1. Linear projections                                               #
+        # ------------------------------------------------------------------ #
+        q = self.q_proj(x)  # (B, T, hidden_dim)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+
+        # ------------------------------------------------------------------ #
+        # 2. Reshape to (B, H, T, head_dim) for multi-head computation       #
+        # ------------------------------------------------------------------ #
+        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)  # (B,H,T,D)
+        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # ------------------------------------------------------------------ #
+        # 3. Apply Rotary Position Embeddings to Q and K                     #
+        # ------------------------------------------------------------------ #
+        # offset=past_len so a cached decode step rotates the new token by its
+        # TRUE absolute position rather than position 0.
+        q, k = apply_rope(q, k, self.rope, offset=past_len)
+
+        # ------------------------------------------------------------------ #
+        # 3b. Prepend the cache. RoPE is applied to the new k BEFORE the
+        #     concat, and cached keys were already rotated when they were first
+        #     computed -- so each key keeps the rotation for its own position.
+        # ------------------------------------------------------------------ #
+        if kv_cache is not None:
+            k = torch.cat([kv_cache[0], k], dim=2)
+            v = torch.cat([kv_cache[1], v], dim=2)
+        present = (k, v) if use_cache else None
+        S = k.size(2)   # total key length (past + current)
+
+        # ------------------------------------------------------------------ #
+        # 4. Build the additive key-padding bias (if any).                    #
+        #    Shape broadcasts over heads and query positions: (B, 1, 1, T).   #
+        # ------------------------------------------------------------------ #
+        pad_bias: Optional[Tensor] = None
+        if attention_mask is not None:
+            # 0 where padding → -inf added to those key columns.
+            pad = (attention_mask == 0)[:, None, None, :]  # (B,1,1,S) bool
+            pad_bias = torch.zeros(
+                (B, 1, 1, pad.size(-1)), dtype=q.dtype, device=q.device
+            ).masked_fill(pad, float("-inf"))
+
+        if not return_attn_weights:
+            # Fused, memory-efficient path (FlashAttention when available).
+            # is_causal=True applies the causal mask without materialising it.
+            if pad_bias is None and past_len == 0:
+                context = F.scaled_dot_product_attention(
+                    q, k, v,
+                    is_causal=True,
+                    dropout_p=self.attn_dropout_p if self.training else 0.0,
+                )
+            elif pad_bias is None and T == 1:
+                # Single-token decode: every cached key is in the past, so the
+                # causal constraint is already satisfied and no mask is needed.
+                context = F.scaled_dot_product_attention(
+                    q, k, v, dropout_p=0.0,
+                )
+            else:
+                # Combine causal + padding into one additive float mask.
+                # Query i sits at absolute position past_len + i and may attend
+                # to keys 0..past_len+i, so the triangle is offset by past_len.
+                causal = _causal_bias(T, S, past_len, q.dtype, q.device)
+                attn_bias = causal[None, None, :, :]
+                if pad_bias is not None:
+                    attn_bias = attn_bias + pad_bias   # (B,1,T,S)
+                context = F.scaled_dot_product_attention(
+                    q, k, v,
+                    attn_mask=attn_bias,
+                    dropout_p=self.attn_dropout_p if self.training else 0.0,
+                )
+            attn_weights = None
+        else:
+            # Explicit path — needed only when the caller wants the weights.
+            scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale  # (B,H,T,S)
+            causal = _causal_bias(T, S, past_len, scores.dtype, scores.device)
+            scores = scores + causal[None, None, :, :]
+            if pad_bias is not None:
+                scores = scores + pad_bias
+            attn_weights = F.softmax(scores, dim=-1, dtype=torch.float32)
+            if self.attn_dropout_p > 0.0 and self.training:
+                attn_weights = F.dropout(attn_weights, p=self.attn_dropout_p)
+            attn_weights = attn_weights.to(v.dtype)
+            context = torch.matmul(attn_weights, v)  # (B,H,T,head_dim)
+
+        # ------------------------------------------------------------------ #
+        # 8. Merge heads: (B, H, T, D) → (B, T, H*D) = (B, T, hidden_dim)  #
+        # ------------------------------------------------------------------ #
+        context = context.transpose(1, 2).contiguous().view(B, T, self.hidden_dim)
+
+        # ------------------------------------------------------------------ #
+        # 9. Output projection                                                #
+        # ------------------------------------------------------------------ #
+        output = self.o_proj(context)
+
+        return output, attn_weights, present
+
+    def extra_repr(self) -> str:
+        return (
+            f"hidden_dim={self.hidden_dim}, "
+            f"num_heads={self.num_heads}, "
+            f"head_dim={self.head_dim}"
+        )

cyberslm/model/block.py

@@ -0,0 +1,144 @@
+"""
+Transformer Decoder Block
+==========================
+A single Pre-Norm residual decoder block consisting of:
+  1. RMSNorm → Multi-Head Self Attention → residual add
+  2. RMSNorm → SwiGLU FFN               → residual add
+
+Pre-Norm architecture
+---------------------
+Post-Norm (original Transformer):
+    x = LayerNorm(x + SubLayer(x))
+
+Pre-Norm (modern: GPT-2 onward, LLaMA, etc.):
+    x = x + SubLayer(LayerNorm(x))
+
+Pre-Norm is strongly preferred for deep networks because:
+  - Gradients flow through the residual connection bypassing the
+    normalisation, preventing vanishing gradients in very deep stacks.
+  - Training is more stable without learning-rate warmup tricks.
+  - Final RMSNorm on the output is added at the model level (not here)
+    to normalise the final residual stream before the LM head.
+
+Residual stream
+---------------
+The residual stream x ∈ ℝ^{B×T×d} is the backbone of the model. Each
+sub-layer reads from it, computes a delta, and adds back:
+
+    Δ_attn = MHSA( RMSNorm(x) )
+    x       = x + Δ_attn
+
+    Δ_ffn  = FFN( RMSNorm(x) )
+    x       = x + Δ_ffn
+
+This additive structure means the gradient of the loss with respect to
+early layers contains a direct path through the identity (residual),
+enabling reliable training of 12+ layer models.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from cyberslm.model.config import CyberSLMConfig
+from cyberslm.model.norm import RMSNorm
+from cyberslm.model.attention import MultiHeadSelfAttention
+from cyberslm.model.ffn import SwiGLUFeedForward
+
+
+class DecoderBlock(nn.Module):
+    """
+    Pre-Norm Transformer Decoder Block.
+
+    Parameters
+    ----------
+    config : CyberSLMConfig
+        Validated model configuration.
+    layer_idx : int
+        Zero-based index of this block in the stack (used for display only).
+
+    Sub-modules
+    -----------
+    attn_norm : RMSNorm
+        Normalises the residual stream before attention.
+    attn : MultiHeadSelfAttention
+        Self attention with RoPE and causal masking.
+    ffn_norm : RMSNorm
+        Normalises the residual stream before the FFN.
+    ffn : SwiGLUFeedForward
+        SwiGLU position-wise feed-forward network.
+
+    Shape
+    -----
+    Input  : ``(batch, seq_len, hidden_dim)``
+    Output : ``(batch, seq_len, hidden_dim)``
+    """
+
+    def __init__(
+        self,
+        config: CyberSLMConfig,
+        layer_idx: int = 0,
+        rope=None,
+    ) -> None:
+        super().__init__()
+        self.layer_idx = layer_idx
+
+        # Pre-norm before attention.
+        self.attn_norm = RMSNorm(config.hidden_dim, eps=config.norm_eps)
+        # Multi-head self attention (owns RoPE + causal mask buffers).
+        self.attn = MultiHeadSelfAttention(config, rope=rope)
+
+        # Pre-norm before FFN.
+        self.ffn_norm = RMSNorm(config.hidden_dim, eps=config.norm_eps)
+        # SwiGLU feed-forward.
+        self.ffn = SwiGLUFeedForward(config)
+
+    def forward(
+        self,
+        x: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        return_attn_weights: bool = False,
+        kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
+        use_cache: bool = False,
+    ) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor, Tensor]]]:
+        """
+        Apply one Pre-Norm residual decoder block.
+
+        Parameters
+        ----------
+        x : Tensor
+            Residual stream of shape ``(batch, seq_len, hidden_dim)``.
+        attention_mask : Optional[Tensor]
+            Key-padding mask ``(batch, seq_len)`` (1=keep, 0=pad), forwarded
+            to the attention sub-layer. ``None`` for packed/unpadded batches.
+        return_attn_weights : bool
+            Propagated to the attention sub-layer.
+
+        Returns
+        -------
+        x : Tensor
+            Updated residual stream, same shape as input.
+        attn_weights : Optional[Tensor]
+            Attention weights if requested, else None.
+        """
+        # ---- Attention sub-layer ----------------------------------------- #
+        attn_out, attn_weights, present = self.attn(
+            self.attn_norm(x),
+            attention_mask=attention_mask,
+            return_attn_weights=return_attn_weights,
+            kv_cache=kv_cache,
+            use_cache=use_cache,
+        )
+        x = x + attn_out
+
+        # ---- FFN sub-layer ------------------------------------------------ #
+        x = x + self.ffn(self.ffn_norm(x))
+
+        return x, attn_weights, present
+
+    def extra_repr(self) -> str:
+        return f"layer_idx={self.layer_idx}"

cyberslm/model/config.py

@@ -0,0 +1,234 @@
+"""
+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

cyberslm/model/ffn.py

@@ -0,0 +1,127 @@
+"""
+SwiGLU Feed-Forward Network (FFN)
+===================================
+Reference: "GLU Variants Improve Transformer" (Noam Shazeer, 2020)
+           https://arxiv.org/abs/2002.05202
+
+Mathematical definition
+-----------------------
+Standard FFN (for contrast):
+    FFN(x) = activation(x W₁) W₂
+
+SwiGLU FFN:
+    SwiGLU(x) = (x W_gate ⊙ swish(x W_gate)) W₂   ← WRONG shorthand
+    
+Correct form with separate gate and value projections:
+    gate(x) = x W_gate        ∈ ℝ^{B×T×ffn_dim}
+    val(x)  = x W_val         ∈ ℝ^{B×T×ffn_dim}
+    hidden  = swish(gate(x)) ⊙ val(x)
+    out     = hidden W_out    ∈ ℝ^{B×T×hidden_dim}
+
+where swish(z) = z · sigmoid(z) = z · σ(z).
+
+Why SwiGLU
+----------
+- The gating mechanism (⊙) gives the network a multiplicative path to
+  control information flow — intuitively "how much" of each feature passes
+  through at each position.
+- swish is smooth and non-monotonic, empirically outperforming ReLU and
+  GELU in large-scale experiments (PaLM, LLaMA, etc.).
+- The factored W_gate / W_val structure adds one extra matrix but improves
+  quality relative to a standard 2-layer FFN of the same parameter budget.
+
+Parameter count
+---------------
+Three matrices: W_gate, W_val, W_out
+    W_gate : hidden_dim × ffn_hidden_dim   (384 × 1024 = 393 216)
+    W_val  : hidden_dim × ffn_hidden_dim   (384 × 1024 = 393 216)
+    W_out  : ffn_hidden_dim × hidden_dim   (1024 × 384 = 393 216)
+    Total per layer: 1 179 648 ≈ 1.18 M
+
+No bias on any projection (consistent with modern practice).
+
+Numerical stability
+-------------------
+- swish(z) = z · σ(z) is numerically safe for all real z.
+- σ(z) = 1/(1+exp(−z)) in PyTorch uses a numerically stable implementation.
+- The element-wise product of two bounded (after sigmoid) quantities does
+  not amplify values explosively.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+from cyberslm.model.config import CyberSLMConfig
+
+
+class SwiGLUFeedForward(nn.Module):
+    """
+    SwiGLU Feed-Forward Network.
+
+    Consists of three bias-free linear projections:
+    ``gate_proj``, ``val_proj``, and ``out_proj``.
+
+    Parameters
+    ----------
+    config : CyberSLMConfig
+        Validated model configuration.
+
+    Shape
+    -----
+    Input  : ``(batch, seq_len, hidden_dim)``
+    Output : ``(batch, seq_len, hidden_dim)``
+    """
+
+    def __init__(self, config: CyberSLMConfig) -> None:
+        super().__init__()
+        self.hidden_dim     = config.hidden_dim
+        self.ffn_hidden_dim = config.ffn_hidden_dim
+
+        # Gate projection: produces the gating signal fed through swish.
+        self.gate_proj = nn.Linear(
+            config.hidden_dim, config.ffn_hidden_dim, bias=False
+        )
+        # Value projection: produces the value signal gated element-wise.
+        self.val_proj = nn.Linear(
+            config.hidden_dim, config.ffn_hidden_dim, bias=False
+        )
+        # Output projection: maps back to residual stream dimension.
+        self.out_proj = nn.Linear(
+            config.ffn_hidden_dim, config.hidden_dim, bias=False
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Apply SwiGLU feed-forward transformation.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input of shape ``(batch, seq_len, hidden_dim)``.
+
+        Returns
+        -------
+        Tensor
+            Output of shape ``(batch, seq_len, hidden_dim)``.
+        """
+        # gate:   (B, T, ffn_hidden_dim)
+        # val:    (B, T, ffn_hidden_dim)
+        gate = self.gate_proj(x)
+        val  = self.val_proj(x)
+
+        # SwiGLU: swish(gate) ⊙ val
+        # F.silu is swish: silu(z) = z * sigmoid(z)
+        hidden = F.silu(gate) * val  # (B, T, ffn_hidden_dim)
+
+        # Project back to hidden_dim
+        return self.out_proj(hidden)  # (B, T, hidden_dim)
+
+    def extra_repr(self) -> str:
+        return (
+            f"hidden_dim={self.hidden_dim}, "
+            f"ffn_hidden_dim={self.ffn_hidden_dim}"
+        )

cyberslm/model/model.py

@@ -0,0 +1,514 @@
+"""
+CyberSLM — Complete Decoder-Only Language Model
+================================================
+Assembles all components into the full model:
+
+  Token Embedding  →  Decoder Stack (×12)  →  Final RMSNorm  →  LM Head
+
+Weight tying
+------------
+The LM head (output projection that maps hidden_dim → vocab_size) shares its
+weight matrix with the token embedding (vocab_size × hidden_dim).
+
+Mathematical justification: both the embedding matrix E and the unembedding
+matrix U operate in the same semantic space. Tying U = Eᵀ forces consistency
+("a token's output representation should be similar to its input representation"),
+reduces parameters by vocab_size × hidden_dim = 32 000 × 384 ≈ 12.3 M, and
+empirically improves perplexity on small models.
+
+Parameter count breakdown
+--------------------------
+Component                                Params
+---------------------------------------- ------
+Token embedding (vocab × hidden)          12 288 000
+  ↳ shared with LM head (no extra cost)          0
+Decoder blocks × 12:
+  attn_norm (RMSNorm)          384 ×12 =    4 608
+  attn Q/K/V/O proj       589 824 ×12 =  7 077 888
+  ffn_norm (RMSNorm)           384 ×12 =    4 608
+  ffn gate/val/out        1 179 648 ×12 = 14 155 776
+Final RMSNorm                               384
+---------------------------------------- ------
+Total                                  ≈ 33 531 264  (≈33.53 M)
+
+Note: RoPE buffers and causal mask buffers are NOT parameters.
+
+Initialisation
+--------------
+- Embeddings: N(0, 0.02)  — small but non-zero, standard practice.
+- All linear weights: N(0, 0.02)
+- All RMSNorm γ: 1.0  (already set by RMSNorm.__init__)
+- Output projection weights = Embedding weights (weight tying).
+- Scaled output projections: attention o_proj and FFN out_proj are
+  scaled by 1/√(2·num_layers) to prevent residual stream variance
+  from growing with depth (following GPT-2 / LLaMA init practice).
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from cyberslm.model.config import CyberSLMConfig, default_config
+from cyberslm.model.block import DecoderBlock
+from cyberslm.model.norm import RMSNorm
+from cyberslm.model.rope import RotaryPositionEmbedding
+
+
+class CyberSLM(nn.Module):
+    """
+    CyberSLM Decoder-Only Transformer.
+
+    Parameters
+    ----------
+    config : CyberSLMConfig
+        Validated model configuration.
+
+    Attributes
+    ----------
+    config : CyberSLMConfig
+    embedding : nn.Embedding
+        Token embedding table, shape ``(vocab_size, hidden_dim)``.
+    layers : nn.ModuleList[DecoderBlock]
+        Stack of ``num_layers`` decoder blocks.
+    final_norm : RMSNorm
+        Applied to the residual stream after the last block.
+    lm_head : nn.Linear
+        Projects hidden_dim → vocab_size.  Weight tied to ``embedding``.
+    """
+
+    def __init__(self, config: CyberSLMConfig) -> None:
+        super().__init__()
+        config.validate()
+        self.config = config
+
+        # ------------------------------------------------------------------ #
+        # Token embedding                                                      #
+        # ------------------------------------------------------------------ #
+        self.embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
+
+        # ------------------------------------------------------------------ #
+        # Decoder stack                                                        #
+        # ------------------------------------------------------------------ #
+        # One RoPE table shared by every layer (identical by construction).
+        self.rope = RotaryPositionEmbedding(
+            head_dim=config.head_dim,
+            max_seq_len=config.max_seq_len,
+            base=config.rope_base,
+        )
+        self.layers = nn.ModuleList(
+            [
+                DecoderBlock(config, layer_idx=i, rope=self.rope)
+                for i in range(config.num_layers)
+            ]
+        )
+
+        # ------------------------------------------------------------------ #
+        # Final normalisation                                                  #
+        # ------------------------------------------------------------------ #
+        self.final_norm = RMSNorm(config.hidden_dim, eps=config.norm_eps)
+
+        # ------------------------------------------------------------------ #
+        # Language model head                                                  #
+        # ------------------------------------------------------------------ #
+        # bias=False: unembedding never needs a bias term.
+        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+
+        # ------------------------------------------------------------------ #
+        # Weight tying: lm_head.weight ≡ embedding.weight                    #
+        # ------------------------------------------------------------------ #
+        if config.tie_weights:
+            self.lm_head.weight = self.embedding.weight
+
+        # ------------------------------------------------------------------ #
+        # Parameter initialisation                                             #
+        # ------------------------------------------------------------------ #
+        self._init_weights()
+
+    # ---------------------------------------------------------------------- #
+    # Initialisation                                                           #
+    # ---------------------------------------------------------------------- #
+
+    def _init_weights(self) -> None:
+        """
+        Initialise all parameters with production-quality values.
+
+        Strategy
+        --------
+        - Embedding          : N(0, 0.02)
+        - All Linear weights : N(0, 0.02)
+        - o_proj and out_proj: scaled down by 1/√(2·L) where L = num_layers
+          to stabilise the residual stream variance at initialisation.
+        - All RMSNorm γ      : 1.0  (already set in RMSNorm.__init__)
+        - All biases         : 0.0  (none exist in this config)
+        """
+        std = 0.02
+        scaled_std = std / math.sqrt(2.0 * self.config.num_layers)
+
+        for name, module in self.named_modules():
+            if isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=std)
+            elif isinstance(module, nn.Linear):
+                # Scaled init for residual output projections.
+                if name.endswith("o_proj") or name.endswith("out_proj"):
+                    nn.init.normal_(module.weight, mean=0.0, std=scaled_std)
+                else:
+                    nn.init.normal_(module.weight, mean=0.0, std=std)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+        # Weight tying must be re-applied after init because _init_weights
+        # initialised embedding.weight; lm_head.weight already points to the
+        # same tensor (Python object reference), so no extra step needed.
+        # Verify it is still tied.
+        if self.config.tie_weights:
+            assert self.lm_head.weight is self.embedding.weight, (
+                "Weight tying broken after _init_weights"
+            )
+
+    # ---------------------------------------------------------------------- #
+    # Forward pass                                                             #
+    # ---------------------------------------------------------------------- #
+
+    def forward(
+        self,
+        input_ids: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        return_all_attn_weights: bool = False,
+    ) -> Tuple[Tensor, List[Optional[Tensor]]]:
+        """
+        Run the full forward pass.
+
+        Parameters
+        ----------
+        input_ids : Tensor
+            Long tensor of shape ``(batch, seq_len)`` with token IDs in
+            ``[0, vocab_size)``.
+        attention_mask : Optional[Tensor]
+            Optional key-padding mask ``(batch, seq_len)`` (1=keep, 0=pad).
+            Pass this when batches contain right-padded sequences so padded
+            positions do not corrupt real tokens; leave ``None`` for packed,
+            unpadded training batches.
+        return_all_attn_weights : bool
+            If True, collect and return attention weights from every layer.
+            Disabled by default for training efficiency.
+
+        Returns
+        -------
+        logits : Tensor
+            Shape ``(batch, seq_len, vocab_size)``.  Raw (pre-softmax) scores.
+        all_attn_weights : List[Optional[Tensor]]
+            One entry per decoder block; each is either the attention weight
+            tensor or ``None``.
+
+        Notes
+        -----
+        For language model training the standard loss is:
+            loss = cross_entropy(logits[:, :-1].reshape(-1, V),
+                                 input_ids[:, 1:].reshape(-1))
+        where we predict the next token at every position.
+        """
+        # ------------------------------------------------------------------ #
+        # 1. Token embedding                                                   #
+        # ------------------------------------------------------------------ #
+        x = self.embedding(input_ids)  # (B, T, hidden_dim)
+
+        # ------------------------------------------------------------------ #
+        # 2. Decoder stack                                                     #
+        # ------------------------------------------------------------------ #
+        all_attn_weights: List[Optional[Tensor]] = []
+        for block in self.layers:
+            # use_cache=False -> `present` is None; forward() deliberately keeps
+            # its (logits, attn_weights) return signature unchanged. Cached
+            # decoding lives in generate() instead of overloading this method.
+            x, attn_w, _ = block(
+                x,
+                attention_mask=attention_mask,
+                return_attn_weights=return_all_attn_weights,
+            )
+            all_attn_weights.append(attn_w)
+
+        # ------------------------------------------------------------------ #
+        # 3. Final normalisation                                               #
+        # ------------------------------------------------------------------ #
+        x = self.final_norm(x)  # (B, T, hidden_dim)
+
+        # ------------------------------------------------------------------ #
+        # 4. LM head (weight-tied unembedding)                                #
+        # ------------------------------------------------------------------ #
+        logits = self.lm_head(x)  # (B, T, vocab_size)
+
+        return logits, all_attn_weights
+
+
+    # ---------------------------------------------------------------------- #
+    # Cached autoregressive generation                                         #
+    # ---------------------------------------------------------------------- #
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: Tensor,
+        max_new_tokens: int = 256,
+        temperature: float = 0.0,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        repetition_penalty: float = 1.0,
+        eos_id: Optional[int] = None,
+    ) -> Tensor:
+        """
+        Generate continuations using a KV cache.
+
+        Why this exists
+        ---------------
+        The previous generation loop re-ran the whole 12-layer stack over the
+        entire prefix for every single token, making decoding O(n^2) in
+        sequence length. With a cache each step attends over the cached keys and
+        only computes the new token, which is O(n) overall.
+
+        Sampling is applied per row, so batched prompts are supported. Rows that
+        have emitted ``eos_id`` are frozen (further tokens are forced to
+        ``eos_id``) and generation stops once every row is finished.
+
+        Parameters
+        ----------
+        input_ids : Tensor  ``(batch, prompt_len)`` of token ids.
+        temperature : 0.0 selects greedy argmax; >0 samples.
+        top_k / top_p : 0 and 1.0 respectively disable the filter.
+        repetition_penalty : >1.0 divides logits of already-present tokens.
+        eos_id : stop token; ``None`` means never stop early.
+
+        Returns
+        -------
+        Tensor ``(batch, prompt_len + generated)`` including the prompt.
+        """
+        self.eval()
+        device = input_ids.device
+        B = input_ids.size(0)
+        max_ctx = self.config.max_seq_len
+
+        if input_ids.size(1) >= max_ctx:
+            input_ids = input_ids[:, -(max_ctx - 1):]
+
+        caches: List[Optional[tuple]] = [None] * len(self.layers)
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+        out = input_ids
+        cur = input_ids
+
+        for _ in range(max_new_tokens):
+            if out.size(1) >= max_ctx:
+                break
+
+            h = self.embedding(cur)
+            new_caches = []
+            for block, layer_cache in zip(self.layers, caches):
+                h, _, present = block(h, kv_cache=layer_cache, use_cache=True)
+                new_caches.append(present)
+            caches = new_caches
+            logits = self.lm_head(self.final_norm(h))[:, -1, :].float()
+
+            if repetition_penalty != 1.0:
+                for b in range(B):
+                    seen = torch.unique(out[b])
+                    lg = logits[b, seen]
+                    # Divide positives, multiply negatives, so the penalty always
+                    # pushes a token DOWN regardless of its logit's sign.
+                    logits[b, seen] = torch.where(
+                        lg > 0, lg / repetition_penalty, lg * repetition_penalty
+                    )
+
+            if temperature == 0.0:
+                nxt = logits.argmax(dim=-1)
+            else:
+                logits = logits / temperature
+                if top_k > 0:
+                    k = min(top_k, logits.size(-1))
+                    thresh = torch.topk(logits, k, dim=-1).values[:, -1, None]
+                    logits = logits.masked_fill(logits < thresh, float("-inf"))
+                if top_p < 1.0:
+                    srt, idx = torch.sort(logits, descending=True, dim=-1)
+                    probs = torch.softmax(srt, dim=-1)
+                    cum = probs.cumsum(dim=-1) - probs   # prob mass strictly before this token
+                    srt = srt.masked_fill(cum > top_p, float("-inf"))
+                    logits = torch.full_like(logits, float("-inf")).scatter(1, idx, srt)
+                nxt = torch.multinomial(torch.softmax(logits, dim=-1), 1).squeeze(-1)
+
+            if eos_id is not None:
+                nxt = torch.where(finished, torch.full_like(nxt, eos_id), nxt)
+                finished = finished | (nxt == eos_id)
+
+            cur = nxt.unsqueeze(1)
+            out = torch.cat([out, cur], dim=1)
+
+            if eos_id is not None and bool(finished.all()):
+                break
+
+        return out
+
+    # ---------------------------------------------------------------------- #
+    # Convenience: next-token logits                                           #
+    # ---------------------------------------------------------------------- #
+
+    def get_next_token_logits(self, input_ids: Tensor) -> Tensor:
+        """
+        Return logits for the next token after the last input position.
+
+        Parameters
+        ----------
+        input_ids : Tensor
+            Shape ``(batch, seq_len)``.
+
+        Returns
+        -------
+        Tensor
+            Shape ``(batch, vocab_size)``.
+        """
+        logits, _ = self.forward(input_ids)
+        return logits[:, -1, :]  # (B, vocab_size)
+
+
+# --------------------------------------------------------------------------- #
+# Parameter counting                                                            #
+# --------------------------------------------------------------------------- #
+
+def count_parameters(model: nn.Module) -> Dict[str, int]:
+    """
+    Count trainable and total parameters.
+
+    Because of weight tying, lm_head.weight is counted only once
+    (it shares storage with embedding.weight).
+
+    Parameters
+    ----------
+    model : nn.Module
+
+    Returns
+    -------
+    dict with keys:
+        ``total``     — total parameter elements (no double-counting)
+        ``trainable`` — trainable parameter elements
+    """
+    seen: set = set()
+    total = 0
+    trainable = 0
+    for param in model.parameters():
+        # data_ptr() is unique per underlying storage tensor.
+        if param.data_ptr() in seen:
+            continue
+        seen.add(param.data_ptr())
+        n = param.numel()
+        total += n
+        if param.requires_grad:
+            trainable += n
+    return {"total": total, "trainable": trainable}
+
+
+# --------------------------------------------------------------------------- #
+# Model summary                                                                 #
+# --------------------------------------------------------------------------- #
+
+def model_summary(model: CyberSLM) -> str:
+    """
+    Return a human-readable model summary string.
+
+    Parameters
+    ----------
+    model : CyberSLM
+
+    Returns
+    -------
+    str
+        Formatted multi-line summary including per-component parameter counts.
+    """
+    cfg = model.config
+    param_info = count_parameters(model)
+
+    lines = [
+        "=" * 60,
+        f"  CyberSLM Model Summary",
+        "=" * 60,
+        f"  Architecture   : Decoder-only Transformer",
+        f"  Hidden dim     : {cfg.hidden_dim}",
+        f"  Num layers     : {cfg.num_layers}",
+        f"  Num heads      : {cfg.num_heads}",
+        f"  Head dim       : {cfg.head_dim}",
+        f"  FFN hidden dim : {cfg.ffn_hidden_dim}",
+        f"  Vocab size     : {cfg.vocab_size:,}",
+        f"  Max seq len    : {cfg.max_seq_len:,}",
+        f"  Weight tied    : {cfg.tie_weights}",
+        f"  RoPE base      : {cfg.rope_base}",
+        "-" * 60,
+        f"  Total params   : {param_info['total']:>14,}",
+        f"  Trainable      : {param_info['trainable']:>14,}",
+        "-" * 60,
+        "  Per-component:",
+    ]
+
+    # Embedding
+    emb_p = model.embedding.weight.numel()
+    lines.append(f"    Embedding              : {emb_p:>12,}")
+
+    # Per-block breakdown (just the first block, all identical)
+    block = model.layers[0]
+    attn_norm_p = sum(p.numel() for p in block.attn_norm.parameters())
+    attn_p      = sum(p.numel() for p in block.attn.parameters())
+    ffn_norm_p  = sum(p.numel() for p in block.ffn_norm.parameters())
+    ffn_p       = sum(p.numel() for p in block.ffn.parameters())
+    block_total = attn_norm_p + attn_p + ffn_norm_p + ffn_p
+    lines.append(f"    Decoder block (×{cfg.num_layers:2d})     : {block_total:>12,}  per block")
+    lines.append(f"      attn_norm            : {attn_norm_p:>12,}")
+    lines.append(f"      attention (Q/K/V/O)  : {attn_p:>12,}")
+    lines.append(f"      ffn_norm             : {ffn_norm_p:>12,}")
+    lines.append(f"      ffn (gate/val/out)   : {ffn_p:>12,}")
+    lines.append(f"    Decoder stack total    : {block_total * cfg.num_layers:>12,}")
+
+    # Final norm
+    final_norm_p = sum(p.numel() for p in model.final_norm.parameters())
+    lines.append(f"    Final RMSNorm          : {final_norm_p:>12,}")
+
+    # LM head — note: weight tied, so 0 additional params
+    lm_tied_note = " (weight-tied, no extra params)" if cfg.tie_weights else ""
+    lines.append(f"    LM head                : {0:>12,}{lm_tied_note}")
+
+    lines.append("=" * 60)
+    return "\n".join(lines)
+
+
+# --------------------------------------------------------------------------- #
+# Model builder                                                                 #
+# --------------------------------------------------------------------------- #
+
+def build_model(
+    config: Optional[CyberSLMConfig] = None,
+    device: Optional[torch.device] = None,
+) -> CyberSLM:
+    """
+    Build, initialise, and optionally place the CyberSLM model.
+
+    Parameters
+    ----------
+    config : CyberSLMConfig, optional
+        Validated config.  Uses :func:`default_config` if None.
+    device : torch.device, optional
+        Target device.  Stays on CPU if None.
+
+    Returns
+    -------
+    CyberSLM
+        Fully initialised model ready for training.
+    """
+    if config is None:
+        config = default_config()
+    else:
+        config.validate()
+
+    model = CyberSLM(config)
+
+    if device is not None:
+        model = model.to(device)
+
+    return model

cyberslm/model/norm.py

@@ -0,0 +1,114 @@
+"""
+RMSNorm — Root Mean Square Layer Normalisation
+===============================================
+Reference: "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019)
+           https://arxiv.org/abs/1910.07467
+
+Mathematical definition
+-----------------------
+Given an input vector **x** ∈ ℝ^d:
+
+    RMS(x) = sqrt( (1/d) * Σ xᵢ² + ε )
+
+    RMSNorm(x) = (x / RMS(x)) * γ
+
+where γ ∈ ℝ^d is a learned per-channel scale (initialised to 1.0) and
+ε > 0 is a small constant for numerical stability.
+
+Key differences from LayerNorm
+-------------------------------
+- No mean subtraction (no re-centering step).
+- No learned bias β (the bias-free variant).
+- ~30 % fewer operations than LayerNorm, which matters across 12 layers.
+- Empirically matches or exceeds LayerNorm in modern transformer training.
+
+Numerical stability
+-------------------
+- The RMS is computed in float32 regardless of input dtype. This prevents
+  underflow/overflow when activations are in bfloat16 or float16.  For our
+  FP32 training runs this cast is a no-op but it is correct and future-proof.
+- ε = 1e-6 (default) prevents division by zero even for near-zero inputs.
+- The scale γ is cast back to the input dtype before multiplication.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class RMSNorm(nn.Module):
+    """
+    Root Mean Square Layer Normalisation without mean-centering or bias.
+
+    Parameters
+    ----------
+    dim : int
+        Feature dimension to normalise over (last dimension of the input).
+    eps : float
+        Small constant added to the RMS denominator for numerical stability.
+        Defaults to 1e-6.
+
+    Shape
+    -----
+    Input  : ``(*, dim)``  — any leading batch / sequence dimensions.
+    Output : ``(*, dim)``  — same shape as input.
+
+    Examples
+    --------
+    >>> norm = RMSNorm(384, eps=1e-6)
+    >>> x = torch.randn(2, 512, 384)
+    >>> y = norm(x)
+    >>> y.shape
+    torch.Size([2, 512, 384])
+    """
+
+    def __init__(self, dim: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        if dim <= 0:
+            raise ValueError(f"dim must be positive, got {dim}")
+        if eps <= 0.0:
+            raise ValueError(f"eps must be positive, got {eps}")
+
+        self.dim = dim
+        self.eps = eps
+        # Learned per-channel scale, initialised to 1 (identity transform).
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _compute_rms(self, x: Tensor) -> Tensor:
+        """
+        Compute the RMS over the last dimension.
+
+        Always promotes to float32 to prevent numerical issues with
+        reduced-precision dtypes. For FP32 training this is a no-op.
+
+        Returns
+        -------
+        Tensor
+            Shape ``(*, 1)`` — one RMS value per token position.
+        """
+        return x.float().pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        Normalise ``x`` by its per-token root mean square.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input of shape ``(*, dim)``.
+
+        Returns
+        -------
+        Tensor
+            Normalised output of the same shape and dtype as ``x``.
+        """
+        rms = self._compute_rms(x)
+        # Normalise in float32, then cast back to original dtype.
+        x_normed = x.float() / rms
+        # Scale by learned weights (cast to input dtype for type safety).
+        return (x_normed * self.weight.float()).to(x.dtype)
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, eps={self.eps}"

cyberslm/model/rope.py

@@ -0,0 +1,291 @@
+"""
+Rotary Position Embedding (RoPE)
+=================================
+Reference: "RoFormer: Enhanced Transformer with Rotary Position Embedding"
+           (Su et al., 2021) — https://arxiv.org/abs/2104.09864
+
+Mathematical definition
+-----------------------
+For a query (or key) vector **q** at position ``m`` with head dimension ``d``:
+
+  1. Partition **q** into pairs: (q₁, q₂), (q₃, q₄), ..., (q_{d-1}, q_d).
+
+  2. For each pair index ``i ∈ {0, 1, ..., d/2 - 1}`` define the frequency:
+
+         θᵢ = 1 / base^(2i / d)         (base = 10 000)
+
+  3. Apply a 2-D rotation to each pair at position ``m``:
+
+         R(m, θᵢ) · (q_{2i}, q_{2i+1}) =
+             (q_{2i} cos(m·θᵢ) − q_{2i+1} sin(m·θᵢ),
+              q_{2i} sin(m·θᵢ) + q_{2i+1} cos(m·θᵢ))
+
+This is equivalent to multiplying **q** (viewed as complex numbers) by
+``exp(i · m · θ)``, which preserves the inner product of relative positions:
+
+    ⟨R(m)q, R(n)k⟩ depends only on (m − n),
+
+giving translation-equivariant attention without absolute position tokens.
+
+Efficient implementation
+------------------------
+The rotation can be expressed without complex arithmetic:
+
+    q_rot = [q_even · cos − q_odd · sin,
+             q_even · sin + q_odd · cos]
+
+where ``q_even = q[..., 0::2]``, ``q_odd = q[..., 1::2]``.
+
+Numerically interleaved form (even/odd) vs. split-half form (first/second half)
+are equivalent; we use the interleaved form for clarity.
+
+Precomputation
+--------------
+``cos`` and ``sin`` tensors of shape ``(max_seq_len, head_dim // 2)`` are
+computed once and registered as non-parameter buffers so they move with the
+module (CPU ↔ GPU) and are not included in ``state_dict`` checkpoints.
+
+Stability notes
+---------------
+- Frequencies are computed in float64 then cast to float32 to minimise
+  floating-point error in ``pow`` and ``arange``.
+- All rotations are executed in float32 to prevent loss of precision.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class RotaryPositionEmbedding(nn.Module):
+    """
+    Precomputed Rotary Position Embedding cache.
+
+    Registers ``cos`` and ``sin`` buffers of shape
+    ``(max_seq_len, head_dim // 2)`` at construction time. Applying RoPE
+    to a query or key tensor costs only element-wise multiplications and
+    additions — no matrix multiplications.
+
+    Parameters
+    ----------
+    head_dim : int
+        Dimension of each attention head.  Must be even.
+    max_seq_len : int
+        Maximum sequence length to pre-compute.  Sequences longer than this
+        will raise an error at runtime.
+    base : int
+        RoPE base frequency (10 000 in the original paper).
+
+    Shape of ``apply``
+    ------------------
+    Input  : ``(batch, num_heads, seq_len, head_dim)``
+    Output : ``(batch, num_heads, seq_len, head_dim)``
+    """
+
+    def __init__(
+        self,
+        head_dim: int,
+        max_seq_len: int,
+        base: int = 10_000,
+    ) -> None:
+        super().__init__()
+
+        if head_dim <= 0 or head_dim % 2 != 0:
+            raise ValueError(
+                f"head_dim must be a positive even integer, got {head_dim}"
+            )
+        if max_seq_len <= 0:
+            raise ValueError(f"max_seq_len must be positive, got {max_seq_len}")
+        if base <= 0:
+            raise ValueError(f"base must be positive, got {base}")
+
+        self.head_dim = head_dim
+        self.max_seq_len = max_seq_len
+        self.base = base
+
+        # Pre-compute and register buffers (not model parameters).
+        cos_cache, sin_cache = self._build_cache(head_dim, max_seq_len, base)
+        self.register_buffer("cos_cache", cos_cache, persistent=False)
+        self.register_buffer("sin_cache", sin_cache, persistent=False)
+
+    @staticmethod
+    def _build_cache(
+        head_dim: int,
+        max_seq_len: int,
+        base: int,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Build ``(cos, sin)`` caches of shape ``(max_seq_len, head_dim // 2)``.
+
+        Computation is performed in float64 for precision, then cast to
+        float32 for storage.
+
+        Returns
+        -------
+        Tuple[Tensor, Tensor]
+            ``cos_cache`` and ``sin_cache``, each of shape
+            ``(max_seq_len, head_dim // 2)``.
+        """
+        half_dim = head_dim // 2
+
+        # θᵢ = 1 / base^(2i / head_dim)  for i ∈ {0, ..., half_dim - 1}
+        # Computed in float64 to avoid precision loss in the exponent.
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, head_dim, 2, dtype=torch.float64) / head_dim)
+        )
+        # Shape: (half_dim,)
+
+        # Position indices m ∈ {0, 1, ..., max_seq_len - 1}
+        positions = torch.arange(max_seq_len, dtype=torch.float64)
+        # Shape: (max_seq_len,)
+
+        # Outer product: angles[m, i] = m * θᵢ
+        angles = torch.outer(positions, inv_freq)
+        # Shape: (max_seq_len, half_dim)
+
+        cos_cache = angles.cos().to(torch.float32)
+        sin_cache = angles.sin().to(torch.float32)
+
+        return cos_cache, sin_cache
+
+    @staticmethod
+    def _rotate_half(x: Tensor) -> Tensor:
+        """
+        Rotate the last dimension by interleaving even/odd pairs.
+
+        For input ``x`` of shape ``(..., head_dim)``:
+
+            x_even = x[..., 0::2]  (positions 0, 2, 4, ...)
+            x_odd  = x[..., 1::2]  (positions 1, 3, 5, ...)
+
+        Returns ``[-x_odd, x_even]`` interleaved back into ``(..., head_dim)``.
+        This is the standard rotation that implements the complex-number trick.
+
+        Parameters
+        ----------
+        x : Tensor
+            Shape ``(..., head_dim)`` where ``head_dim`` is even.
+
+        Returns
+        -------
+        Tensor
+            Same shape as ``x``.
+        """
+        x_even = x[..., 0::2]  # (..., head_dim // 2)
+        x_odd = x[..., 1::2]   # (..., head_dim // 2)
+        # Interleave: stack along new dim then flatten.
+        rotated = torch.stack([-x_odd, x_even], dim=-1)
+        # Shape: (..., head_dim // 2, 2) → (..., head_dim)
+        return rotated.flatten(start_dim=-2)
+
+    def apply(self, x: Tensor, offset: int = 0) -> Tensor:
+        """
+        Apply Rotary Position Embeddings to ``x``.
+
+        Parameters
+        ----------
+        x : Tensor
+            Query or key tensor of shape
+            ``(batch, num_heads, seq_len, head_dim)``.
+        offset : int
+            Absolute position of ``x[..., 0, :]`` in the full sequence.
+
+            This is what makes incremental decoding correct. With a KV cache the
+            model feeds one token at a time, so ``seq_len == 1`` and the naive
+            ``cos_cache[:1]`` would rotate every generated token as if it were at
+            position 0 -- destroying all positional information after the prompt.
+            Passing ``offset=len(cache)`` selects the true absolute position.
+
+        Returns
+        -------
+        Tensor
+            Rotated tensor with the same shape and dtype as ``x``.
+
+        Raises
+        ------
+        ValueError
+            If ``seq_len`` exceeds ``max_seq_len``.
+        """
+        seq_len = x.size(2)
+        if offset < 0:
+            raise ValueError(f"offset must be >= 0, got {offset}")
+        if offset + seq_len > self.max_seq_len:
+            raise ValueError(
+                f"Positions [{offset}, {offset + seq_len}) exceed RoPE cache size "
+                f"{self.max_seq_len}.  Re-instantiate with a larger max_seq_len."
+            )
+
+        # Retrieve cached values for this absolute position span.
+        # cos_cache: (seq_len, head_dim // 2)
+        # sin_cache: (seq_len, head_dim // 2)
+        cos = self.cos_cache[offset : offset + seq_len]  # type: ignore[index]
+        sin = self.sin_cache[offset : offset + seq_len]  # type: ignore[index]
+
+        # Expand to broadcast over batch and head dimensions:
+        # (1, 1, seq_len, head_dim // 2) → broadcasts with (B, H, T, D/2)
+        cos = cos.unsqueeze(0).unsqueeze(0)
+        sin = sin.unsqueeze(0).unsqueeze(0)
+
+        # Interleave cos and sin to match full head_dim.
+        # Each of (cos, sin) has shape (1, 1, T, D/2).
+        # We need (1, 1, T, D) by interleaving even positions with cos,
+        # odd positions with sin.  The _rotate_half trick handles this:
+        #
+        #   x_rot = x * cos_full + rotate_half(x) * sin_full
+        #
+        # where cos_full[..., 0::2] = cos and cos_full[..., 1::2] = cos,
+        # i.e. each cos value applies to both the even AND its paired odd slot.
+        # Achieved by repeating each half-dim value into both slots.
+        cos_full = cos.repeat_interleave(2, dim=-1)  # (1, 1, T, D)
+        sin_full = sin.repeat_interleave(2, dim=-1)  # (1, 1, T, D)
+
+        # Work in float32 for stability, then restore original dtype.
+        x_fp32 = x.float()
+        x_rot = x_fp32 * cos_full + self._rotate_half(x_fp32) * sin_full
+        return x_rot.to(x.dtype)
+
+    def forward(self, x: Tensor, offset: int = 0) -> Tensor:
+        """Alias for :meth:`apply` to support ``nn.Sequential`` usage."""
+        return self.apply(x, offset=offset)
+
+    def extra_repr(self) -> str:
+        return (
+            f"head_dim={self.head_dim}, "
+            f"max_seq_len={self.max_seq_len}, "
+            f"base={self.base}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Functional helper
+# ---------------------------------------------------------------------------
+
+def apply_rope(
+    q: Tensor,
+    k: Tensor,
+    rope: RotaryPositionEmbedding,
+    offset: int = 0,
+) -> Tuple[Tensor, Tensor]:
+    """
+    Apply the same RoPE instance to both query and key tensors.
+
+    Parameters
+    ----------
+    q : Tensor
+        Query tensor of shape ``(batch, num_heads, seq_len, head_dim)``.
+    k : Tensor
+        Key tensor of shape ``(batch, num_heads, seq_len, head_dim)``.
+    rope : RotaryPositionEmbedding
+        Pre-built RoPE module (carries the cos/sin cache on the correct device).
+
+    Returns
+    -------
+    Tuple[Tensor, Tensor]
+        ``(q_rot, k_rot)`` — rotated queries and keys.
+    """
+    return rope.apply(q, offset=offset), rope.apply(k, offset=offset)