from __future__ import annotations import math from dataclasses import dataclass from pathlib import Path from typing import Any import torch import torch.nn.functional as F from safetensors.torch import load_model, save_model from torch import nn from torch.nn.attention import SDPBackend, sdpa_kernel from torch.utils.checkpoint import checkpoint from tiny_gdn.config import TinyGDNConfig try: # Import the module directly — `from fla.layers import GatedDeltaNet2` # executes layers/__init__.py and eagerly loads every attention kernel. from fla.layers.gdn2 import GatedDeltaNet2 except ImportError as import_error: GatedDeltaNet2 = None FLA_IMPORT_ERROR: ImportError | None = import_error else: FLA_IMPORT_ERROR = None @dataclass class TinyGDNOutput: loss: torch.Tensor | None logits: torch.Tensor | None main_loss: torch.Tensor | None mtp_loss: torch.Tensor | None z_loss: torch.Tensor | None hidden_states: torch.Tensor | None = None class RMSNorm(nn.Module): """Zero-centered RMSNorm as used by Qwen3-Next.""" def __init__(self, hidden_size: int, eps: float) -> None: super().__init__() self.weight = nn.Parameter(torch.zeros(hidden_size)) self.eps = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype normalized = hidden_states.float() normalized = normalized * torch.rsqrt(normalized.square().mean(dim=-1, keepdim=True) + self.eps) normalized = normalized * (1.0 + self.weight.float()) return normalized.to(dtype=input_dtype) class RotaryEmbedding(nn.Module): def __init__(self, rotary_dim: int, rope_theta: float) -> None: super().__init__() inverse_frequency = 1.0 / ( rope_theta ** ( torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim ) ) self.rotary_dim = rotary_dim self.register_buffer("inverse_frequency", inverse_frequency, persistent=False) def forward( self, sequence_length: int, device: torch.device, dtype: torch.dtype, position_offset: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: positions = torch.arange( position_offset, position_offset + sequence_length, device=device, dtype=torch.float32, ) frequencies = torch.outer(positions, self.inverse_frequency.float()) embeddings = torch.cat((frequencies, frequencies), dim=-1) return embeddings.cos().to(dtype=dtype), embeddings.sin().to(dtype=dtype) def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor: first, second = hidden_states.chunk(2, dim=-1) return torch.cat((-second, first), dim=-1) def apply_rotary_embedding( query: torch.Tensor, key: torch.Tensor, cosine: torch.Tensor, sine: torch.Tensor, rotary_dim: int, ) -> tuple[torch.Tensor, torch.Tensor]: cosine = cosine[None, None, :, :] sine = sine[None, None, :, :] query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:] key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:] query_rotary = query_rotary * cosine + rotate_half(query_rotary) * sine key_rotary = key_rotary * cosine + rotate_half(key_rotary) * sine return ( torch.cat((query_rotary, query_pass), dim=-1), torch.cat((key_rotary, key_pass), dim=-1), ) class GatedGroupedQueryAttention(nn.Module): """QK-normalized, partially rotary GQA with a learned sigmoid output gate.""" def __init__(self, config: TinyGDNConfig) -> None: super().__init__() self.num_heads = config.num_attention_heads self.num_key_value_heads = config.num_key_value_heads self.head_dim = config.attention_head_dim self.rotary_dim = config.rotary_dim self.dropout = config.attention_dropout query_size = self.num_heads * self.head_dim key_value_size = self.num_key_value_heads * self.head_dim self.q_gate_proj = nn.Linear(config.hidden_size, query_size * 2, bias=False) self.k_proj = nn.Linear(config.hidden_size, key_value_size, bias=False) self.v_proj = nn.Linear(config.hidden_size, key_value_size, bias=False) self.o_proj = nn.Linear(query_size, config.hidden_size, bias=False) self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) self.rotary = RotaryEmbedding(self.rotary_dim, config.rope_theta) def _attention_mask( self, attention_mask: torch.Tensor | None, sequence_length: int, device: torch.device, ) -> torch.Tensor | None: if attention_mask is None: return None if attention_mask.ndim != 2: raise ValueError("attention_mask must have shape [batch, sequence]") if attention_mask.shape[1] != sequence_length: raise ValueError("attention_mask sequence length does not match input") causal = torch.ones( sequence_length, sequence_length, dtype=torch.bool, device=device, ).tril() valid_keys = attention_mask[:, None, None, :].to(dtype=torch.bool, device=device) return causal[None, None, :, :] & valid_keys def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, sequence_length, _ = hidden_states.shape query_and_gate = self.q_gate_proj(hidden_states) query, output_gate = query_and_gate.chunk(2, dim=-1) query = query.view(batch_size, sequence_length, self.num_heads, self.head_dim) key = self.k_proj(hidden_states).view( batch_size, sequence_length, self.num_key_value_heads, self.head_dim, ) value = self.v_proj(hidden_states).view( batch_size, sequence_length, self.num_key_value_heads, self.head_dim, ) query = self.q_norm(query).transpose(1, 2) key = self.k_norm(key).transpose(1, 2) value = value.transpose(1, 2) cosine, sine = self.rotary( sequence_length, device=hidden_states.device, dtype=query.dtype, ) query, key = apply_rotary_embedding( query, key, cosine, sine, rotary_dim=self.rotary_dim, ) sdpa_mask = self._attention_mask( attention_mask, sequence_length, hidden_states.device, ) sdpa_options = { "attn_mask": sdpa_mask, "dropout_p": self.dropout if self.training else 0.0, "is_causal": sdpa_mask is None, "enable_gqa": True, } # Prefer Flash / mem-efficient when available; fall back to MATH for # Windows PyTorch builds that ship without FlashAttention kernels. backends = ( [ SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH, ] if query.is_cuda else [SDPBackend.MATH] ) with sdpa_kernel(backends): attention_output = F.scaled_dot_product_attention( query, key, value, **sdpa_options, ) attention_output = attention_output.transpose(1, 2).reshape( batch_size, sequence_length, -1, ) attention_output = attention_output * torch.sigmoid(output_gate) return self.o_proj(attention_output) class SwiGLU(nn.Module): def __init__(self, config: TinyGDNConfig) -> None: super().__init__() self.gate_up_proj = nn.Linear( config.hidden_size, config.intermediate_size * 2, bias=False, ) self.down_proj = nn.Linear( config.intermediate_size, config.hidden_size, bias=False, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1) return self.down_proj(F.silu(gate) * up) class TinyGDNBlock(nn.Module): def __init__(self, config: TinyGDNConfig, layer_index: int) -> None: super().__init__() layer_type = config.layer_types[layer_index] self.layer_type = layer_type self.token_mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) if layer_type == "gdn2": if GatedDeltaNet2 is None: raise ImportError( "Gated DeltaNet-2 requires the pinned flash-linear-attention dependency" ) from FLA_IMPORT_ERROR self.token_mixer = GatedDeltaNet2( hidden_size=config.hidden_size, expand_v=config.linear_expand_v, head_dim=config.linear_head_dim, num_heads=config.linear_num_heads, num_v_heads=config.linear_num_value_heads, mode="chunk", use_short_conv=True, allow_neg_eigval=config.allow_negative_eigenvalues, conv_size=config.linear_conv_kernel_dim, conv_bias=False, layer_idx=layer_index, norm_eps=config.rms_norm_eps, ) elif layer_type == "full_attention": self.token_mixer = GatedGroupedQueryAttention(config) else: raise ValueError(f"Unsupported layer type: {layer_type}") self.mlp = SwiGLU(config) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: residual = hidden_states normalized = self.token_mixer_norm(hidden_states) if self.layer_type == "gdn2": mixed, _, _ = self.token_mixer( normalized, attention_mask=attention_mask, use_cache=False, ) else: mixed = self.token_mixer(normalized, attention_mask=attention_mask) hidden_states = residual + mixed hidden_states = hidden_states + self.mlp(self.mlp_norm(hidden_states)) return hidden_states class MultiTokenPredictionAdapter(nn.Module): """A lightweight residual adapter for one additional prediction horizon.""" def __init__(self, config: TinyGDNConfig) -> None: super().__init__() self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.down_proj = nn.Linear( config.hidden_size, config.mtp_adapter_rank, bias=False, ) self.up_proj = nn.Linear( config.mtp_adapter_rank, config.hidden_size, bias=False, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: adapted = self.up_proj(F.silu(self.down_proj(self.norm(hidden_states)))) return hidden_states + adapted class TinyGDNForCausalLM(nn.Module): def __init__(self, config: TinyGDNConfig) -> None: super().__init__() self.config = config self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList( TinyGDNBlock(config, layer_index) for layer_index in range(config.num_hidden_layers) ) self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.mtp_adapters = nn.ModuleList( MultiTokenPredictionAdapter(config) for _ in range(config.mtp_num_heads) ) self.gradient_checkpointing = False self.apply(self._initialize_module) self._initialize_residual_projections() def _initialize_module(self, module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.normal_( module.weight, mean=0.0, std=self.config.initializer_range, ) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_( module.weight, mean=0.0, std=self.config.initializer_range, ) def _initialize_residual_projections(self) -> None: residual_std = self.config.initializer_range / math.sqrt( 2 * self.config.num_hidden_layers ) for layer in self.layers: nn.init.normal_( layer.token_mixer.o_proj.weight, mean=0.0, std=residual_std, ) nn.init.normal_( layer.mlp.down_proj.weight, mean=0.0, std=residual_std, ) for adapter in self.mtp_adapters: nn.init.normal_(adapter.up_proj.weight, mean=0.0, std=residual_std) def enable_gradient_checkpointing(self, enabled: bool = True) -> None: self.gradient_checkpointing = enabled def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor: return F.linear(hidden_states, self.embed_tokens.weight) def _run_layer( self, layer: TinyGDNBlock, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None, ) -> torch.Tensor: if self.gradient_checkpointing and self.training: return checkpoint( layer, hidden_states, attention_mask, use_reentrant=False, ) return layer(hidden_states, attention_mask) def _causal_loss( self, hidden_states: torch.Tensor, labels: torch.Tensor, target_offset: int, adapter: nn.Module | None = None, compute_z_loss: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: if target_offset < 0: raise ValueError("target_offset cannot be negative") if target_offset and hidden_states.shape[1] <= target_offset: raise ValueError( f"Sequence length must exceed target offset {target_offset}" ) if target_offset: prediction_states = hidden_states[:, :-target_offset, :] targets = labels[:, target_offset:].contiguous() else: prediction_states = hidden_states targets = labels.contiguous() if adapter is not None: prediction_states = adapter(prediction_states) logits = self.project_to_vocabulary(prediction_states) cross_entropy = F.cross_entropy( logits.reshape(-1, self.config.vocab_size), targets.reshape(-1), ignore_index=-100, ) z_loss = None if compute_z_loss: valid_targets = targets.ne(-100) log_partition = torch.logsumexp(logits.float(), dim=-1) z_loss = log_partition.square()[valid_targets].mean() return cross_entropy, z_loss def forward( self, input_ids: torch.Tensor, labels: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, *, return_logits: bool = True, return_hidden_states: bool = False, labels_are_shifted: bool = False, include_mtp_loss: bool = True, mtp_loss_weight: float | None = None, z_loss_coefficient: float = 0.0, logits_to_keep: int | None = None, ) -> TinyGDNOutput: if input_ids.ndim != 2: raise ValueError("input_ids must have shape [batch, sequence]") if input_ids.shape[1] > self.config.max_position_embeddings: raise ValueError("Input exceeds max_position_embeddings") if labels is not None and labels.shape != input_ids.shape: raise ValueError("labels must have the same shape as input_ids") if z_loss_coefficient < 0.0: raise ValueError("z_loss_coefficient cannot be negative") if logits_to_keep is not None and logits_to_keep <= 0: raise ValueError("logits_to_keep must be positive") effective_mtp_weight = ( self.config.mtp_loss_weight if mtp_loss_weight is None else mtp_loss_weight ) if not 0.0 <= effective_mtp_weight <= 1.0: raise ValueError("mtp_loss_weight must be between zero and one") hidden_states = self.embed_tokens(input_ids) shared_layer_indices = set(self.config.shared_layer_indices) for layer_index, layer in enumerate(self.layers): hidden_states = self._run_layer( layer, hidden_states, attention_mask, ) if layer_index in shared_layer_indices: hidden_states = self._run_layer( layer, hidden_states, attention_mask, ) hidden_states = self.final_norm(hidden_states) main_loss = None mtp_loss = None z_loss = None total_loss = None if labels is not None: main_target_offset = 0 if labels_are_shifted else 1 main_loss, z_loss = self._causal_loss( hidden_states, labels, target_offset=main_target_offset, compute_z_loss=z_loss_coefficient > 0.0, ) if self.mtp_adapters and include_mtp_loss: auxiliary_losses = [ self._causal_loss( hidden_states, labels, target_offset=( head_index + 1 if labels_are_shifted else head_index + 2 ), adapter=adapter, )[0] for head_index, adapter in enumerate(self.mtp_adapters) ] mtp_loss = torch.stack(auxiliary_losses).mean() total_loss = main_loss + effective_mtp_weight * mtp_loss else: total_loss = main_loss if z_loss is not None: total_loss = total_loss + z_loss_coefficient * z_loss output_states = ( hidden_states if logits_to_keep is None else hidden_states[:, -logits_to_keep:, :] ) logits = self.project_to_vocabulary(output_states) if return_logits else None return TinyGDNOutput( loss=total_loss, logits=logits, main_loss=main_loss, mtp_loss=mtp_loss, z_loss=z_loss, hidden_states=hidden_states if return_hidden_states else None, ) def parameter_report(self) -> dict[str, int]: total = sum(parameter.numel() for parameter in self.parameters()) mtp = sum(parameter.numel() for parameter in self.mtp_adapters.parameters()) embeddings = self.embed_tokens.weight.numel() return { "deployable_core": total - mtp, "training_total": total, "embedding": embeddings, "mtp_auxiliary": mtp, "non_embedding_core": total - mtp - embeddings, } def save_checkpoint(self, output_dir: Path) -> None: output_dir.mkdir(parents=True, exist_ok=True) self.config.save_json(output_dir / "config.json") save_model(self, output_dir / "model.safetensors") @classmethod def from_checkpoint( cls, checkpoint_dir: Path, *, device: str | torch.device = "cpu", dtype: torch.dtype | None = None, ) -> TinyGDNForCausalLM: config = TinyGDNConfig.from_json(checkpoint_dir / "config.json") model = cls(config).to(device=device, dtype=dtype) load_model(model, checkpoint_dir / "model.safetensors", device=str(device)) return model def extra_repr(self) -> str: report = self.parameter_report() return ( f"core_parameters={report['deployable_core']:,}, " f"training_parameters={report['training_total']:,}" ) def get_architecture_metadata(self) -> dict[str, Any]: return { "architecture": self.config.architecture, "layer_types": list(self.config.layer_types), "effective_num_layers": self.config.effective_num_layers, "shared_layer_indices": list(self.config.shared_layer_indices), "parameter_report": self.parameter_report(), "features": [ "32-layer deep-thin parameter allocation", "Gated DeltaNet-2 recurrent memory", "3:1 recurrent-to-full-attention hybrid", "gated grouped-query attention", "QK normalization", "partial rotary embeddings", "zero-centered RMSNorm", "SwiGLU", "tied input-output embeddings", "optional multi-token prediction auxiliaries", ], }