| """ANKA-50M RMW3 model architecture. |
| |
| This module contains the exact architecture used by the submitted Base 5B |
| checkpoint. It intentionally depends only on PyTorch. |
| """ |
|
|
| from dataclasses import asdict, dataclass |
| from typing import Optional |
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class RMW3Config: |
| vocab_size: int = 16_384 |
| d_model: int = 512 |
| n_physical_layers: int = 13 |
| recurrent_passes: int = 2 |
| n_heads: int = 8 |
| n_kv_heads: int = 2 |
| d_ff: int = 1_280 |
| reflect_rank: int = 32 |
| morph_rank: int = 368 |
| morph_stages: int = 4 |
| morph_experts: int = 4 |
| max_seq_len: int = 1_024 |
| rope_theta: float = 10_000.0 |
| norm_eps: float = 1e-5 |
|
|
| def to_dict(self) -> dict: |
| return asdict(self) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dimension: int, epsilon: float): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dimension)) |
| self.epsilon = epsilon |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| inverse_rms = torch.rsqrt( |
| x.float().pow(2).mean(dim=-1, keepdim=True) + self.epsilon |
| ) |
| normalized = (x.float() * inverse_rms).to(x.dtype) |
| return normalized * self.weight |
|
|
|
|
| def build_rope_cache( |
| maximum_sequence_length: int, |
| head_dimension: int, |
| theta: float, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| inverse_frequency = 1.0 / ( |
| theta |
| ** ( |
| torch.arange(0, head_dimension, 2, dtype=torch.float32) |
| / head_dimension |
| ) |
| ) |
| positions = torch.arange(maximum_sequence_length, dtype=torch.float32) |
| angles = torch.outer(positions, inverse_frequency) |
| return angles.cos(), angles.sin() |
|
|
|
|
| def apply_rope( |
| x: torch.Tensor, |
| cosine: torch.Tensor, |
| sine: torch.Tensor, |
| ) -> torch.Tensor: |
| sequence_length = x.shape[-2] |
| cosine = cosine[:sequence_length].to(device=x.device, dtype=x.dtype)[ |
| None, None, :, : |
| ] |
| sine = sine[:sequence_length].to(device=x.device, dtype=x.dtype)[ |
| None, None, :, : |
| ] |
| even = x[..., 0::2] |
| odd = x[..., 1::2] |
| rotated = torch.stack( |
| (even * cosine - odd * sine, even * sine + odd * cosine), dim=-1 |
| ) |
| return rotated.flatten(-2) |
|
|
|
|
| class GroupedQueryAttention(nn.Module): |
| def __init__(self, config: RMW3Config): |
| super().__init__() |
| self.number_of_heads = config.n_heads |
| self.number_of_kv_heads = config.n_kv_heads |
| self.head_dimension = config.d_model // config.n_heads |
|
|
| 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.query_projection = nn.Linear( |
| config.d_model, |
| config.n_heads * self.head_dimension, |
| bias=False, |
| ) |
| self.key_projection = nn.Linear( |
| config.d_model, |
| config.n_kv_heads * self.head_dimension, |
| bias=False, |
| ) |
| self.value_projection = nn.Linear( |
| config.d_model, |
| config.n_kv_heads * self.head_dimension, |
| bias=False, |
| ) |
| self.output_projection = nn.Linear( |
| config.n_heads * self.head_dimension, |
| config.d_model, |
| bias=False, |
| ) |
|
|
| cosine, sine = build_rope_cache( |
| config.max_seq_len, |
| self.head_dimension, |
| config.rope_theta, |
| ) |
| self.register_buffer("rope_cosine", cosine, persistent=False) |
| self.register_buffer("rope_sine", sine, persistent=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| batch_size, sequence_length, _ = x.shape |
| query = self.query_projection(x).view( |
| batch_size, |
| sequence_length, |
| self.number_of_heads, |
| self.head_dimension, |
| ).transpose(1, 2) |
| key = self.key_projection(x).view( |
| batch_size, |
| sequence_length, |
| self.number_of_kv_heads, |
| self.head_dimension, |
| ).transpose(1, 2) |
| value = self.value_projection(x).view( |
| batch_size, |
| sequence_length, |
| self.number_of_kv_heads, |
| self.head_dimension, |
| ).transpose(1, 2) |
|
|
| query = apply_rope(query, self.rope_cosine, self.rope_sine) |
| key = apply_rope(key, self.rope_cosine, self.rope_sine) |
| output = F.scaled_dot_product_attention( |
| query, |
| key, |
| value, |
| dropout_p=0.0, |
| is_causal=True, |
| enable_gqa=True, |
| ) |
| output = output.transpose(1, 2).contiguous().view( |
| batch_size, sequence_length, -1 |
| ) |
| return self.output_projection(output) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, config: RMW3Config): |
| super().__init__() |
| self.gate_projection = nn.Linear(config.d_model, config.d_ff, bias=False) |
| self.up_projection = nn.Linear(config.d_model, config.d_ff, bias=False) |
| self.down_projection = nn.Linear(config.d_ff, config.d_model, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.down_projection( |
| F.silu(self.gate_projection(x)) * self.up_projection(x) |
| ) |
|
|
|
|
| class RMW3Block(nn.Module): |
| def __init__(self, config: RMW3Config): |
| super().__init__() |
| self.attention_norm = RMSNorm(config.d_model, config.norm_eps) |
| self.attention = GroupedQueryAttention(config) |
| self.ffn_norm = RMSNorm(config.d_model, config.norm_eps) |
| self.ffn = SwiGLU(config) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = x + self.attention(self.attention_norm(x)) |
| x = x + self.ffn(self.ffn_norm(x)) |
| return x |
|
|
|
|
| class ReflectUnit(nn.Module): |
| def __init__(self, config: RMW3Config): |
| super().__init__() |
| self.down_projection = nn.Linear( |
| config.d_model, config.reflect_rank, bias=False |
| ) |
| self.up_projection = nn.Linear( |
| config.reflect_rank, config.d_model, bias=False |
| ) |
| self.gate_logit = nn.Parameter(torch.tensor(-2.1972246, dtype=torch.float32)) |
|
|
| def forward( |
| self, |
| current: torch.Tensor, |
| previous: torch.Tensor, |
| ) -> torch.Tensor: |
| discrepancy = current - previous |
| correction = self.up_projection( |
| F.silu(self.down_projection(discrepancy)) |
| ) |
| gate = torch.sigmoid(self.gate_logit).to(current.dtype) |
| return current + gate * correction |
|
|
|
|
| class MorphExpert(nn.Module): |
| def __init__(self, config: RMW3Config): |
| super().__init__() |
| self.down_projection = nn.Linear( |
| config.d_model, config.morph_rank, bias=False |
| ) |
| self.up_projection = nn.Linear( |
| config.morph_rank, config.d_model, bias=False |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.up_projection(F.silu(self.down_projection(x))) |
|
|
|
|
| class SparseMorphStage(nn.Module): |
| def __init__(self, config: RMW3Config): |
| super().__init__() |
| self.number_of_experts = config.morph_experts |
| self.router = nn.Linear( |
| config.d_model, config.morph_experts, bias=False |
| ) |
| self.experts = nn.ModuleList( |
| [MorphExpert(config) for _ in range(config.morph_experts)] |
| ) |
| self.scale_logit = nn.Parameter(torch.tensor(-4.0, dtype=torch.float32)) |
| self.last_loads: Optional[torch.Tensor] = None |
| self.last_entropy: Optional[torch.Tensor] = None |
|
|
| def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| original_shape = x.shape |
| flat = x.reshape(-1, original_shape[-1]) |
| probabilities = torch.softmax(self.router(flat).float(), dim=-1) |
| choices = probabilities.argmax(dim=-1) |
| output = torch.zeros_like(flat) |
| loads = [] |
|
|
| for expert_index, expert in enumerate(self.experts): |
| selected_mask = choices == expert_index |
| loads.append(selected_mask.float().mean()) |
| indices = torch.nonzero(selected_mask, as_tuple=False).flatten() |
| if indices.numel() == 0: |
| continue |
|
|
| selected_input = flat.index_select(0, indices) |
| expert_output = expert(selected_input).to(output.dtype) |
| selected_probability = probabilities.index_select(0, indices)[ |
| :, expert_index |
| ].to(output.dtype) |
| straight_through = ( |
| 1.0 + selected_probability - selected_probability.detach() |
| ) |
| output.index_copy_( |
| 0, |
| indices, |
| expert_output * straight_through.unsqueeze(-1), |
| ) |
|
|
| load_tensor = torch.stack(loads) |
| importance = probabilities.mean(dim=0) |
| balance_loss = self.number_of_experts * torch.sum( |
| load_tensor * importance |
| ) |
| entropy = -( |
| probabilities * probabilities.clamp_min(1e-9).log() |
| ).sum(dim=-1).mean() |
| self.last_loads = load_tensor.detach() |
| self.last_entropy = entropy.detach() |
| scale = torch.sigmoid(self.scale_logit).to(output.dtype) |
| return x + scale * output.view(original_shape), balance_loss |
|
|
|
|
| class RMW3(nn.Module): |
| def __init__(self, config: Optional[RMW3Config] = None): |
| super().__init__() |
| self.config = config if config is not None else RMW3Config() |
| config = self.config |
|
|
| self.token_embeddings = nn.Embedding(config.vocab_size, config.d_model) |
| self.blocks = nn.ModuleList( |
| [RMW3Block(config) for _ in range(config.n_physical_layers)] |
| ) |
| self.reflect_units = nn.ModuleList( |
| [ReflectUnit(config) for _ in range(config.n_physical_layers)] |
| ) |
| self.morph_stages = nn.ModuleList( |
| [SparseMorphStage(config) for _ in range(config.morph_stages)] |
| ) |
| self.final_norm = RMSNorm(config.d_model, config.norm_eps) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
|
|
| self.apply(self._initialize_weights) |
| self.lm_head.weight = self.token_embeddings.weight |
| self._scale_residual_outputs() |
|
|
| @staticmethod |
| def _initialize_weights(module: nn.Module) -> None: |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
| def _scale_residual_outputs(self) -> None: |
| residual_scale = 1.0 / math.sqrt( |
| 2 |
| * self.config.n_physical_layers |
| * self.config.recurrent_passes |
| ) |
| with torch.no_grad(): |
| for block in self.blocks: |
| block.attention.output_projection.weight.mul_(residual_scale) |
| block.ffn.down_projection.weight.mul_(residual_scale) |
| for reflect_unit in self.reflect_units: |
| reflect_unit.up_projection.weight.mul_(residual_scale) |
| for stage in self.morph_stages: |
| for expert in stage.experts: |
| expert.up_projection.weight.mul_(residual_scale) |
|
|
| def forward( |
| self, |
| input_ids: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| if input_ids.ndim != 2: |
| raise ValueError("input_ids must have shape [batch, sequence]") |
| if input_ids.shape[1] > self.config.max_seq_len: |
| raise ValueError( |
| f"Sequence length {input_ids.shape[1]} exceeds " |
| f"the limit {self.config.max_seq_len}" |
| ) |
|
|
| x = self.token_embeddings(input_ids) |
| first_pass_states = [] |
| balance_losses = [] |
| effective_layer_index = 0 |
| total_effective_layers = ( |
| self.config.n_physical_layers * self.config.recurrent_passes |
| ) |
|
|
| for recurrent_pass in range(self.config.recurrent_passes): |
| for layer_index, block in enumerate(self.blocks): |
| if recurrent_pass == 1: |
| x = self.reflect_units[layer_index]( |
| x, first_pass_states[layer_index] |
| ) |
|
|
| x = block(x) |
| if recurrent_pass == 0: |
| first_pass_states.append(x) |
|
|
| current_stage = min( |
| self.config.morph_stages - 1, |
| effective_layer_index |
| * self.config.morph_stages |
| // total_effective_layers, |
| ) |
| next_stage = min( |
| self.config.morph_stages - 1, |
| (effective_layer_index + 1) |
| * self.config.morph_stages |
| // total_effective_layers, |
| ) |
| stage_finished = ( |
| next_stage != current_stage |
| or effective_layer_index == total_effective_layers - 1 |
| ) |
| if stage_finished: |
| x, balance_loss = self.morph_stages[current_stage](x) |
| balance_losses.append(balance_loss) |
|
|
| effective_layer_index += 1 |
|
|
| logits = self.lm_head(self.final_norm(x)) |
| auxiliary_loss = ( |
| torch.stack(balance_losses).mean() |
| if balance_losses |
| else logits.new_zeros(()) |
| ) |
| return logits, auxiliary_loss |
|
|
|
|
| RMW3_CONFIG = RMW3Config() |
|
|
|
|
| def count_parameters(model: nn.Module) -> int: |
| """Count unique trainable parameters, including tied embeddings once.""" |
| return sum(parameter.numel() for parameter in model.parameters()) |
|
|
|
|
| __all__ = ["RMW3", "RMW3Config", "RMW3_CONFIG", "count_parameters"] |
|
|
