| """Standalone evaluation/inference implementation for ymodel31. |
| |
| This file intentionally contains a self-contained inference path so exported |
| checkpoints can be loaded without importing the training implementation. |
| Training-only features such as gradient checkpointing and self-distillation are |
| omitted here on purpose. |
| """ |
|
|
| 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 safetensors.torch import load_file as load_safetensors |
| from transformers import GenerationMixin, PreTrainedModel |
| from transformers.activations import ACT2FN |
| from transformers.configuration_utils import PretrainedConfig |
| from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
|
|
| def normalize_gradient_checkpointing_level(value: Union[bool, int, str, None]) -> int: |
| if isinstance(value, bool): |
| return 1 if value else 0 |
| if value is None: |
| return 0 |
| if isinstance(value, int): |
| return max(0, value) |
| text = str(value).strip().lower() |
| if text in {"", "false", "off", "no", "none"}: |
| return 0 |
| if text in {"true", "on", "yes"}: |
| return 1 |
| try: |
| return max(0, int(text)) |
| except ValueError as exc: |
| raise ValueError(f"Unsupported gradient_checkpointing level: {value!r}") from exc |
|
|
|
|
| class YConfig31(PretrainedConfig): |
| model_type = "ynet31" |
|
|
| def __init__( |
| self, |
| dropout: float = 0.0, |
| bos_token_id: int = 151644, |
| eos_token_id: int = 151645, |
| pad_token_id: int = 151643, |
| hidden_act: str = "silu", |
| hidden_size: int = 768, |
| num_hidden_layers: int = 8, |
| 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", |
| self_distill: bool = True, |
| intermediate_size: int = 1536, |
| num_heads: int = 12, |
| mla_kv_lora_rank: int = 64, |
| mla_qk_nope_head_dim: int = 64, |
| mla_qk_rope_head_dim: int = 32, |
| mla_attn_impl: str = "absorb", |
| qkv_lora: bool = False, |
| gradient_checkpointing: Union[bool, int, str] = 0, |
| use_sengram: bool = True, |
| sengram_bucket_size: Optional[int] = 4096, |
| sengram_topk: int = 2, |
| engram_bucket_size: Optional[int] = None, |
| engram_topk: Optional[int] = None, |
| **kwargs, |
| ): |
| super().__init__( |
| bos_token_id=bos_token_id, |
| eos_token_id=eos_token_id, |
| pad_token_id=pad_token_id, |
| **kwargs, |
| ) |
| self.dropout = dropout |
| self.hidden_act = hidden_act |
| self.hidden_size = hidden_size |
| self.num_hidden_layers = num_hidden_layers |
| self.max_position_embeddings = max_position_embeddings |
| self.vocab_size = vocab_size |
| self.rms_norm_eps = rms_norm_eps |
| self.rope_theta = rope_theta |
| self.rope_scaling = rope_scaling |
| self.dtype = dtype |
| self.self_distill = self_distill |
| self.intermediate_size = intermediate_size |
| self.num_heads = num_heads |
| self.mla_kv_lora_rank = mla_kv_lora_rank |
| self.mla_qk_nope_head_dim = mla_qk_nope_head_dim |
| self.mla_qk_rope_head_dim = mla_qk_rope_head_dim |
| self.mla_attn_impl = mla_attn_impl |
| self.qkv_lora = qkv_lora |
| self.gradient_checkpointing = normalize_gradient_checkpointing_level(gradient_checkpointing) |
| self.use_sengram = bool(use_sengram) |
| if engram_bucket_size is not None: |
| sengram_bucket_size = engram_bucket_size |
| if engram_topk is not None: |
| sengram_topk = engram_topk |
| self.sengram_bucket_size = sengram_bucket_size |
| self.sengram_topk = sengram_topk |
| self.engram_bucket_size = self.sengram_bucket_size |
| self.engram_topk = self.sengram_topk |
|
|
| @property |
| def head_dim(self) -> int: |
| return self.mla_qk_nope_head_dim + self.mla_qk_rope_head_dim |
|
|
| @property |
| def qk_head_dim(self) -> int: |
| return self.head_dim |
|
|
| def scale_lvl(self, lvl: int = 0): |
| if lvl == 0: |
| self.hidden_size = 768 |
| self.num_hidden_layers = 12 |
| self.num_heads = 8 |
| self.mla_kv_lora_rank = 256 |
| self.mla_qk_nope_head_dim = 128 |
| self.mla_qk_rope_head_dim = 64 |
| self.intermediate_size = 2048 |
| self.use_sengram = True |
| self.sengram_bucket_size = 8192 |
| self.sengram_topk = 8 |
| elif lvl == -1: |
| self.hidden_size = 768 |
| self.num_hidden_layers = 8 |
| self.num_heads = 6 |
| self.mla_kv_lora_rank = 128 |
| self.mla_qk_nope_head_dim = 64 |
| self.mla_qk_rope_head_dim = 64 |
| self.intermediate_size = 1536 |
| self.use_sengram = True |
| elif lvl == -2: |
| self.hidden_size = 512 |
| self.num_hidden_layers = 4 |
| self.num_heads = 4 |
| self.mla_kv_lora_rank = 128 |
| self.mla_qk_nope_head_dim = 64 |
| self.mla_qk_rope_head_dim = 64 |
| self.intermediate_size = 1024 |
| self.use_sengram = True |
| else: |
| raise ValueError(f"invalid ymodel31 scale level: {lvl}") |
| return self |
|
|
|
|
| def _yarn_linear_ramp(low: float, high: float, dim: int) -> torch.Tensor: |
| if low == high: |
| high += 0.001 |
| linear = (torch.arange(dim, dtype=torch.float32) - low) / (high - low) |
| return torch.clamp(linear, 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 precompute_freqs_cis( |
| dim: int, |
| end: int, |
| theta: float, |
| rope_scaling: Optional[dict] = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| freqs = 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", end)) |
| 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) |
| freqs = freqs / factor * (1.0 - ramp) + freqs * ramp |
| attention_factor = float(rope_scaling.get("attention_factor", 1.0)) |
| t = torch.arange(end) |
| freqs = torch.outer(t, freqs).float() |
| freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor |
| freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor |
| return freqs_cos, freqs_sin |
|
|
|
|
| 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 = cos.unsqueeze(0).unsqueeze(0) |
| sin = sin.unsqueeze(0).unsqueeze(0) |
| elif cos.dim() == 3: |
| cos = cos.unsqueeze(1) |
| sin = sin.unsqueeze(1) |
| return (x * cos) + (rotate_half(x) * sin) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) |
| self.eps = eps |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| out = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) |
| return (out * self.weight.float()).to(x.dtype) |
|
|
|
|
| class SEBlock(nn.Module): |
| def __init__(self, dim: int, reduction: int = 16, act: Optional[nn.Module] = None): |
| super().__init__() |
| reduction = max(reduction, dim // reduction) |
| self.se = nn.Sequential( |
| nn.Linear(dim, reduction, bias=False), |
| act or nn.SiLU(), |
| nn.Linear(reduction, dim, bias=False), |
| nn.Sigmoid(), |
| ) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return x * self.se(x) |
|
|
|
|
| class MLGA(nn.Module): |
| """Multihead Latent Gated Attention""" |
|
|
| def __init__(self, config: YConfig31, layer_id: int): |
| super().__init__() |
| self.layer_id = layer_id |
| self.hidden_size = config.hidden_size |
| self.num_heads = config.num_heads |
| self.dropout = config.dropout |
| self.kv_lora_rank = config.mla_kv_lora_rank |
| self.qk_nope_head_dim = config.mla_qk_nope_head_dim |
| self.qk_rope_head_dim = config.mla_qk_rope_head_dim |
| self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim |
| self.attn_impl = config.mla_attn_impl |
| self.softmax_scale = self.qk_head_dim**-0.5 |
| self.out_dim = self.num_heads * self.kv_lora_rank |
|
|
| self.wq = nn.Linear(self.hidden_size, self.num_heads * self.qk_head_dim, bias=False) |
| self.wkv_a = nn.Linear(self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False) |
| self.kv_norm = RMSNorm(self.kv_lora_rank, config.rms_norm_eps) |
| self.wkv_b = nn.Linear(self.kv_lora_rank, self.num_heads * self.qk_nope_head_dim, bias=False) |
| self.z_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False) |
| self.o_proj = nn.Linear(self.out_dim, self.hidden_size, bias=False) |
|
|
| def _project_q(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| bsz, seq_len, _ = x.shape |
| q = self.wq(x).reshape(bsz, seq_len, self.num_heads, self.qk_head_dim) |
| return q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) |
|
|
| def _project_kv(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| raw = self.wkv_a(x) |
| c_kv, k_pe = raw.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) |
| c_kv = self.kv_norm(c_kv) |
| k_pe = apply_rope_to_single(k_pe.unsqueeze(1), cos, sin).permute(0, 2, 1, 3) |
| return c_kv, k_pe |
|
|
| def _explicit_kv(self, c_kv: torch.Tensor, k_pe: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: |
| bsz, seq_len, _ = c_kv.shape |
| k_nope = self.wkv_b(c_kv).reshape(bsz, seq_len, self.num_heads, self.qk_nope_head_dim) |
| k = torch.cat([k_nope, k_pe.expand(-1, -1, self.num_heads, -1)], dim=-1) |
| v = c_kv.unsqueeze(2).expand(-1, -1, self.num_heads, -1) |
| return k, v |
|
|
| 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:] |
| mask = attention_mask.reshape(bsz, 1, 1, total_len).bool() |
| return mask.expand(bsz, self.num_heads, seq_len, total_len) |
|
|
| def _forward_sdpa( |
| self, |
| q_nope: torch.Tensor, |
| q_pe: torch.Tensor, |
| c_kv: torch.Tensor, |
| k_pe: torch.Tensor, |
| z: torch.Tensor, |
| attention_mask: Optional[torch.Tensor], |
| ) -> torch.Tensor: |
| bsz, seq_len, _, _ = q_nope.shape |
| total_len = c_kv.shape[1] |
| k, v = self._explicit_kv(c_kv, k_pe) |
| q = torch.cat([q_nope, q_pe], dim=-1).permute(0, 2, 1, 3) |
| k = k.permute(0, 2, 1, 3) |
| v = v.permute(0, 2, 1, 3) |
| attn_mask = self._attention_mask(attention_mask, bsz, seq_len, total_len) |
| is_causal = attention_mask is None and seq_len == total_len |
| out = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=attn_mask, |
| dropout_p=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 = out * torch.sigmoid(z) |
| return self.o_proj(out) |
|
|
| def _forward_absorb( |
| self, |
| q_nope: torch.Tensor, |
| q_pe: torch.Tensor, |
| c_kv: torch.Tensor, |
| k_pe: torch.Tensor, |
| z: torch.Tensor, |
| attention_mask: Optional[torch.Tensor], |
| ) -> torch.Tensor: |
| bsz, seq_len, _, _ = q_nope.shape |
| total_len = c_kv.shape[1] |
| w = self.wkv_b.weight.reshape(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank) |
| q_nope_c = torch.einsum("bshd,hdc->bshc", q_nope, w) |
| scores = torch.einsum("bshc,btc->bsht", q_nope_c, c_kv) |
| scores = scores + torch.einsum("bshr,btr->bsht", q_pe, k_pe.squeeze(2)) |
| scores = scores * self.softmax_scale |
|
|
| causal = torch.full((seq_len, seq_len), float("-inf"), device=scores.device, dtype=scores.dtype) |
| causal = torch.triu(causal, diagonal=1).reshape(1, seq_len, 1, seq_len) |
| scores = scores + F.pad(causal, (total_len - seq_len, 0), value=0.0) |
| if attention_mask is not None: |
| if attention_mask.shape[-1] != total_len: |
| attention_mask = attention_mask[..., -total_len:] |
| scores = scores + (1.0 - attention_mask.reshape(bsz, 1, 1, total_len).float()) * -1e9 |
| probs = torch.softmax(scores.float(), dim=-1).to(q_nope.dtype) |
| out = torch.einsum("bsht,btc->bshc", probs, c_kv).reshape(bsz, seq_len, self.out_dim) |
| out = out * torch.sigmoid(z) |
| return self.o_proj(out) |
|
|
| 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]]]: |
| bsz, seq_len, _ = x.shape |
| cos, sin = position_embeddings |
| if cos.dim() == 2: |
| cos = cos[:seq_len, : self.qk_rope_head_dim] |
| sin = sin[:seq_len, : self.qk_rope_head_dim] |
| else: |
| cos = cos[:, :seq_len, : self.qk_rope_head_dim] |
| sin = sin[:, :seq_len, : self.qk_rope_head_dim] |
| q_nope, q_pe = self._project_q(x) |
| q_pe = apply_rope_to_single(q_pe.permute(0, 2, 1, 3), cos, sin).permute(0, 2, 1, 3) |
| c_kv, k_pe = self._project_kv(x, cos, sin) |
| z = self.z_proj(x) |
|
|
| if past_key_values is not None: |
| past_c, past_pe = past_key_values |
| c_kv = torch.cat([past_c, c_kv], dim=1) |
| k_pe = torch.cat([past_pe, k_pe], dim=1) |
| new_past = (c_kv, k_pe) if use_cache else None |
|
|
| if self.attn_impl == "naive": |
| out = self._forward_sdpa(q_nope, q_pe, c_kv, k_pe, z, attention_mask) |
| else: |
| out = self._forward_absorb(q_nope, q_pe, c_kv, k_pe, z, attention_mask) |
| return out, new_past |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, config: YConfig31, 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) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| up, gate = self.up_proj(x), self.gate_proj(x) |
| up = nn.functional.silu(gate) * up |
| return self.down_proj(up) |
|
|
|
|
| class SengramIndexer(nn.Module): |
| def __init__(self, config: YConfig31): |
| 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]: |
| bucket_logits = self.bucket_proj(hidden_states) |
| route_scores = torch.softmax(bucket_logits.float(), dim=-1) |
| topk_ids = torch.topk(route_scores, k=self.topk, dim=-1, sorted=False).indices |
| topk_scores = route_scores.gather(-1, topk_ids) |
| denom = topk_scores.sum(dim=-1, keepdim=True).clamp_min(1e-20) |
| topk_scores = (topk_scores / denom).to(bucket_logits.dtype) |
| return topk_ids, topk_scores |
|
|
|
|
| class SengramPLE(nn.Module): |
| def __init__(self, config: YConfig31): |
| super().__init__() |
| self.hidden_size = int(config.hidden_size) |
| self.embedding = nn.Embedding(int(config.sengram_bucket_size or 4096), self.hidden_size) |
| self.key_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False) |
| self.memory_norm = RMSNorm(self.hidden_size, config.rms_norm_eps) |
| self.key_norm = RMSNorm(self.hidden_size, config.rms_norm_eps) |
| self.query_norm = RMSNorm(self.hidden_size, config.rms_norm_eps) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| topk_ids: torch.Tensor, |
| topk_scores: torch.Tensor, |
| ) -> torch.Tensor: |
| topk_embed = F.embedding(topk_ids, self.embedding.weight) |
| return (topk_embed * topk_scores.unsqueeze(-1).to(topk_embed.dtype)).sum(dim=-2) |
|
|
|
|
| class YBlock31(nn.Module): |
| def __init__(self, config: YConfig31, 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.attn = MLGA(config, layer_id) |
| self.ffn = SwiGLU(config) |
| self.se1 = SEBlock(config.hidden_size, act=ACT2FN[config.hidden_act]) |
| self.se2 = SEBlock(config.hidden_size, act=ACT2FN[config.hidden_act]) |
|
|
| def forward( |
| self, |
| x: torch.Tensor, |
| position_embeddings: tuple[torch.Tensor, torch.Tensor], |
| past_key_values: Optional[tuple[torch.Tensor, torch.Tensor]] = None, |
| use_cache: bool = False, |
| attention_mask: Optional[torch.Tensor] = None, |
| route_ids: Optional[torch.Tensor] = None, |
| route_scores: Optional[torch.Tensor] = None, |
| **kwargs, |
| ): |
| if self.use_sengram and route_ids is not None and route_scores is not None and self.sengram_ple is not None: |
| x = x + self.sengram_ple(x, route_ids, route_scores) |
| x0 = self.se1(self.input_layernorm(x)) |
| attn_out, past = self.attn( |
| x0, |
| position_embeddings, |
| past_key_values=past_key_values, |
| attention_mask=attention_mask, |
| use_cache=use_cache, |
| ) |
| x = x + attn_out |
| x0 = self.se2(self.post_attention_layernorm(x)) |
| x = x + self.ffn(x0) |
| return x, past |
|
|
|
|
| class YModel31(nn.Module): |
| def __init__(self, config: YConfig31): |
| super().__init__() |
| self.config = config |
| self.vocab_size = config.vocab_size |
| self.num_layers = config.num_hidden_layers |
| self.dropout = config.dropout |
| self.use_sengram = bool(config.use_sengram) |
| self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) |
| self.sengram_indexer = SengramIndexer(config) if self.use_sengram else None |
| self.layers = nn.ModuleList([YBlock31(config, i) for i in range(config.num_hidden_layers)]) |
| self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps) |
| freqs_cos, freqs_sin = precompute_freqs_cis( |
| dim=config.mla_qk_rope_head_dim, |
| end=config.max_position_embeddings, |
| theta=config.rope_theta, |
| rope_scaling=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 _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): |
| for key in list(state_dict.keys()): |
| if key.startswith(prefix + "sengram."): |
| state_dict.pop(key) |
| super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs) |
|
|
| def forward( |
| self, |
| input_ids: Optional[torch.Tensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| past_key_values: Optional[list] = None, |
| use_cache: bool = False, |
| cache_position: Optional[torch.LongTensor] = None, |
| position_ids: Optional[torch.LongTensor] = None, |
| **kwargs, |
| ): |
| bsz, seq_len = input_ids.shape |
| if use_cache and past_key_values is None: |
| past_key_values = [None] * self.num_layers |
| if cache_position is None: |
| if past_key_values is not None and past_key_values[0] is not None: |
| past_seen = past_key_values[0][0].shape[1] |
| else: |
| past_seen = 0 |
| cache_position = torch.arange(past_seen, past_seen + seq_len, device=input_ids.device) |
|
|
| x = self.embed_tokens(input_ids) |
| if position_ids is None: |
| position_ids = cache_position |
| position_embeddings = (self.freqs_cos[position_ids].to(x.device), self.freqs_sin[position_ids].to(x.device)) |
| route_ids = None |
| route_scores = None |
| if self.use_sengram and self.sengram_indexer is not None: |
| route_ids, route_scores = self.sengram_indexer(x) |
| new_past = [] if use_cache else None |
|
|
| for i, layer in enumerate(self.layers): |
| past = past_key_values[i] if past_key_values is not None else None |
| x, layer_past = layer( |
| x, |
| position_embeddings=position_embeddings, |
| past_key_values=past, |
| attention_mask=attention_mask, |
| use_cache=use_cache, |
| route_ids=route_ids, |
| route_scores=route_scores, |
| ) |
| if use_cache: |
| new_past.append(layer_past) |
| return self.norm(x), new_past |
|
|
|
|
| class YForCausalLM31(PreTrainedModel, GenerationMixin): |
| config_class = YConfig31 |
|
|
| def __init__(self, config: Optional[YConfig31] = None): |
| self.config = config or YConfig31() |
| super().__init__(self.config) |
| self.model = YModel31(self.config) |
| self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) |
| self.model.embed_tokens.weight = self.lm_head.weight |
| self.OUT = CausalLMOutputWithPast() |
| dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}.get(self.config.dtype) |
| if dtype is not None: |
| self.to(dtype) |
|
|
| def forward( |
| self, |
| input_ids: Optional[torch.Tensor] = None, |
| attention_mask: Optional[torch.Tensor] = None, |
| past_key_values: Optional[list] = None, |
| use_cache: bool = False, |
| logits_to_keep: Union[int, torch.Tensor] = 0, |
| cache_position: Optional[torch.LongTensor] = None, |
| **kwargs, |
| ): |
| h, past_kvs = self.model( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| past_key_values=past_key_values, |
| use_cache=use_cache, |
| cache_position=cache_position, |
| position_ids=kwargs.get("position_ids", None), |
| ) |
| slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep |
| logits = self.lm_head(h[:, slice_indices, :]) |
| self.OUT.__setitem__("last_hidden_state", h) |
| self.OUT.__setitem__("logits", logits) |
| self.OUT.__setitem__("past_key_values", past_kvs) |
| return self.OUT |
|
|
| def generate( |
| self, |
| inputs, |
| 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=True, |
| num_return_sequences=1, |
| do_sample=True, |
| repetition_penalty=1.0, |
| **kwargs, |
| ): |
| input_ids = kwargs.get("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", None) |
| past_key_values = None |
| if streamer: |
| streamer.put(input_ids.cpu()) |
| with torch.no_grad(): |
| for _ in range(max_new_tokens): |
| if use_cache and past_key_values is not None: |
| outputs = self.forward(input_ids[:, -1:], None, past_key_values, use_cache=use_cache) |
| else: |
| outputs = self.forward(input_ids, attention_mask, past_key_values, use_cache=use_cache) |
| logits = outputs.logits[:, -1, :] / temperature |
| if repetition_penalty != 1.0: |
| for i in range(input_ids.shape[0]): |
| logits[i, torch.unique(input_ids[i])] /= repetition_penalty |
| if logits_processor is not None: |
| logits = logits_processor(input_ids, logits) |
| if top_k > 0: |
| logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float("inf") |
| if top_p < 1.0: |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p |
| mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0 |
| logits[mask.scatter(1, sorted_indices, 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) |
| past_key_values = outputs.past_key_values |
| if streamer: |
| streamer.put(next_token.cpu()) |
| if eos_token_id and (next_token == eos_token_id).any(): |
| break |
| if streamer: |
| streamer.end() |
| return input_ids |
|
|
|
|
| def count_parameters(config: YConfig31) -> int: |
| return sum(p.numel() for p in YForCausalLM31(config).parameters()) |
|
|
|
|
| def _load_state_dict(path: Union[str, Path]) -> dict[str, torch.Tensor]: |
| path = Path(path) |
| if path.is_dir(): |
| safetensors_path = path / "model.safetensors" |
| bin_path = path / "pytorch_model.bin" |
| if safetensors_path.exists(): |
| path = safetensors_path |
| elif bin_path.exists(): |
| path = bin_path |
| else: |
| raise FileNotFoundError(f"no model.safetensors or pytorch_model.bin found in {path}") |
| if path.suffix == ".safetensors": |
| return load_safetensors(str(path), device="cpu") |
| return torch.load(path, map_location="cpu", weights_only=True) |
|
|
|
|
| def load_ymodel31_eval(path: Union[str, Path], config: Optional[YConfig31] = None, strict: bool = True) -> YForCausalLM31: |
| path = Path(path) |
| if config is None: |
| config_path = path / "config.json" if path.is_dir() else path.with_name("config.json") |
| if not config_path.exists(): |
| raise FileNotFoundError("config is required when config.json is not next to the checkpoint") |
| config = YConfig31.from_json_file(str(config_path)) |
| model = YForCausalLM31(config) |
| state = _load_state_dict(path) |
| model.load_state_dict(state, strict=strict) |
| model.eval() |
| return model |
|
|
|
|
| YModel31Eval = YModel31 |
| YForCausalLM31Eval = YForCausalLM31 |
|
|
|
|
| __all__ = [ |
| "MLGA", |
| "RMSNorm", |
| "SEBlock", |
| "SengramIndexer", |
| "SengramPLE", |
| "SwiGLU", |
| "YBlock31", |
| "YConfig31", |
| "YForCausalLM31", |
| "YForCausalLM31Eval", |
| "YModel31", |
| "YModel31Eval", |
| "apply_rope_to_single", |
| "count_parameters", |
| "load_ymodel31_eval", |
| "precompute_freqs_cis", |
| ] |
|
|
