import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, PretrainedConfig from transformers.generation import GenerationMixin from transformers.modeling_outputs import CausalLMOutput class TinyQwen3NoveltyConfig(PretrainedConfig): model_type = "tinyqwen3_novelty" def __init__( self, vocab_size=4098, hidden_size=256, intermediate_size=896, num_hidden_layers=8, num_attention_heads=8, num_key_value_heads=4, head_dim=32, rms_norm_eps=1e-6, rope_theta=2500.0, max_position_embeddings=1024, tie_word_embeddings=True, initializer_range=0.02, bos_token_id=1, eos_token_id=2, pad_token_id=2, novelty_gate_floor=0.05, novelty_gate_type="math_rms_abs_delta", im_start_token_id=4096, im_end_token_id=4097, **kwargs, ): self.vocab_size = vocab_size self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.head_dim = head_dim self.rms_norm_eps = rms_norm_eps self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings self.tie_word_embeddings = tie_word_embeddings self.initializer_range = initializer_range self.novelty_gate_floor = novelty_gate_floor self.novelty_gate_type = novelty_gate_type self.im_start_token_id = im_start_token_id self.im_end_token_id = im_end_token_id super().__init__( bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs, ) class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x): return self.weight.to(dtype=x.dtype) * x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) def rotate_half(x): x1, x2 = x.chunk(2, dim=-1) return torch.cat((-x2, x1), dim=-1) class RotaryEmbedding(nn.Module): def __init__(self, head_dim, theta): super().__init__() inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) self.register_buffer("inv_freq", inv_freq) def forward(self, seq_len, device): pos = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype) freqs = torch.outer(pos, self.inv_freq.to(device)) emb = torch.cat((freqs, freqs), dim=-1) return emb.cos()[None, None, :, :], emb.sin()[None, None, :, :] def apply_rope(x, cos, sin): return (x * cos.to(dtype=x.dtype)) + (rotate_half(x) * sin.to(dtype=x.dtype)) def causal_attention(q, k, v): scores = (q.float() @ k.float().transpose(-2, -1)) / (q.size(-1) ** 0.5) causal_mask = torch.ones( scores.size(-2), scores.size(-1), dtype=torch.bool, device=scores.device, ).triu(1) scores = scores.masked_fill(causal_mask, torch.finfo(scores.dtype).min) probs = F.softmax(scores, dim=-1).to(dtype=v.dtype) return probs @ v class MathNoveltyGate(nn.Module): def __init__(self, head_dim, floor=0.05): super().__init__() del head_dim self.floor = floor self.last_gate = None def forward(self, heads): context = (heads.sum(dim=1, keepdim=True) - heads) / (heads.size(1) - 1) scale = heads.pow(2).mean(dim=-1, keepdim=True).sqrt() + context.pow(2).mean(dim=-1, keepdim=True).sqrt() + 1e-6 score = (heads - context).abs() / scale gate = self.floor + (1.0 - self.floor) * score.clamp(0.0, 1.0) compiler = getattr(torch, "compiler", None) if compiler is None or not compiler.is_compiling(): self.last_gate = gate.detach() return heads * gate class NoveltyGQA(nn.Module): def __init__(self, config): super().__init__() dim = config.hidden_size n_heads = config.num_attention_heads n_kv_heads = config.num_key_value_heads self.dim = dim self.n_heads = n_heads self.n_kv_heads = n_kv_heads self.head_dim = dim // n_heads self.kv_dim = n_kv_heads * self.head_dim self.kv_repeat = n_heads // n_kv_heads self.q_proj = nn.Linear(dim, dim, bias=False) self.k_proj = nn.Linear(dim, self.kv_dim, bias=False) self.v_proj = nn.Linear(dim, self.kv_dim, bias=False) self.o_proj = nn.Linear(dim, dim, bias=False) self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) rope_theta = getattr(config, "rope_theta", 1000000.0) self.rope = RotaryEmbedding(self.head_dim, rope_theta) self.novelty = MathNoveltyGate(self.head_dim, floor=getattr(config, "novelty_gate_floor", 0.05)) def forward(self, x): bsz, seq_len, _ = x.shape q = self.q_proj(x).view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) q = self.q_norm(q) k = self.k_norm(k) cos, sin = self.rope(seq_len, x.device) q = apply_rope(q, cos, sin) k = apply_rope(k, cos, sin) k = k.repeat_interleave(self.kv_repeat, dim=1) v = v.repeat_interleave(self.kv_repeat, dim=1) heads = causal_attention(q, k, v) heads = self.novelty(heads) out = heads.transpose(1, 2).contiguous().view(bsz, seq_len, self.dim) return self.o_proj(out) class SwiGLU(nn.Module): def __init__(self, dim, hidden_dim): super().__init__() self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) self.up_proj = nn.Linear(dim, hidden_dim, bias=False) self.down_proj = nn.Linear(hidden_dim, dim, bias=False) def forward(self, x): return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) class TinyQwen3NoveltyBlock(nn.Module): def __init__(self, config): super().__init__() self.input_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attn = NoveltyGQA(config) self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.mlp = SwiGLU(config.hidden_size, config.intermediate_size) def forward(self, x): x = x + self.attn(self.input_norm(x)) x = x + self.mlp(self.post_attn_norm(x)) return x class TinyQwen3NoveltyForCausalLM(PreTrainedModel, GenerationMixin): config_class = TinyQwen3NoveltyConfig base_model_prefix = "" _no_split_modules = ["TinyQwen3NoveltyBlock"] _tied_weights_keys = {} def __init__(self, config): super().__init__(config) self.config = config self.all_tied_weights_keys = {} self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.layers = nn.ModuleList(TinyQwen3NoveltyBlock(config) for _ in range(config.num_hidden_layers)) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def get_output_embeddings(self): return self.embed_tokens def set_output_embeddings(self, value): self.embed_tokens = value def prepare_inputs_for_generation(self, input_ids, attention_mask=None, token_type_ids=None, **kwargs): max_len = getattr(self.config, "max_position_embeddings", None) if max_len is not None and input_ids.shape[-1] > max_len: input_ids = input_ids[:, -max_len:] if attention_mask is not None: attention_mask = attention_mask[:, -max_len:] if token_type_ids is not None: token_type_ids = token_type_ids[:, -max_len:] model_inputs = {"input_ids": input_ids} if attention_mask is not None: model_inputs["attention_mask"] = attention_mask if token_type_ids is not None: model_inputs["token_type_ids"] = token_type_ids return model_inputs def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None, return_dict=True, **kwargs): del attention_mask, token_type_ids, kwargs x = self.embed_tokens(input_ids) for layer in self.layers: x = layer(x) x = self.norm(x) logits = x @ self.embed_tokens.weight.t() loss = None if labels is not None: loss = F.cross_entropy(logits[:, :-1, :].contiguous().view(-1, self.config.vocab_size), labels[:, 1:].contiguous().view(-1)) if not return_dict: return (loss, logits) if loss is not None else (logits,) return CausalLMOutput(loss=loss, logits=logits)