| |
| |
| |
| """Llama-style decoder-only Transformer for German Commons pretraining. |
| |
| This module is the larger German-language base-model counterpart to the |
| original ``MiniLeipzigGPT``. It intentionally uses only PyTorch and standard |
| Python. There are no imports from Hugging Face ``transformers`` or any other |
| pre-built model library, and no pre-trained weights are consumed. |
| |
| Architecture |
| ------------ |
| The implementation follows the decoder-only Transformer family introduced by |
| Vaswani et al. (2017), using modern small-model choices that are common in |
| Llama-style language models: |
| |
| * causal self-attention with PyTorch scaled-dot-product attention; |
| * RoPE rotary position embeddings instead of learned absolute positions; |
| * RMSNorm pre-normalization; |
| * SwiGLU feed-forward layers; |
| * bias-free linear layers; |
| * tied token embedding / output projection by default; |
| * optional grouped-query attention via ``n_kv_heads``. |
| |
| Training provenance |
| ------------------- |
| The model is designed to train from random weights on locally prepared German |
| Commons token shards produced with the repository's own byte-level BPE |
| tokenizer. The implementation keeps the architecture and training inputs |
| inspectable within the project. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| from typing import Dict, Optional, Tuple |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class GermanGPTConfig: |
| """Configuration for :class:`GermanGPT`. |
| |
| The field names mirror ``configs.ModelConfig`` so checked training plans can |
| be passed directly into this model. |
| """ |
|
|
| name: str = "german-commons" |
| vocab_size: int = 32768 |
| d_model: int = 1024 |
| n_layers: int = 24 |
| n_heads: int = 16 |
| n_kv_heads: int = 16 |
| d_ff: int = 2816 |
| context_len: int = 2048 |
| rope_theta: float = 10000.0 |
| norm: str = "rmsnorm" |
| mlp: str = "swiglu" |
| tie_embeddings: bool = True |
| dropout: float = 0.0 |
|
|
| @property |
| def head_dim(self) -> int: |
| """Return the per-head hidden dimension.""" |
|
|
| return self.d_model // self.n_heads |
|
|
| @classmethod |
| def from_dict(cls, data: Dict[str, object]) -> "GermanGPTConfig": |
| """Build a config from a plain checkpoint dictionary.""" |
|
|
| return cls( |
| name=str(data.get("name", "german-commons")), |
| vocab_size=int(data["vocab_size"]), |
| d_model=int(data.get("d_model", 1024)), |
| n_layers=int(data.get("n_layers", 24)), |
| n_heads=int(data.get("n_heads", 16)), |
| n_kv_heads=int(data.get("n_kv_heads", data.get("n_heads", 16))), |
| d_ff=int(data.get("d_ff", 2816)), |
| context_len=int(data.get("context_len", 2048)), |
| rope_theta=float(data.get("rope_theta", 10000.0)), |
| norm=str(data.get("norm", "rmsnorm")), |
| mlp=str(data.get("mlp", "swiglu")), |
| tie_embeddings=bool(data.get("tie_embeddings", True)), |
| dropout=float(data.get("dropout", 0.0)), |
| ) |
|
|
| def to_dict(self) -> Dict[str, object]: |
| """Return a JSON/checkpoint-friendly dictionary.""" |
|
|
| return asdict(self) |
|
|
|
|
| def config_from_model_config(config: object, vocab_size: Optional[int] = None) -> GermanGPTConfig: |
| """Convert ``configs.ModelConfig`` or a similar object to ``GermanGPTConfig``.""" |
|
|
| data = config.__dict__.copy() |
| if vocab_size is not None: |
| data["vocab_size"] = vocab_size |
| return GermanGPTConfig.from_dict(data) |
|
|
|
|
| class RMSNorm(nn.Module): |
| """Root-mean-square normalization used in Llama-style blocks.""" |
|
|
| def __init__(self, dim: int, eps: float = 1e-6) -> None: |
| super().__init__() |
| self.eps = eps |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Normalize the final dimension of ``x`` without mean subtraction.""" |
|
|
| dtype = x.dtype |
| x_float = x.float() |
| variance = x_float.pow(2).mean(dim=-1, keepdim=True) |
| normalized = x_float * torch.rsqrt(variance + self.eps) |
| return (normalized.to(dtype) * self.weight) |
|
|
|
|
| class RotaryEmbedding(nn.Module): |
| """Precomputed RoPE cosine/sine tables for causal attention.""" |
|
|
| def __init__(self, dim: int, max_seq_len: int, theta: float = 10000.0) -> None: |
| super().__init__() |
| if dim % 2 != 0: |
| raise ValueError("RoPE head dimension must be even") |
| inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) |
| positions = torch.arange(max_seq_len, dtype=torch.float32) |
| freqs = torch.outer(positions, inv_freq) |
| self.register_buffer("cos", freqs.cos(), persistent=False) |
| self.register_buffer("sin", freqs.sin(), persistent=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Apply rotary embeddings to ``x`` shaped ``(B, T, H, D)``.""" |
|
|
| seq_len = x.size(1) |
| cos = self.cos[:seq_len].to(dtype=x.dtype, device=x.device).view(1, seq_len, 1, -1) |
| sin = self.sin[:seq_len].to(dtype=x.dtype, device=x.device).view(1, seq_len, 1, -1) |
| even = x[..., 0::2] |
| odd = x[..., 1::2] |
| rotated_even = even * cos - odd * sin |
| rotated_odd = even * sin + odd * cos |
| return torch.stack((rotated_even, rotated_odd), dim=-1).flatten(-2) |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| """Causal multi-head attention with optional grouped-query attention.""" |
|
|
| def __init__(self, config: GermanGPTConfig) -> None: |
| super().__init__() |
| if config.d_model % config.n_heads != 0: |
| raise ValueError("d_model must be divisible by n_heads") |
| if config.n_heads % config.n_kv_heads != 0: |
| raise ValueError("n_heads must be divisible by n_kv_heads") |
|
|
| self.n_heads = config.n_heads |
| self.n_kv_heads = config.n_kv_heads |
| self.head_dim = config.head_dim |
| self.kv_repeat = config.n_heads // config.n_kv_heads |
|
|
| self.q_proj = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False) |
| self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False) |
| self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=False) |
| self.o_proj = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=False) |
| self.dropout = float(config.dropout) |
| self.rope = RotaryEmbedding(config.head_dim, config.context_len, config.rope_theta) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Apply causal self-attention to hidden states ``x``.""" |
|
|
| batch_size, seq_len, _ = x.shape |
| q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim) |
| k = self.k_proj(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim) |
| v = self.v_proj(x).view(batch_size, seq_len, self.n_kv_heads, self.head_dim) |
|
|
| q = self.rope(q) |
| k = self.rope(k) |
|
|
| if self.kv_repeat > 1: |
| k = k.repeat_interleave(self.kv_repeat, dim=2) |
| v = v.repeat_interleave(self.kv_repeat, dim=2) |
|
|
| q = q.transpose(1, 2) |
| k = k.transpose(1, 2) |
| v = v.transpose(1, 2) |
|
|
| y = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=None, |
| dropout_p=self.dropout if self.training else 0.0, |
| is_causal=True, |
| ) |
| y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, -1) |
| return self.o_proj(y) |
|
|
|
|
| class SwiGLU(nn.Module): |
| """Bias-free SwiGLU feed-forward layer.""" |
|
|
| def __init__(self, config: GermanGPTConfig) -> None: |
| super().__init__() |
| self.gate_proj = nn.Linear(config.d_model, config.d_ff, bias=False) |
| self.up_proj = nn.Linear(config.d_model, config.d_ff, bias=False) |
| self.down_proj = nn.Linear(config.d_ff, config.d_model, bias=False) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Apply ``SiLU(gate) * up`` followed by down projection.""" |
|
|
| return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| """Pre-normalized decoder block.""" |
|
|
| def __init__(self, config: GermanGPTConfig) -> None: |
| super().__init__() |
| self.attn_norm = RMSNorm(config.d_model) |
| self.attn = CausalSelfAttention(config) |
| self.ffn_norm = RMSNorm(config.d_model) |
| self.ffn = SwiGLU(config) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| """Apply attention and feed-forward residual paths.""" |
|
|
| x = x + self.attn(self.attn_norm(x)) |
| x = x + self.ffn(self.ffn_norm(x)) |
| return x |
|
|
|
|
| class GermanGPT(nn.Module): |
| """Decoder-only German base language model.""" |
|
|
| def __init__(self, config: GermanGPTConfig) -> None: |
| super().__init__() |
| validate_config(config) |
| self.config = config |
| self.token_embedding = nn.Embedding(config.vocab_size, config.d_model) |
| self.dropout = nn.Dropout(config.dropout) |
| self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)]) |
| self.norm = RMSNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| if config.tie_embeddings: |
| self.lm_head.weight = self.token_embedding.weight |
|
|
| self.apply(self._init_weights) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| labels: Optional[torch.Tensor] = None, |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
| """Run the model and optionally compute next-token loss.""" |
|
|
| if input_ids.ndim != 2: |
| raise ValueError("input_ids must have shape (batch, time)") |
| _, seq_len = input_ids.shape |
| if seq_len > self.config.context_len: |
| raise ValueError( |
| f"Sequence length {seq_len} exceeds context_len {self.config.context_len}" |
| ) |
|
|
| x = self.dropout(self.token_embedding(input_ids)) |
| for block in self.blocks: |
| x = block(x) |
| x = self.norm(x) |
| logits = self.lm_head(x) |
|
|
| loss = None |
| if labels is not None: |
| loss = F.cross_entropy( |
| logits.reshape(-1, logits.size(-1)), |
| labels.reshape(-1), |
| ignore_index=-100, |
| ) |
| return logits, loss |
|
|
| @staticmethod |
| def _init_weights(module: nn.Module) -> None: |
| """Initialize trainable weights with a GPT/Llama-compatible normal init.""" |
|
|
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
|
|
| def validate_config(config: GermanGPTConfig) -> None: |
| """Validate architectural invariants before allocating the model.""" |
|
|
| if config.vocab_size <= 0: |
| raise ValueError("vocab_size must be positive") |
| if config.d_model <= 0 or config.n_layers <= 0: |
| raise ValueError("d_model and n_layers must be positive") |
| if config.context_len <= 0: |
| raise ValueError("context_len must be positive") |
| if config.d_model % config.n_heads != 0: |
| raise ValueError("d_model must be divisible by n_heads") |
| if config.n_heads % config.n_kv_heads != 0: |
| raise ValueError("n_heads must be divisible by n_kv_heads") |
| if config.head_dim % 2 != 0: |
| raise ValueError("head_dim must be even for RoPE") |
| if config.norm != "rmsnorm": |
| raise ValueError("Only rmsnorm is implemented") |
| if config.mlp != "swiglu": |
| raise ValueError("Only swiglu is implemented") |
| if not 0.0 <= config.dropout < 1.0: |
| raise ValueError("dropout must be in [0, 1)") |
|
|
|
|
| def count_parameters(model: nn.Module) -> int: |
| """Return the number of unique trainable parameters.""" |
|
|
| seen = set() |
| total = 0 |
| for parameter in model.parameters(): |
| if id(parameter) in seen: |
| continue |
| seen.add(id(parameter)) |
| if parameter.requires_grad: |
| total += parameter.numel() |
| return total |
|
|
