"""Standalone inference implementation for YModel3.2. This module contains the complete model definition. Its only third-party dependencies are PyTorch and Transformers. """ from __future__ import annotations import math from pathlib import Path from typing import Optional, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers import GenerationMixin, PreTrainedModel from transformers.configuration_utils import PretrainedConfig from transformers.modeling_outputs import CausalLMOutputWithPast def _normalize_group_sizes(group_size: Union[int, list[int], tuple[int, ...]]) -> tuple[int, ...]: if isinstance(group_size, bool): raise TypeError("group_size must be a positive integer or a list of integers") if isinstance(group_size, int): values = (group_size,) elif isinstance(group_size, (list, tuple)): values = tuple(group_size) else: raise TypeError("group_size must be a positive integer or a list of integers") if not values or any( isinstance(value, bool) or not isinstance(value, int) or value < 1 for value in values ): raise ValueError("group_size must contain positive integers") if len(set(values)) != len(values): raise ValueError("group_size choices must be unique") return values class YConfig32(PretrainedConfig): """Architecture config for YModel3.2.""" model_type = "ynet32" MIN_FULL_BLOCKS = 3 def __init__( self, dropout: float = 0.0, bos_token_id: int = 151644, eos_token_id: int = 151645, pad_token_id: int = 151643, hidden_size: int = 1024, num_hidden_layers: Optional[int] = None, max_position_embeddings: int = 8192, vocab_size: int = 6400, rms_norm_eps: float = 1e-6, rope_theta: float = 5e4, rope_scaling: Optional[dict] = None, dtype: str = "float32", intermediate_size: int = 2048, num_heads: int = 8, num_key_value_heads: Optional[int] = None, head_dim: int = 192, group_size: Union[int, list[int], tuple[int, ...]] = 4, layer_pattern: Optional[str] = "fgggfgggf", use_sengram: bool = False, sengram_bucket_size: Optional[int] = 4096, sengram_topk: int = 2, use_dcp: bool = False, dynamic_compression_probe_hidden_size: int = 768, tie_word_embeddings: bool = True, **kwargs, ): # Derived from group_size; must not be set independently. kwargs.pop("group_size_options", None) kwargs.pop("num_layers", None) 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, ) self.dropout = float(dropout) self.hidden_size = int(hidden_size) self.num_hidden_layers = int(num_hidden_layers) if num_hidden_layers is not None else 8 self.max_position_embeddings = int(max_position_embeddings) self.vocab_size = int(vocab_size) self.rms_norm_eps = float(rms_norm_eps) self.rope_theta = float(rope_theta) self.rope_scaling = rope_scaling self.dtype = str(dtype) self.intermediate_size = int(intermediate_size) self.num_heads = int(num_heads) self.num_key_value_heads = ( int(num_key_value_heads) if num_key_value_heads is not None else self.num_heads ) self.head_dim = int(head_dim) self.group_size = list(group_size) if isinstance(group_size, tuple) else group_size self.group_size_options = list(_normalize_group_sizes(group_size)) self.layer_pattern = layer_pattern self.use_sengram = bool(use_sengram) self.sengram_bucket_size = int(sengram_bucket_size or 4096) self.sengram_topk = int(sengram_topk) self.use_dcp = bool(use_dcp) self.dynamic_compression_probe_hidden_size = int(dynamic_compression_probe_hidden_size) if self.dynamic_compression_probe_hidden_size < 1: raise ValueError("dynamic_compression_probe_hidden_size must be positive") self.num_hidden_layers = len(resolve_layer_pattern(self)) def scale_lvl(self, lvl: int = 0): """Apply a published YModel32 architecture preset.""" if lvl == 2: self.hidden_size, self.num_heads, self.num_key_value_heads = 2048, 8, 2 self.head_dim, self.intermediate_size = 256, 6144 self.group_size, self.layer_pattern = 8, "fgggggfgggggfgggggfgggggf" self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 32768, 12 self.use_dcp = False elif lvl == 1: self.hidden_size, self.num_heads, self.num_key_value_heads = 1024, 8, 2 self.head_dim, self.intermediate_size = 256, 3072 self.group_size, self.layer_pattern = 8, "fgggggfgggggfgggggfgggggfgggggf" self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 16384, 12 self.use_dcp = False self.dynamic_compression_probe_hidden_size = 1024 elif lvl == 0: self.hidden_size, self.num_heads, self.num_key_value_heads = 768, 12, 6 self.head_dim, self.intermediate_size = 128, 2048 self.group_size, self.layer_pattern = 8, "fggggfggggfggggfgggf" self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 8192, 12 self.use_dcp = False elif lvl == -1: self.hidden_size, self.num_heads, self.num_key_value_heads = 768, 12, 6 self.head_dim, self.intermediate_size = 128, 2048 self.group_size, self.layer_pattern = list(range(4, 17)), "fggggfggggf" self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 4096, 12 self.dynamic_compression_probe_hidden_size = 768 self.use_dcp = True elif lvl == -2: self.hidden_size, self.num_heads, self.num_key_value_heads = 512, 4, 2 self.head_dim, self.intermediate_size = 128, 1024 self.group_size, self.layer_pattern = 8, "fggfggf" self.use_sengram, self.sengram_bucket_size, self.sengram_topk = True, 1024, 8 self.use_dcp = False else: raise ValueError(f"invalid YModel3.2 scale level: {lvl}") if self.layer_pattern.count("f") < self.MIN_FULL_BLOCKS: raise ValueError("YModel3.2 scale patterns require at least three full blocks") self.group_size_options = list(_normalize_group_sizes(self.group_size)) self.num_hidden_layers = len(resolve_layer_pattern(self)) return self class RMSNorm(nn.RMSNorm): def __init__(self, dim: int, eps: float = 1e-6): super().__init__(dim, eps=eps) def forward(self, x: torch.Tensor) -> torch.Tensor: return super().forward(x).to(x.dtype) def _yarn_linear_ramp(low: float, high: float, dim: int) -> torch.Tensor: if low == high: high += 0.001 return torch.clamp((torch.arange(dim, dtype=torch.float32) - low) / (high - low), 0.0, 1.0) def _yarn_correction_dim(num_rotations: float, dim: int, theta: float, max_position_embeddings: int) -> float: return dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) / (2 * math.log(theta)) def resolve_rope_frequencies( dim: int, theta: float, rope_scaling: Optional[dict] = None, max_position_embeddings: Optional[int] = None, ) -> tuple[torch.Tensor, float]: inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) attention_factor = 1.0 if rope_scaling and str(rope_scaling.get("type", "yarn")).lower() == "yarn": factor = float(rope_scaling.get("factor", 1.0)) if factor > 1.0: original = int(rope_scaling.get("original_max_position_embeddings", max_position_embeddings or 0)) beta_fast = float(rope_scaling.get("beta_fast", 32.0)) beta_slow = float(rope_scaling.get("beta_slow", 1.0)) low = math.floor(_yarn_correction_dim(beta_fast, dim, theta, original)) high = math.ceil(_yarn_correction_dim(beta_slow, dim, theta, original)) ramp = _yarn_linear_ramp(low, high, dim // 2) inv_freq = inv_freq / factor * (1.0 - ramp) + inv_freq * ramp attention_factor = float(rope_scaling.get("attention_factor", 1.0)) return inv_freq, attention_factor def precompute_freqs_cis( dim: int, end: int, theta: float, rope_scaling: Optional[dict] = None, ) -> tuple[torch.Tensor, torch.Tensor]: inv_freq, attention_factor = resolve_rope_frequencies( dim, theta, rope_scaling, max_position_embeddings=end ) freqs = torch.outer(torch.arange(end, dtype=torch.float32), inv_freq).float() return ( torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor, torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor, ) def rope_embeddings_from_positions( positions: torch.Tensor, dim: int, theta: float, rope_scaling: Optional[dict] = None, max_position_embeddings: Optional[int] = None, ) -> tuple[torch.Tensor, torch.Tensor]: inv_freq, attention_factor = resolve_rope_frequencies( dim, theta, rope_scaling, max_position_embeddings ) positions = positions.to(dtype=torch.float32) freqs = torch.outer(positions.reshape(-1), inv_freq.to(positions.device)).float() return ( torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor, torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor, ) def rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x[..., x.shape[-1] // 2 :], x[..., : x.shape[-1] // 2]), dim=-1) def apply_rope_to_single(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: if cos.dim() == 2: cos, sin = cos.unsqueeze(0).unsqueeze(0), sin.unsqueeze(0).unsqueeze(0) elif cos.dim() == 3: cos, sin = cos.unsqueeze(1), sin.unsqueeze(1) return (x * cos) + (rotate_half(x) * sin) class GatedAttention(nn.Module): """YModel32 gated grouped-query attention with Q/K RMS normalization.""" def __init__(self, config: YConfig32, layer_id: int): super().__init__() del layer_id self.hidden_size = config.hidden_size self.num_heads = int(config.num_heads) self.num_kv_heads = int(config.num_key_value_heads or self.num_heads) if self.num_heads % self.num_kv_heads: raise ValueError("num_heads must be divisible by num_key_value_heads") self.num_kv_groups = self.num_heads // self.num_kv_heads self.dropout = config.dropout self.head_dim = int(config.head_dim) self.softmax_scale = self.head_dim**-0.5 self.out_dim = self.num_heads * self.head_dim self.q_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False) self.gate_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False) self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False) self.o_proj = nn.Linear(self.out_dim, self.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) def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): q_key, gate_key = prefix + "q_proj.weight", prefix + "gate_proj.weight" legacy_qg = state_dict.get(q_key) if legacy_qg is not None and legacy_qg.shape[0] == self.out_dim * 2 and gate_key not in state_dict: legacy_qg = legacy_qg.reshape(self.num_heads, 2, self.head_dim, self.hidden_size) state_dict[q_key] = legacy_qg[:, 0].reshape(self.out_dim, self.hidden_size) state_dict[gate_key] = legacy_qg[:, 1].reshape(self.out_dim, self.hidden_size) super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ) def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor: if self.num_kv_groups == 1: return x bsz, num_kv_heads, seq_len, head_dim = x.shape x = x[:, :, None, :, :].expand(bsz, num_kv_heads, self.num_kv_groups, seq_len, head_dim) return x.reshape(bsz, self.num_heads, seq_len, head_dim) def _attention_mask(self, attention_mask: Optional[torch.Tensor], bsz: int, seq_len: int, total_len: int): if attention_mask is None: return None if attention_mask.shape[-1] != total_len: attention_mask = attention_mask[..., -total_len:] return attention_mask.reshape(bsz, 1, 1, total_len).bool().expand(bsz, self.num_heads, seq_len, total_len) def forward( self, x: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], past_key_values: Optional[tuple[torch.Tensor, torch.Tensor]] = None, attention_mask: Optional[torch.Tensor] = None, use_cache: bool = False, **kwargs, ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]: del kwargs bsz, seq_len, _ = x.shape cos, sin = position_embeddings if cos.dim() == 2: cos, sin = cos[:seq_len, : self.head_dim], sin[:seq_len, : self.head_dim] else: cos, sin = cos[:, :seq_len, : self.head_dim], sin[:, :seq_len, : self.head_dim] q = self.q_proj(x).reshape(bsz, seq_len, self.num_heads, self.head_dim) gate = self.gate_proj(x) q = self.q_norm(q).permute(0, 2, 1, 3) k = self.k_norm( self.k_proj(x).reshape(bsz, seq_len, self.num_kv_heads, self.head_dim) ).permute(0, 2, 1, 3) v = self.v_proj(x).reshape(bsz, seq_len, self.num_kv_heads, self.head_dim).permute(0, 2, 1, 3) q = apply_rope_to_single(q, cos, sin).contiguous() k = apply_rope_to_single(k, cos, sin).contiguous() v = v.contiguous() if past_key_values is not None: past_k, past_v = past_key_values k, v = torch.cat([past_k, k], dim=2), torch.cat([past_v, v], dim=2) new_past = (k, v) if use_cache else None total_len = k.shape[2] k, v = self._repeat_kv(k).contiguous(), self._repeat_kv(v).contiguous() attn_mask = self._attention_mask(attention_mask, bsz, seq_len, total_len) is_causal = attention_mask is None and seq_len == total_len if q.dtype != v.dtype: q = q.to(dtype=v.dtype) if k.dtype != v.dtype: k = k.to(dtype=v.dtype) out = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=is_causal, scale=self.softmax_scale, ) out = out.permute(0, 2, 1, 3).reshape(bsz, seq_len, self.out_dim) out = self.o_proj(out * torch.sigmoid(gate) * 2.0) return F.dropout(out, p=self.dropout, training=self.training), new_past def resolve_layer_pattern(config: YConfig32) -> str: pattern = getattr(config, "layer_pattern", None) if pattern is None: pattern = "f" + "g" * max(0, int(config.num_hidden_layers) - 1) pattern = "".join(ch for ch in str(pattern).lower() if ch in {"f", "g"}) if not pattern or pattern[0] != "f": raise ValueError(f"YModel3.2 layer_pattern must start with a full layer, got {pattern!r}") return pattern class SwiGLU(nn.Module): def __init__(self, config: YConfig32, intermediate_size: Optional[int] = None): super().__init__() inter = intermediate_size or config.intermediate_size self.up_proj = nn.Linear(config.hidden_size, inter, bias=False) self.gate_proj = nn.Linear(config.hidden_size, inter, bias=False) self.down_proj = nn.Linear(inter, config.hidden_size, bias=False) self.dropout = config.dropout def forward(self, x: torch.Tensor) -> torch.Tensor: up, gate = self.up_proj(x), self.gate_proj(x) return self.down_proj(F.dropout(F.silu(gate) * up, p=self.dropout, training=self.training)) class SEBlock(nn.Module): def __init__(self, dim: int, reduction: int = 16): super().__init__() reduction = max(reduction, dim // reduction) self.se = nn.Sequential( nn.Linear(dim, reduction, bias=False), nn.SiLU(), nn.Linear(reduction, dim, bias=False), nn.Sigmoid(), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return x * (2.0 * self.se(x)) class SengramIndexer(nn.Module): """Inference routing only: choose bucket top-k and normalized scores.""" def __init__(self, config: YConfig32): super().__init__() self.hidden_size = int(config.hidden_size) self.bucket_size = int(config.sengram_bucket_size or 4096) self.topk = max(1, min(int(config.sengram_topk), self.bucket_size)) self.bucket_proj = nn.Linear(self.hidden_size, self.bucket_size, bias=False) def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: router_input = hidden_states.float() router_input = router_input * torch.rsqrt(router_input.square().mean(dim=-1, keepdim=True) + 1e-6) bucket_logits = self.bucket_proj(router_input.to(dtype=hidden_states.dtype)) scores = torch.sigmoid(bucket_logits.float()) _, topk_ids = torch.topk(scores, k=self.topk, dim=-1, sorted=False) topk_scores = scores.gather(-1, topk_ids) topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True).clamp_min(1e-20) return topk_ids, topk_scores.to(bucket_logits.dtype) class SengramPLE(nn.Module): def __init__(self, config: YConfig32): super().__init__() self.hidden_size = int(config.hidden_size) self.embedding = nn.Embedding(int(config.sengram_bucket_size or 4096), self.hidden_size) def forward(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_scores: torch.Tensor) -> torch.Tensor: topk_embed = self.embedding(topk_ids) return (topk_embed * topk_scores.unsqueeze(-1)).sum(dim=-2) class DynamicCompressionProbe(nn.Module): """Inference probe used only when DCP controls token-to-group assignment.""" def __init__(self, config: YConfig32): super().__init__() hidden_size = int(config.hidden_size) probe_size = int(config.dynamic_compression_probe_hidden_size) self.norm = RMSNorm(hidden_size, config.rms_norm_eps) self.gate_proj = nn.Linear(hidden_size, probe_size, bias=False) self.up_proj = nn.Linear(hidden_size, probe_size, bias=False) self.down_proj = nn.Linear(probe_size, 1, bias=True) def forward(self, hidden: torch.Tensor) -> torch.Tensor: hidden = self.norm(hidden.detach()) hidden = F.silu(self.gate_proj(hidden)) * self.up_proj(hidden) return self.down_proj(hidden).squeeze(-1) class G2F(nn.Module): """Add shifted group memory directly to the full-token stream.""" def __init__(self, config: YConfig32): super().__init__() del config def forward(self, hidden: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: return hidden + memory class F2G(nn.Module): """Compress full-token states into groups, optionally DCP-weighted.""" def __init__(self, config: YConfig32): super().__init__() self.hidden_size = int(config.hidden_size) self.group_size = _normalize_group_sizes(config.group_size)[0] def _aligned_weights(self, importance_logits, tokens: int, device: torch.device): if importance_logits is None: return None logits = importance_logits.to(device=device) if logits.ndim != 2: raise ValueError("DCP importance logits must have shape [batch, tokens]") required = max(0, tokens - 1) logits = logits[:, :required] if logits.shape[1] < required: logits = F.pad(logits, (0, required - logits.shape[1]), value=0.0) return torch.sigmoid(logits.float()) def _weighted_reduce(self, h, attention_mask, group_ids, num_groups, importance_logits=None): bsz, tokens, hidden = h.shape if group_ids.shape != (bsz, tokens): raise ValueError(f"group_ids shape {tuple(group_ids.shape)} does not match hidden shape {(bsz, tokens)}") if attention_mask.shape != (bsz, tokens): raise ValueError(f"attention_mask shape {tuple(attention_mask.shape)} does not match hidden shape {(bsz, tokens)}") rest_h, rest_ids = h[:, 1:], group_ids[:, 1:] rest_valid = attention_mask[:, 1:].bool().reshape(-1) offsets = torch.arange(bsz, device=h.device, dtype=torch.int64).unsqueeze(1) * int(num_groups) flat_ids = (rest_ids.to(torch.int64) + offsets).reshape(-1) flat_h = rest_h.reshape(-1, hidden) aligned_weights = self._aligned_weights(importance_logits, tokens, h.device) if aligned_weights is None: flat_weights = torch.ones((bsz, max(0, tokens - 1)), device=h.device, dtype=torch.float32).reshape(-1) else: flat_weights = aligned_weights.reshape(-1) flat_ids, flat_h, flat_weights = flat_ids[rest_valid], flat_h[rest_valid], flat_weights[rest_valid] weighted = h.new_zeros((bsz * int(num_groups), hidden)) weighted.index_add_(0, flat_ids, flat_h * flat_weights.to(h.dtype).unsqueeze(-1)) denominators = torch.zeros((bsz * int(num_groups), 1), device=h.device, dtype=torch.float32) denominators.index_add_(0, flat_ids, flat_weights.unsqueeze(-1)) compressed = weighted / denominators.clamp_min(1e-12).to(h.dtype) compressed = torch.where(denominators.gt(0.0), compressed, torch.zeros_like(compressed)) compressed = compressed.reshape(bsz, int(num_groups), hidden) first = h[:, :1] * attention_mask[:, :1].to(h.dtype).unsqueeze(-1) return torch.cat([first, compressed[:, 1:]], dim=1) def forward( self, h: torch.Tensor, attention_mask: torch.Tensor, group_ids: Optional[torch.Tensor] = None, num_groups: Optional[int] = None, group_size: Optional[int] = None, importance_logits: Optional[torch.Tensor] = None, ) -> torch.Tensor: if group_ids is not None: if num_groups is None: raise ValueError("indexed F2G requires num_groups") return self.forward_indexed(h, attention_mask, group_ids, num_groups, importance_logits) group_size = self.group_size if group_size is None else int(group_size) bsz, padded_len, _ = h.shape rest_len = max(0, padded_len - 1) num_groups = 1 + (rest_len + group_size - 1) // group_size positions = torch.arange(padded_len, device=h.device, dtype=torch.int64) fixed_ids = torch.zeros((bsz, padded_len), device=h.device, dtype=torch.int64) if padded_len > 1: fixed_ids[:, 1:] = (positions[1:] - 1) // group_size + 1 return self._weighted_reduce(h, attention_mask, fixed_ids, num_groups, importance_logits) def forward_indexed(self, h, attention_mask, group_ids, num_groups, importance_logits=None): return self._weighted_reduce(h, attention_mask, group_ids, int(num_groups), importance_logits) class FullBlock(nn.Module): def __init__(self, config: YConfig32, layer_id: int, has_group_context: bool): super().__init__() self.has_group_context = has_group_context self.group_size = _normalize_group_sizes(config.group_size)[0] self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.g2f = G2F(config) if has_group_context else None self.self_attn = GatedAttention(config, layer_id) self.mlp = SwiGLU(config) @property def attn(self): return self.self_attn @property def ffn(self): return self.mlp def _add_prev_group_residual(self, h, g, group_attention_mask=None, token_group_ids=None, group_size=None): if token_group_ids is not None: previous = (token_group_ids - 1).clamp_min(0) memory = g.gather(1, previous.unsqueeze(-1).expand(-1, -1, g.shape[-1])) memory = memory * token_group_ids.gt(0).unsqueeze(-1).to(memory.dtype) return self.g2f(h, memory) group_size = self.group_size if group_size is None else int(group_size) bsz, padded_len, hidden = h.shape rest_len = max(0, padded_len - 1) rest_groups = (rest_len + group_size - 1) // group_size groups = 1 + rest_groups if g.shape[1] != groups or rest_groups == 0: return h grouped_h = h[:, 1:].reshape(bsz, rest_groups, group_size, hidden) memory = g[:, :-1, :].unsqueeze(2) if group_attention_mask is not None: memory = memory * group_attention_mask[:, :-1].to(g.dtype).reshape(bsz, rest_groups, 1, 1) fused = self.g2f(grouped_h, memory).reshape(bsz, rest_groups * group_size, hidden) return torch.cat([h[:, :1], fused[:, :rest_len]], dim=1) def forward(self, h, g, token_position_embeddings, group_attention_mask=None, token_group_ids=None, group_size=None): if self.has_group_context and g is not None: h = self._add_prev_group_residual(h, g, group_attention_mask, token_group_ids, group_size) attn_out, _ = self.self_attn(self.input_layernorm(h), token_position_embeddings) h = h + attn_out return h + self.mlp(self.post_attention_layernorm(h)) class GroupBlock(nn.Module): def __init__(self, config: YConfig32, layer_id: int): super().__init__() self.use_sengram = bool(config.use_sengram) self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps) self.sengram_ple = SengramPLE(config) if self.use_sengram else None self.self_attn = GatedAttention(config, layer_id) self.mlp = SwiGLU(config) self.se1 = SEBlock(config.hidden_size) self.se2 = SEBlock(config.hidden_size) @property def attn(self): return self.self_attn @property def ffn(self): return self.mlp def forward(self, g, group_routes, group_position_embeddings): if self.use_sengram and group_routes is not None and self.sengram_ple is not None: g = g + self.sengram_ple(g, group_routes[0], group_routes[1]) g0 = self.se1(self.input_layernorm(g)) attn_out, _ = self.self_attn(g0, group_position_embeddings) g = g + attn_out return g + self.mlp(self.se2(self.post_attention_layernorm(g))) class YModel32(nn.Module): """YModel3.2 backbone.""" def __init__(self, config: YConfig32): super().__init__() self.config = config self.group_size_options = _normalize_group_sizes(config.group_size) self.group_size = self.group_size_options[0] self.max_group_size = max(self.group_size_options) self.active_group_size = self.group_size self.layer_pattern = resolve_layer_pattern(config) self.num_layers = len(self.layer_pattern) self.dropout = config.dropout self.use_sengram = bool(config.use_sengram) self.use_dcp = bool(config.use_dcp) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.sengram_indexer = SengramIndexer(config) if self.use_sengram else None # Kept in the state dict even when DCP is disabled so checkpoint # loading never drops learned parameters. self.dynamic_compression_probe = DynamicCompressionProbe(config) f2g_count = 0 in_g = False for kind in self.layer_pattern: if kind == "g" and not in_g: f2g_count += 1 in_g = True elif kind == "f": in_g = False self.f2g_layers = nn.ModuleList([F2G(config) for _ in range(f2g_count)]) self.f2g_for_full = [] f2g_ptr = 0 for layer_pos, kind in enumerate(self.layer_pattern): if kind != "f": continue has_group_run = layer_pos + 1 < len(self.layer_pattern) and self.layer_pattern[layer_pos + 1] == "g" self.f2g_for_full.append(f2g_ptr if has_group_run else -1) if has_group_run: f2g_ptr += 1 layers = [] seen_full = 0 for layer_id, kind in enumerate(self.layer_pattern): if kind == "f": layers.append(FullBlock(config, layer_id, has_group_context=seen_full > 0)) seen_full += 1 else: layers.append(GroupBlock(config, layer_id)) self.layers = nn.ModuleList(layers) self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) freqs_cos, freqs_sin = precompute_freqs_cis( config.head_dim, config.max_position_embeddings + self.max_group_size, config.rope_theta, config.rope_scaling, ) self.register_buffer("freqs_cos", freqs_cos, persistent=False) self.register_buffer("freqs_sin", freqs_sin, persistent=False) @property def sengram(self): return self.sengram_indexer def _sample_group_size(self) -> int: # Deterministic at inference: uses the first configured size. return self.group_size_options[0] def _pad_to_groups(self, input_ids, attention_mask, group_size, pad_to_groups=True): bsz, seq_len = input_ids.shape rest_len = max(0, seq_len - 1) pad_len = (-rest_len) % group_size if pad_to_groups else 0 if pad_len: pad_id = self.config.pad_token_id if self.config.pad_token_id < self.config.vocab_size else 0 input_ids = torch.cat([input_ids, input_ids.new_full((bsz, pad_len), pad_id)], dim=1) if attention_mask is not None: attention_mask = torch.cat([attention_mask, attention_mask.new_zeros((bsz, pad_len))], dim=1) if attention_mask is None: pad_id = self.config.pad_token_id if self.config.pad_token_id < self.config.vocab_size else -1 attention_mask = input_ids.ne(pad_id).long() return input_ids, attention_mask, seq_len def _group_count(self, padded_len: int, group_size: int) -> int: return 1 + max(0, padded_len - 1) // group_size def _group_attention_mask(self, attention_mask, group_size): bsz, padded_len = attention_mask.shape rest_len = max(0, padded_len - 1) rest_groups = (rest_len + group_size - 1) // group_size first = attention_mask[:, :1] if rest_groups == 0: return first.long() rest_mask = attention_mask[:, 1:] pad_len = rest_groups * group_size - rest_len if pad_len: rest_mask = torch.cat([rest_mask, rest_mask.new_zeros((bsz, pad_len))], dim=1) return torch.cat([first, rest_mask.reshape(bsz, rest_groups, group_size).any(dim=-1).long()], dim=1) def _group_positions(self, groups: int, device: torch.device, group_size: int): if groups <= 1: return torch.zeros((groups,), device=device, dtype=torch.float32) rest = 1.0 + (torch.arange(groups - 1, device=device, dtype=torch.float32) + 0.5) * group_size return torch.cat([rest.new_zeros((1,)), rest], dim=0) def _dcp_group_ids(self, probe_logits, attention_mask, group_size=None): bsz, tokens = probe_logits.shape group_ids = torch.zeros((bsz, tokens), device=probe_logits.device, dtype=torch.int64) if tokens <= 1: return group_ids, 1 group_size = self.group_size if group_size is None else int(group_size) target_groups = max(1, (tokens + group_size - 1) // group_size) target_rest_groups = target_groups - 1 if target_rest_groups == 0: return group_ids, target_groups score = torch.sigmoid(probe_logits[:, :-1]).float() score = score * attention_mask[:, 1:].to(score.dtype) budget = float(group_size) * 0.5 raw_integral = score.sum(dim=1, keepdim=True) / budget scale = target_rest_groups / raw_integral.clamp_min(1e-6) scaled_score = score * scale integral_before_token = (torch.cumsum(scaled_score, dim=1) - scaled_score) / budget integral_bucket = torch.floor(integral_before_token).to(torch.int64) integral_bucket = integral_bucket.clamp_(0, target_rest_groups - 1) valid_rest = attention_mask[:, 1:].bool() rest_positions = torch.arange(tokens - 1, device=probe_logits.device).unsqueeze(0) last_valid = rest_positions.masked_fill(~valid_rest, -1).amax(dim=1) safe_last_valid = last_valid.clamp_min(0) rows = torch.arange(bsz, device=probe_logits.device) current_last_bucket = integral_bucket[rows, safe_last_valid] final_bucket = torch.where( last_valid.ge(0), torch.full_like(current_last_bucket, target_rest_groups - 1), current_last_bucket, ) integral_bucket[rows, safe_last_valid] = final_bucket group_ids[:, 1:] = integral_bucket + 1 return group_ids, target_groups def _indexed_group_metadata(self, group_ids, attention_mask, num_groups): bsz, tokens = group_ids.shape offsets = torch.arange(bsz, device=group_ids.device, dtype=torch.int64).unsqueeze(1) * num_groups flat_ids = (group_ids + offsets).reshape(-1) valid = attention_mask.reshape(-1).bool() counts = torch.zeros(bsz * num_groups, device=group_ids.device, dtype=torch.float32) counts.index_add_(0, flat_ids[valid], torch.ones_like(flat_ids[valid], dtype=torch.float32)) positions = torch.arange(tokens, device=group_ids.device, dtype=torch.float32).expand(bsz, -1).reshape(-1) position_sums = torch.zeros_like(counts) position_sums.index_add_(0, flat_ids[valid], positions[valid]) group_mask = counts.reshape(bsz, num_groups).gt(0) centers = (position_sums / counts.clamp_min(1.0)).reshape(bsz, num_groups) return group_mask.long(), centers def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, **kwargs, ) -> torch.Tensor: del kwargs if input_ids is None: raise ValueError("input_ids is required") group_size = self._sample_group_size() self.active_group_size = group_size input_ids, attention_mask, orig_len = self._pad_to_groups( input_ids, attention_mask, group_size, pad_to_groups=not self.use_dcp ) bsz, padded_len = input_ids.shape h = F.dropout(self.embed_tokens(input_ids), p=self.dropout, training=False) if position_ids is None: position_ids = torch.arange(padded_len, device=input_ids.device) token_position_embeddings = ( self.freqs_cos[position_ids[:padded_len]].to(h.device), self.freqs_sin[position_ids[:padded_len]].to(h.device), ) groups = self._group_count(padded_len, group_size) group_attention_mask = self._group_attention_mask(attention_mask, group_size) group_center_positions = self._group_positions(groups, input_ids.device, group_size) group_position_embeddings = rope_embeddings_from_positions( group_center_positions, self.config.head_dim, self.config.rope_theta, self.config.rope_scaling, self.config.max_position_embeddings, ) group_position_embeddings = (group_position_embeddings[0].to(h.device), group_position_embeddings[1].to(h.device)) g = None pending_group_update = None group_routes = None token_group_ids = None dcp_num_groups = None full_idx = 0 difficulty_logits = None for layer_pos, kind in enumerate(self.layer_pattern): if kind == "f": h = self.layers[layer_pos]( h, g, token_position_embeddings, group_attention_mask=group_attention_mask, token_group_ids=token_group_ids, group_size=group_size, ) if difficulty_logits is None and self.use_dcp: difficulty_logits = self.dynamic_compression_probe(h) token_group_ids, dcp_num_groups = self._dcp_group_ids( difficulty_logits, attention_mask, group_size ) group_attention_mask, group_center_positions = self._indexed_group_metadata( token_group_ids, attention_mask, dcp_num_groups ) group_position_embeddings = rope_embeddings_from_positions( group_center_positions, self.config.head_dim, self.config.rope_theta, self.config.rope_scaling, self.config.max_position_embeddings, ) group_position_embeddings = ( group_position_embeddings[0].reshape(bsz, dcp_num_groups, -1).to(h.device), group_position_embeddings[1].reshape(bsz, dcp_num_groups, -1).to(h.device), ) groups = dcp_num_groups f2g_idx = self.f2g_for_full[full_idx] if f2g_idx >= 0: if token_group_ids is not None and dcp_num_groups is not None: pending_group_update = self.f2g_layers[f2g_idx]( h, attention_mask, token_group_ids, dcp_num_groups, importance_logits=difficulty_logits, ) else: pending_group_update = self.f2g_layers[f2g_idx]( h, attention_mask, group_size=group_size, importance_logits=difficulty_logits, ) if self.sengram_indexer is not None and pending_group_update is not None: group_routes = self.sengram_indexer(pending_group_update) full_idx += 1 else: if pending_group_update is not None: g = pending_group_update if g is None else g + pending_group_update pending_group_update = None elif g is None: g = h.new_zeros(h.shape[0], groups, h.shape[-1]) g = self.layers[layer_pos](g, group_routes, group_position_embeddings) return self.norm(h)[:, :orig_len] class YForCausalLM32(PreTrainedModel, GenerationMixin): config_class = YConfig32 base_model_prefix = "model" _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} def __init__(self, config: Optional[YConfig32] = None): config = config or YConfig32() super().__init__(config) self.model = YModel32(config) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.model.embed_tokens.weight = self.lm_head.weight dtype_name = config.dtype if dtype_name not in {"float16", "bfloat16", "float32"}: hf_dtype = getattr(config, "torch_dtype", None) dtype_name = str(hf_dtype).replace("torch.", "") if hf_dtype is not None else "float32" dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}.get(dtype_name) if dtype is not None: self.to(dtype) def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs, ) -> CausalLMOutputWithPast: position_ids = kwargs.pop("position_ids", None) # This architecture currently recomputes the full prefix. Cache inputs # are accepted for API compatibility but are not consumed. kwargs.pop("past_key_values", None) kwargs.pop("use_cache", None) kwargs.pop("cache_position", None) hidden = self.model(input_ids, attention_mask, position_ids=position_ids) indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep output = CausalLMOutputWithPast( logits=self.lm_head(hidden[:, indices, :]), past_key_values=None, ) output["last_hidden_state"] = hidden return output def generate( self, inputs=None, attention_mask=None, max_new_tokens=8192, temperature=0.85, top_p=0.85, top_k=50, eos_token_id=None, streamer=None, use_cache=False, num_return_sequences=1, do_sample=True, repetition_penalty=1.0, **kwargs, ): del use_cache if inputs is None: inputs = kwargs.pop("input_ids", None) if inputs is None: raise ValueError("inputs or input_ids must be provided") input_ids = inputs.repeat(num_return_sequences, 1) attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None logits_processor = kwargs.get("logits_processor") eos_token_id = self.config.eos_token_id if eos_token_id is None else eos_token_id if streamer is not None: streamer.put(input_ids.cpu()) with torch.inference_mode(): for _ in range(max_new_tokens): outputs = self.forward( input_ids=input_ids, attention_mask=attention_mask, logits_to_keep=1, use_cache=False, ) logits = outputs.logits[:, -1, :] if temperature and temperature > 0: logits = logits / temperature if repetition_penalty != 1.0: for i in range(input_ids.shape[0]): seen = torch.unique(input_ids[i]) logits[i, seen] = torch.where( logits[i, seen] < 0, logits[i, seen] * repetition_penalty, logits[i, seen] / repetition_penalty, ) if logits_processor is not None: logits = logits_processor(input_ids, logits) if top_k > 0: k = min(int(top_k), logits.shape[-1]) cutoff = torch.topk(logits, k, dim=-1)[0][..., -1, None] logits = logits.masked_fill(logits < cutoff, -float("inf")) if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True) sorted_mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p sorted_mask[..., 1:] = sorted_mask[..., :-1].clone() sorted_mask[..., 0] = False remove_mask = torch.zeros_like(sorted_mask).scatter(1, sorted_indices, sorted_mask) logits = logits.masked_fill(remove_mask, -float("inf")) next_token = ( torch.multinomial(torch.softmax(logits, dim=-1), 1) if do_sample else torch.argmax(logits, dim=-1, keepdim=True) ) input_ids = torch.cat([input_ids, next_token], dim=-1) if attention_mask is not None: attention_mask = torch.cat( [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 ) if streamer is not None: streamer.put(next_token.cpu()) if eos_token_id is not None and bool((next_token == eos_token_id).all()): break if streamer is not None: streamer.end() return input_ids def count_parameters(config: YConfig32) -> int: return sum(parameter.numel() for parameter in YForCausalLM32(config).parameters()) def _state_dict_from_checkpoint(payload) -> dict[str, torch.Tensor]: if not isinstance(payload, dict): state_dict = getattr(payload, "state_dict", None) if callable(state_dict): payload = state_dict() else: raise TypeError("checkpoint must contain a PyTorch state-dict mapping") for key in ("state_dict", "model_state_dict", "model", "student", "student_state_dict"): nested = payload.get(key) if isinstance(nested, dict) and nested and all(isinstance(k, str) for k in nested): if any(torch.is_tensor(value) for value in nested.values()): payload = nested break if not payload or not all(isinstance(key, str) for key in payload): raise TypeError("checkpoint does not look like a PyTorch state dict") return payload def _load_state_dict(path: Union[str, Path]) -> dict[str, torch.Tensor]: path = Path(path) if path.is_dir(): candidates = ("sft.pt", "distill.pt", "model.pt", "pytorch_model.bin") path = next((path / name for name in candidates if (path / name).is_file()), None) if path is None: raise FileNotFoundError("no .pt/.bin PyTorch checkpoint found in the directory") if path.suffix.lower() == ".safetensors": raise ValueError(".safetensors loading is not included; export a PyTorch .pt state dict") try: payload = torch.load(path, map_location="cpu", weights_only=True) except TypeError: payload = torch.load(path, map_location="cpu") return _state_dict_from_checkpoint(payload) def _key_variants(key: str) -> tuple[str, ...]: variants = [key] current = key while current.startswith("module.") or current.startswith("_orig_mod."): current = current.split(".", 1)[1] variants.append(current) for source in tuple(variants): aliases = ( source.replace(".attn.", ".self_attn."), source.replace(".ffn.", ".mlp."), source.replace(".ffn.dense.", ".mlp."), source.replace("model.sengram.", "model.sengram_indexer."), ) variants.extend(aliases) # Some compiled checkpoints insert _orig_mod below the outer model prefix. variants.extend( key.replace("model._orig_mod.", "model.").replace("._orig_mod.", ".") for key in tuple(variants) ) return tuple(dict.fromkeys(variants)) def _load_compatible_weights(model: YForCausalLM32, source: dict[str, torch.Tensor], strict: bool = True): target = model.state_dict() normalized = {} for key, value in source.items(): if not torch.is_tensor(value): continue for variant in _key_variants(key): normalized.setdefault(variant, value) matched = {} shape_mismatches = [] for key, tensor in target.items(): found = next((normalized[name] for name in _key_variants(key) if name in normalized), None) if found is None: continue if tuple(found.shape) != tuple(tensor.shape): # Early checkpoints packed Q and gate projections into one # tensor; split them in the same head-major layout used by the # checkpoint loader. if ( key.endswith(".q_proj.weight") and found.ndim == 2 and found.shape[0] == tensor.shape[0] * 2 and found.shape[1] == tensor.shape[1] ): packed = found.reshape( model.config.num_heads, 2, model.config.head_dim, model.config.hidden_size, ) matched[key] = packed[:, 0].reshape_as(tensor) gate_key = key[: -len("q_proj.weight")] + "gate_proj.weight" matched[gate_key] = packed[:, 1].reshape_as(target[gate_key]) continue shape_mismatches.append((key, tuple(found.shape), tuple(tensor.shape))) continue matched[key] = found result = model.load_state_dict(matched, strict=False) missing = list(result.missing_keys) if model.config.tie_word_embeddings: tied = {"model.embed_tokens.weight", "lm_head.weight"} if tied.intersection(matched): missing = [key for key in missing if key not in tied] if shape_mismatches and strict: preview = "; ".join(f"{key}: checkpoint{src} != model{dst}" for key, src, dst in shape_mismatches[:8]) raise RuntimeError(f"checkpoint tensor shape mismatch: {preview}") if missing and strict: preview = ", ".join(missing[:20]) raise RuntimeError(f"checkpoint is missing {len(missing)} model tensors: {preview}") unique_target = {} for key, tensor in target.items(): identity = (tensor.device.type, tensor.data_ptr(), tuple(tensor.shape)) unique_target.setdefault(identity, []).append(key) loaded_unique = sum(any(key in matched for key in keys) for keys in unique_target.values()) print(f"loaded: {loaded_unique}/{len(unique_target)} unique model tensors") if shape_mismatches: print(f"shape-mismatched tensors ignored: {len(shape_mismatches)}") if missing: print(f"missing model tensors: {len(missing)}") return model def load_ymodel32_eval( path: Union[str, Path], config: Optional[YConfig32] = None, strict: bool = True, device: Union[str, torch.device] = "cpu", ) -> YForCausalLM32: """Load a standalone evaluator from a .pt/.bin state dict and config JSON.""" path = Path(path) if config is None: config_candidates = ( path / "config.json" if path.is_dir() else path.with_name("config.json"), path.parent.parent / "config.json" if not path.is_dir() else path.parent / "config.json", ) config_path = next((candidate for candidate in config_candidates if candidate.is_file()), None) if config_path is None: raise FileNotFoundError("pass config=YConfig32(...) or place config.json next to the checkpoint") config = YConfig32.from_json_file(str(config_path)) model = YForCausalLM32(config) model = _load_compatible_weights(model, _load_state_dict(path), strict=strict) return model.to(device).eval() YModel32Eval = YModel32 YForCausalLM32Eval = YForCausalLM32 __all__ = [ "DynamicCompressionProbe", "F2G", "FullBlock", "G2F", "GatedAttention", "GroupBlock", "RMSNorm", "SengramIndexer", "SengramPLE", "SwiGLU", "YConfig32", "YForCausalLM32", "YForCausalLM32Eval", "YModel32", "YModel32Eval", "apply_rope_to_single", "count_parameters", "load_ymodel32_eval", "precompute_freqs_cis", "resolve_layer_pattern", "resolve_rope_frequencies", "rope_embeddings_from_positions", ]