| """Explicit attention conversion adapter. Does not edit or retrain TinyGPT.""" |
| import math |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class ConversionGraph(nn.Module): |
| def __init__(self, model): |
| super().__init__() |
| self.model = model.eval() |
| config = model.config |
| self.heads = config.n_heads |
| self.head_dim = config.d_model // config.n_heads |
| self.width = config.d_model |
| mask = torch.full((config.context_length, config.context_length), -float("inf")) |
| self.register_buffer("causal_mask", torch.triu(mask, diagonal=1)) |
|
|
| def forward(self, input_ids): |
| batch, length = input_ids.shape |
| positions = torch.arange(length, device=input_ids.device) |
| hidden = self.model.token_embedding(input_ids) + self.model.position_embedding(positions) |
| for block in self.model.blocks: |
| normalized = block.attention_norm(hidden) |
| qkv = block.attention.qkv(normalized).reshape(batch, length, 3, self.heads, self.head_dim) |
| query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0) |
| attention = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| attention = torch.softmax(attention + self.causal_mask[:length, :length], dim=-1) |
| attended = torch.matmul(attention, value).transpose(1, 2).reshape(batch, length, self.width) |
| hidden = hidden + block.attention.projection(attended) |
| hidden = hidden + block.ffn(block.ffn_norm(hidden)) |
| return self.model.lm_head(self.model.final_norm(hidden)) |
|
|
|
|
| @torch.inference_mode() |
| def trace_graph(model): |
| graph = ConversionGraph(model).eval() |
| sample = torch.zeros((1, min(16, model.config.context_length)), dtype=torch.int32) |
| traced = torch.jit.trace(graph, sample) |
| |
| generator = torch.Generator().manual_seed(716) |
| for length in sorted({1, min(7, model.config.context_length), model.config.context_length}): |
| ids = torch.randint(model.config.vocab_size, (1, length), generator=generator, dtype=torch.int32) |
| original = model(ids.long()) |
| for adapted in (graph(ids), traced(ids)): |
| torch.testing.assert_close(adapted, original, atol=3e-5, rtol=3e-4) |
| return traced |
|
|
