| """Fixed V4-a capacity allocation, with a separate FP32 research KV path.""" |
| import math |
| from dataclasses import dataclass, fields |
|
|
| import torch |
|
|
| from vimeml.training.model_v2 import GPTV2Config, TinyGPTV2 |
|
|
| PARAMETERS = 15_070_848 |
| LAYER_PARAMETERS = 1_696_512 |
|
|
|
|
| @dataclass(frozen=True) |
| class GPTV4Config(GPTV2Config): |
| vocab_size: int = 8192 |
| context_length: int = 128 |
| d_model: int = 384 |
| n_heads: int = 6 |
| n_layers: int = 7 |
| d_ff: int = 960 |
| dropout: float = 0.0 |
|
|
| def __post_init__(self): |
| super().__post_init__() |
| for field in fields(self): |
| if type(getattr(self, field.name)) is not type(field.default) or getattr(self, field.name) != field.default: |
| raise ValueError(f"V4-a {field.name} is fixed to {field.default}") |
|
|
|
|
| @dataclass |
| class KVCacheV4: |
| key: torch.Tensor |
| value: torch.Tensor |
| length: int = 0 |
|
|
| def clone(self): |
| return KVCacheV4(self.key.clone(), self.value.clone(), self.length) |
|
|
| @torch.inference_mode() |
| def restore(self, snapshot): |
| if (self.key.shape != snapshot.key.shape or self.value.shape != snapshot.value.shape |
| or self.key.device != snapshot.key.device or self.value.device != snapshot.value.device |
| or self.key.dtype != snapshot.key.dtype or self.value.dtype != snapshot.value.dtype |
| or type(snapshot.length) is not int or not 0 <= snapshot.length <= 128): |
| raise ValueError("Incompatible KV snapshot") |
| self.key.copy_(snapshot.key) |
| self.value.copy_(snapshot.value) |
| self.length = snapshot.length |
|
|
|
|
| class TinyGPTV4(TinyGPTV2): |
| """Same proven operations as V2; independent fixed architecture identity.""" |
| def __init__(self, config=GPTV4Config()): |
| if type(config) is not GPTV4Config: |
| raise ValueError("V4 requires its fixed GPTV4Config") |
| super().__init__(config) |
| if self.parameter_count() != PARAMETERS or any( |
| sum(p.numel() for p in block.parameters()) != LAYER_PARAMETERS for block in self.blocks): |
| raise ValueError("V4 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 = (7, batch_size, 6, 128, 64) |
| device = self.token_embedding.weight.device |
| return KVCacheV4(torch.zeros(shape, device=device, dtype=torch.float32), |
| torch.zeros(shape, device=device, dtype=torch.float32)) |
|
|
| @torch.inference_mode() |
| def forward_cached(self, input_ids, cache): |
| if self.training or self.token_embedding.weight.dtype != torch.float32: |
| raise ValueError("Research 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 |
| shape = (7, batch, 6, 128, 64) |
| if (cache.key.shape != shape or cache.value.shape != shape |
| or cache.key.dtype != torch.float32 or cache.value.dtype != torch.float32 |
| or cache.key.device != input_ids.device or cache.value.device != input_ids.device |
| or input_ids.device != self.token_embedding.weight.device |
| or type(cache.length) is not int or cache.length < 0 |
| 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()) >= 8192: |
| raise ValueError("Token outside fixed 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): |
| normalized = block.attention_norm(hidden) |
| qkv = block.attention.qkv(normalized).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) |
| probabilities = torch.softmax(scores.masked_fill(~causal, -torch.inf), dim=-1) |
| attended = torch.matmul(probabilities, cache.value[layer, :, :, :end]) |
| attended = attended.transpose(1, 2).reshape(batch, size, 384) |
| hidden = hidden + block.attention.projection(attended) |
| hidden = hidden + block.ffn(block.ffn_norm(hidden)) |
| cache.length = end |
| return self.lm_head(self.final_norm(hidden)) |
|
|
