| """vimeml-tiny-zh-v1 decoder: the V4 block family with two registered shapes. |
| |
| d384, 6 heads, SwiGLU 960, learned positions 128, tied embeddings; only the |
| vocabulary and layer count vary, fixed to the 8K/L10 and 12K/L9 pairs of |
| docs/zh-CN/training-zh-v1-plan.md. |
| """ |
| import math |
| from dataclasses import dataclass |
|
|
| import torch |
|
|
| from vimeml.training.model_v2 import GPTV2Config, TinyGPTV2 |
|
|
| LAYER_PARAMETERS = 1_696_512 |
| |
| SHAPES = {8192: (10, 20_160_384), 12288: (9, 20_036_736)} |
|
|
|
|
| @dataclass(frozen=True) |
| class GPTZh13Config(GPTV2Config): |
| vocab_size: int = 12288 |
| context_length: int = 128 |
| d_model: int = 384 |
| n_heads: int = 6 |
| n_layers: int = 9 |
| d_ff: int = 960 |
| dropout: float = 0.0 |
|
|
| def __post_init__(self): |
| super().__post_init__() |
| if self.vocab_size not in SHAPES or self.n_layers != SHAPES[self.vocab_size][0]: |
| raise ValueError("vimeml-tiny-zh-v1 shapes are 8192/L10 and 12288/L9") |
| fixed = {"context_length": 128, "d_model": 384, "n_heads": 6, "d_ff": 960, "dropout": 0.0} |
| for name, value in fixed.items(): |
| if type(getattr(self, name)) is not type(value) or getattr(self, name) != value: |
| raise ValueError(f"vimeml-tiny-zh-v1 {name} is fixed to {value}") |
|
|
|
|
| @dataclass |
| class KVCacheZh13: |
| key: torch.Tensor |
| value: torch.Tensor |
| length: int = 0 |
|
|
| def clone(self): |
| return KVCacheZh13(self.key.clone(), self.value.clone(), self.length) |
|
|
|
|
| class TinyGPTZh13(TinyGPTV2): |
| def __init__(self, config=GPTZh13Config()): |
| if type(config) is not GPTZh13Config: |
| raise ValueError("vimeml-tiny-zh-v1 requires GPTZh13Config") |
| super().__init__(config) |
| if self.parameter_count() != SHAPES[config.vocab_size][1] or any( |
| sum(p.numel() for p in block.parameters()) != LAYER_PARAMETERS for block in self.blocks): |
| raise ValueError("vimeml-tiny-zh-v1 parameter budget does not match its contract") |
|
|
| def new_cache(self, batch_size=1): |
| if type(batch_size) is not int or batch_size < 1: |
| raise ValueError("Positive integer cache batch required") |
| shape = (self.config.n_layers, batch_size, 6, 128, 64) |
| device = self.token_embedding.weight.device |
| return KVCacheZh13(torch.zeros(shape, device=device), torch.zeros(shape, device=device)) |
|
|
| @torch.inference_mode() |
| def forward_cached(self, input_ids, cache): |
| """FP32 research KV path; matches forward() on the same prefix.""" |
| if self.training or self.token_embedding.weight.dtype != torch.float32: |
| raise ValueError("KV path requires eval mode and FP32 weights") |
| if input_ids.ndim != 2 or input_ids.dtype not in (torch.int32, torch.int64): |
| raise ValueError("Expected integer [batch,time] tokens") |
| batch, size = input_ids.shape |
| if (cache.key.shape != (self.config.n_layers, batch, 6, 128, 64) or size < 1 |
| or cache.length + size > 128): |
| raise ValueError("Invalid KV cache or context overflow") |
| if int(input_ids.min()) < 0 or int(input_ids.max()) >= self.config.vocab_size: |
| raise ValueError("Token outside vocabulary") |
| start, end = cache.length, cache.length + size |
| positions = torch.arange(start, end, device=input_ids.device) |
| hidden = self.token_embedding(input_ids) + self.position_embedding(positions) |
| causal = positions[:, None] >= torch.arange(end, device=input_ids.device)[None, :] |
| for layer, block in enumerate(self.blocks): |
| qkv = block.attention.qkv(block.attention_norm(hidden)).reshape(batch, size, 3, 6, 64) |
| query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0) |
| cache.key[layer, :, :, start:end].copy_(key) |
| cache.value[layer, :, :, start:end].copy_(value) |
| scores = torch.matmul(query, cache.key[layer, :, :, :end].transpose(-2, -1)) / math.sqrt(64) |
| attended = torch.matmul(torch.softmax(scores.masked_fill(~causal, -torch.inf), dim=-1), |
| cache.value[layer, :, :, :end]) |
| hidden = hidden + block.attention.projection(attended.transpose(1, 2).reshape(batch, size, 384)) |
| hidden = hidden + block.ffn(block.ffn_norm(hidden)) |
| cache.length = end |
| return self.lm_head(self.final_norm(hidden)) |
|
|
