| """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, |
| ): |
| |
| 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 |
| |
| |
| 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: |
| |
| 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) |
| |
| |
| 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) |
| |
| 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): |
| |
| |
| |
| 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", |
| ] |
|
|
