File size: 5,575 Bytes
70d134c | 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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutput
from .configuration_tinygpt import TinyGPTConfig
# ─────────────────────────────────────────────
# These three classes are copied verbatim (module
# names included) from the original training script,
# so a raw_model.state_dict() from training loads
# straight into TinyGPTForCausalLM with strict=True.
# ─────────────────────────────────────────────
class MultiHeadAttention(nn.Module):
def __init__(self, n_embd, n_head):
super().__init__()
assert n_embd % n_head == 0
self.n_head = n_head
self.head_size = n_embd // n_head
self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
self.proj = nn.Linear(n_embd, n_embd, bias=False)
def forward(self, x):
B, T, C = x.shape
q, k, v = self.c_attn(x).split(C, dim=2)
q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
out = out.transpose(1, 2).contiguous().view(B, T, C)
return self.proj(out)
class FeedForward(nn.Module):
def __init__(self, n_embd):
super().__init__()
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.GELU(),
nn.Linear(4 * n_embd, n_embd),
)
def forward(self, x):
return self.net(x)
class Block(nn.Module):
def __init__(self, n_embd, n_head):
super().__init__()
self.sa = MultiHeadAttention(n_embd, n_head)
self.ffwd = FeedForward(n_embd)
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
def forward(self, x):
x = x + self.sa(self.ln1(x))
x = x + self.ffwd(self.ln2(x))
return x
class TinyGPTForCausalLM(PreTrainedModel, GenerationMixin):
config_class = TinyGPTConfig
main_input_name = "input_ids"
_no_split_modules = ["Block"]
_tied_weights_keys = {"lm_head.weight": "token_embedding_table.weight"}
def __init__(self, config: TinyGPTConfig):
super().__init__(config)
self.token_embedding_table = nn.Embedding(config.vocab_size, config.n_embd)
self.position_embedding_table = nn.Embedding(config.block_size, config.n_embd)
self.blocks = nn.Sequential(
*[Block(config.n_embd, config.n_head) for _ in range(config.n_layer)]
)
self.ln_f = nn.LayerNorm(config.n_embd)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
# runs custom _init_weights below on every submodule, then ties
# lm_head.weight <-> token_embedding_table.weight via
# get_input_embeddings()/get_output_embeddings() + _tied_weights_keys
# (same net effect as the manual assignment in the training script,
# but done the way HF's save/load machinery expects)
self.post_init()
# this model has no KV-cache implementation: force generate() to
# always recompute the full (truncated) sequence each step
self.generation_config.use_cache = False
# kept identical to the training script; only used if you ever call
# .init_weights() on a *fresh* (non-loaded) model
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
# ── HF plumbing for tied weights / embeddings ──
def get_input_embeddings(self):
return self.token_embedding_table
def set_input_embeddings(self, value):
self.token_embedding_table = value
def get_output_embeddings(self):
return self.lm_head
def set_output_embeddings(self, new_embeddings):
self.lm_head = new_embeddings
# ── forward pass, same math as the training script ──
def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs):
B, T = input_ids.shape
device = input_ids.device
tok_emb = self.token_embedding_table(input_ids)
pos_emb = self.position_embedding_table(torch.arange(T, device=device))
x = tok_emb + pos_emb
x = self.blocks(x)
x = self.ln_f(x)
logits = self.lm_head(x)
loss = None
if labels is not None:
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)),
labels.view(-1),
)
return CausalLMOutput(loss=loss, logits=logits)
# ── generation support (no KV cache, so we just resend the
# truncated running sequence every step — matches the
# behaviour of the model.generate() used at training time) ──
def prepare_inputs_for_generation(self, input_ids, **kwargs):
if input_ids.shape[1] > self.config.block_size:
input_ids = input_ids[:, -self.config.block_size :]
return {"input_ids": input_ids}
|