File size: 4,940 Bytes
c83c617 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | """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))
|