| """RFC-005 decoder using Transformers' Llama attention, RoPE and generation. |
| |
| The explicit model class is required when loading: ordinary Llama does not |
| implement this recipe's parameter-free QK normalization or two-matrix ReLU². |
| """ |
| from dataclasses import asdict |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from transformers import LlamaConfig, LlamaForCausalLM |
| from .recipe import Architecture |
|
|
|
|
| class EnglishBaseConfig(LlamaConfig): |
| model_type = 'fabryka_english_base' |
|
|
| def __init__(self, recipe=None, qk_norm=True, **kwargs): |
| architecture = Architecture(**(recipe or {})) |
| controlled = dict( |
| vocab_size=architecture.vocab_size, hidden_size=architecture.width, |
| intermediate_size=architecture.ffn_width, num_hidden_layers=architecture.layers, |
| num_attention_heads=architecture.query_heads, num_key_value_heads=architecture.kv_heads, |
| head_dim=architecture.head_dim, max_position_embeddings=architecture.context, |
| hidden_act='silu', attention_bias=False, mlp_bias=False, |
| attention_dropout=0.0, tie_word_embeddings=architecture.tied_embeddings, |
| ) |
| for key, value in controlled.items(): |
| if key in kwargs and kwargs[key] != value: |
| raise ValueError(f'{key} conflicts with recorded architecture') |
| kwargs[key] = value |
| if type(qk_norm) is not bool: |
| raise ValueError('qk_norm must be boolean') |
| super().__init__(**kwargs) |
| self.recipe = asdict(architecture) |
| self.qk_norm = qk_norm |
|
|
|
|
| class HeadNormalizedLinear(nn.Linear): |
| """Project, then normalize each head in FP32 before RoPE; no learned gain.""" |
| def __init__(self, source, head_dim, eps): |
| super().__init__(source.in_features, source.out_features, bias=False, |
| device=source.weight.device, dtype=source.weight.dtype) |
| self.weight = source.weight |
| self.head_dim, self.eps = head_dim, eps |
|
|
| def forward(self, x): |
| projected = F.linear(x, self.weight) |
| heads = projected.unflatten(-1, (-1, self.head_dim)) |
| if hasattr(F, 'rms_norm'): |
| normalized = F.rms_norm(heads.float(), (self.head_dim,), eps=self.eps) |
| else: |
| |
| values = heads.float() |
| normalized = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + self.eps) |
| return normalized.to(projected.dtype).flatten(-2) |
|
|
|
|
| class ReluSquaredMLP(nn.Module): |
| def __init__(self, original): |
| super().__init__() |
| self.up_proj = original.up_proj |
| self.down_proj = original.down_proj |
|
|
| def forward(self, x): |
| return self.down_proj(F.relu(self.up_proj(x)).square()) |
|
|
|
|
| class EnglishBaseForCausalLM(LlamaForCausalLM): |
| config_class = EnglishBaseConfig |
|
|
| def __init__(self, config): |
| if not isinstance(config, EnglishBaseConfig): |
| raise TypeError('EnglishBaseConfig is required; byte/Llama checkpoints are distinct') |
| super().__init__(config) |
| for block in self.model.layers: |
| if config.recipe['activation'] == 'relu2': |
| block.mlp = ReluSquaredMLP(block.mlp) |
| if config.qk_norm: |
| for name in ('q_proj', 'k_proj'): |
| projection = getattr(block.self_attn, name) |
| setattr(block.self_attn, name, HeadNormalizedLinear( |
| projection, config.head_dim, config.rms_norm_eps)) |
|
|
| def parameter_report(self): |
| total = sum(p.numel() for p in self.parameters()) |
| expected = Architecture(**self.config.recipe).parameter_count |
| if total != expected: |
| raise ValueError(f'Implemented parameter count {total} differs from expected {expected}') |
| embedding_ids = {id(self.get_input_embeddings().weight), id(self.get_output_embeddings().weight)} |
| non_embedding = sum(p.numel() for p in self.parameters() if id(p) not in embedding_ids) |
| return {'total_parameters': total, 'non_embedding_parameters': non_embedding, |
| 'tied_embeddings': self.get_input_embeddings().weight is self.get_output_embeddings().weight, |
| 'qk_norm': self.config.qk_norm} |
|
|
|
|
| def mark_tied_weights_as_initialized(self, loading_info): |
| """Preserve parameter-tying missing keys for Transformers 5.3 loading. |
| |
| This architecture uses parameter tying only, never shared modules. |
| The remote-code compatibility path in 5.3 removes initialized missing |
| keys before tie_weights, including the lm_head it has just marked. |
| Keep that key so the loader ties it to the loaded input embedding. |
| """ |
| for name in self.all_tied_weights_keys: |
| self.get_parameter(name)._is_hf_initialized = True |
|
|
