"""Fixed, untrained five-token model: input IDs, context vectors and logits. Reference: the Pre-Norm structure in Karpathy's GPT exercise. Run with PyTorch installed: python gpt_model_interface_demo.py Writes gpt_model_interface_demo.json beside this script. CPU float32. """ import json from pathlib import Path import torch from torch import nn from torch.nn import functional as F class TeachingGPT(nn.Module): def __init__(self): super().__init__() self.token_embedding = nn.Embedding(5, 4) self.position_embedding = nn.Embedding(8, 4) self.ln1 = nn.LayerNorm(4) self.query = nn.Linear(4, 4, bias=False) self.key = nn.Linear(4, 4, bias=False) self.value = nn.Linear(4, 4, bias=False) self.attention_projection = nn.Linear(4, 4, bias=False) self.ln2 = nn.LayerNorm(4) self.ff1 = nn.Linear(4, 16) self.ff2 = nn.Linear(16, 4) self.ln_final = nn.LayerNorm(4) # Independent weights and output bias, as in the V7 exercise. self.lm_head = nn.Linear(4, 5) self.set_teaching_weights() def set_teaching_weights(self): """Set illustrative weights; they do not encode learned language patterns.""" with torch.no_grad(): self.token_embedding.weight.copy_(torch.tensor([ [1., 0., 1., 0.], # ID 0: a [0., 1., 0., 1.], # ID 1: b [1., 1., 0., 0.], # ID 2: c [0., 0., 1., 1.], # ID 3: d [.5, .5, .5, .5], # ID 4: e ])) self.position_embedding.weight.zero_() self.position_embedding.weight[:3].copy_(torch.tensor([ [.1, .2, 0., 0.], [0., .1, .2, 0.], [0., 0., .1, .2], ])) eye = torch.eye(4) for layer in (self.query, self.key, self.value, self.attention_projection): layer.weight.copy_(eye) self.ff1.weight.copy_(torch.cat([eye, -eye, .5 * eye, -.5 * eye])) self.ff2.weight.copy_(torch.cat([.2 * eye, .05 * eye, .1 * eye, -.02 * eye], dim=1)) self.ff1.bias.zero_() self.ff2.bias.zero_() self.lm_head.weight.copy_(torch.tensor([ [1., 0., 0., 0.], [1., 0., 1., 0.], [0., 1., 0., 0.], [-1., 0., 0., 0.], [0., 0., 0., 1.], ])) self.lm_head.bias.copy_(torch.tensor([0., .1, 0., -.2, 0.])) def forward(self, ids): positions = torch.arange(ids.shape[1], device=ids.device) token_vectors = self.token_embedding(ids) position_vectors = self.position_embedding(positions) x0 = token_vectors + position_vectors normalized1 = self.ln1(x0) q, k, v = self.query(normalized1), self.key(normalized1), self.value(normalized1) attention_scores = q @ k.transpose(-2, -1) / (4 ** .5) future_mask = torch.ones(ids.shape[1], ids.shape[1], dtype=torch.bool, device=ids.device).triu(1) masked_scores = attention_scores.masked_fill(future_mask, float('-inf')) attention_weights = masked_scores.softmax(dim=-1) attention_values = attention_weights @ v attention_output = self.attention_projection(attention_values) x1 = x0 + attention_output normalized2 = self.ln2(x1) ff_hidden = F.relu(self.ff1(normalized2)) ff_output = self.ff2(ff_hidden) x2 = x1 + ff_output hidden = self.ln_final(x2) logits = self.lm_head(hidden) return dict( positions=positions, token_vectors=token_vectors, position_vectors=position_vectors, x0=x0, normalized1=normalized1, q=q, k=k, v=v, attention_scores=attention_scores, attention_weights=attention_weights, attention_values=attention_values, attention_output=attention_output, x1=x1, normalized2=normalized2, ff_hidden=ff_hidden, ff_output=ff_output, x2=x2, hidden=hidden, logits=logits, last_hidden=hidden[:, -1, :], last_logits=logits[:, -1, :], last_probabilities=logits[:, -1, :].softmax(dim=-1), ) def main(): model = TeachingGPT() ids = torch.tensor([[0, 1, 0]]) targets = torch.tensor([[1, 0, 1]]) # Text a | b | a | b; targets shift one token ahead. trace = model(ids) losses = F.cross_entropy(trace['logits'].reshape(-1, 5), targets.reshape(-1), reduction='none') loss = losses.mean() assert torch.allclose(trace['x0'][0], torch.tensor([[1.1, .2, 1., 0.], [0., 1.1, .2, 1.], [1., 0., 1.1, .2]])) assert torch.allclose(trace['attention_weights'].triu(1), torch.zeros(1, 3, 3)) assert torch.allclose(trace['attention_weights'].sum(-1), torch.ones(1, 3)) assert torch.allclose(trace['last_logits'], trace['last_hidden'] @ model.lm_head.weight.T + model.lm_head.bias) assert torch.allclose(trace['last_probabilities'].sum(-1), torch.ones(1)) assert model.lm_head.weight is not model.token_embedding.weight loss.backward() gradient_checks = { name: bool(module.weight.grad.abs().sum().item() > 0) for name, module in ( ('token_embedding', model.token_embedding), ('position_embedding', model.position_embedding), ('query', model.query), ('key', model.key), ('value', model.value), ('ff1', model.ff1), ('ff2', model.ff2), ('lm_head', model.lm_head), ) } assert all(gradient_checks.values()) def serial(tensor): return tensor.detach().cpu().tolist() record = { 'description': 'Fixed, untrained weights. One forward pass supplies both detail diagrams. Displayed values are rounded; calculations use full precision.', 'vocab': {'a': 0, 'b': 1, 'c': 2, 'd': 3, 'e': 4}, 'config': {'B': 1, 'T': 3, 'D': 4, 'V': 5, 'max_positions': 8, 'layers': 1, 'heads': 1}, 'input_ids': serial(ids), 'target_ids': serial(targets), 'embedding': serial(model.token_embedding.weight), 'position_embedding': serial(model.position_embedding.weight), 'lm_head_weight': serial(model.lm_head.weight), 'lm_head_bias': serial(model.lm_head.bias), 'trace': {name: serial(tensor) for name, tensor in trace.items()}, 'shapes': {name: list(tensor.shape) for name, tensor in trace.items()}, 'loss_per_position': serial(losses), 'loss_mean': loss.item(), 'last_greedy_id': int(trace['last_logits'].argmax(-1).item()), 'gradient_checks': gradient_checks, } path = Path(__file__).with_suffix('.json') path.write_text(json.dumps(record, ensure_ascii=False, indent=2) + '\n', encoding='utf-8') for name in ('x0', 'attention_weights', 'x1', 'ff_output', 'hidden', 'last_logits', 'last_probabilities'): print(name, serial(trace[name])) print('loss_per_position', serial(losses)) print('mean_loss', loss.item()) print('gradient_checks', gradient_checks) print('saved', path) if __name__ == '__main__': main()