mnauf/redditGPT
0
1"""2Full definition of a GPT Language Model, all of it in this single file.3References:41) the official GPT-2 TensorFlow implementation released by OpenAI:5https://github.com/openai/gpt-2/blob/master/src/model.py62) huggingface/transformers PyTorch implementation:7https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py8"""9 10import math11import inspect12from dataclasses import dataclass13 14import torch15import torch.nn as nn16from torch.nn import functional as F17 18# @torch.jit.script # good to enable when not using torch.compile, disable when using (our default)19def new_gelu(x):20 """21 Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).22 Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.0841523 """24 return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))25 26class LayerNorm(nn.Module):27 """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """28 29 def __init__(self, ndim, bias):30 super().__init__()31 self.weight = nn.Parameter(torch.ones(ndim))32 self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None33 34 def forward(self, input):35 return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)36 37class CausalSelfAttention(nn.Module):38 39 def __init__(self, config):40 super().__init__()41 assert config.n_embd % config.n_head == 042 # key, query, value projections for all heads, but in a batch43 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)44 # output projection45 self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)46 # regularization47 self.attn_dropout = nn.Dropout(config.dropout)48 self.resid_dropout = nn.Dropout(config.dropout)49 self.n_head = config.n_head50 self.n_embd = config.n_embd51 self.dropout = config.dropout52 # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.053 self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')54 if not self.flash:55 print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")56 # causal mask to ensure that attention is only applied to the left in the input sequence57 self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))58 .view(1, 1, config.block_size, config.block_size))59 60 def forward(self, x):61 B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)62 63 # calculate query, key, values for all heads in batch and move head forward to be the batch dim64 q, k, v = self.c_attn(x).split(self.n_embd, dim=2)65 k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)66 q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)67 v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)68 69 # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)70 if self.flash:71 # efficient attention using Flash Attention CUDA kernels72 y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)73 else:74 # manual implementation of attention75 att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))76 att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))77 att = F.softmax(att, dim=-1)78 att = self.attn_dropout(att)79 y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)80 y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side81 82 # output projection83 y = self.resid_dropout(self.c_proj(y))84 return y85 86class MLP(nn.Module):87 88 def __init__(self, config):89 super().__init__()90 self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)91 self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)92 self.dropout = nn.Dropout(config.dropout)93 94 def forward(self, x):95 x = self.c_fc(x)96 x = new_gelu(x)97 x = self.c_proj(x)98 x = self.dropout(x)99 return x100 101class Block(nn.Module):102 103 def __init__(self, config):104 super().__init__()105 self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)106 self.attn = CausalSelfAttention(config)107 self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)108 self.mlp = MLP(config)109 110 def forward(self, x):111 x = x + self.attn(self.ln_1(x))112 x = x + self.mlp(self.ln_2(x))113 return x114 115@dataclass116class GPTConfig:117 block_size: int = 1024118 vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency119 n_layer: int = 12120 n_head: int = 12121 n_embd: int = 768122 dropout: float = 0.0123 bias: bool = True # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster124 125class GPT(nn.Module):126 127 def __init__(self, config):128 super().__init__()129 assert config.vocab_size is not None130 assert config.block_size is not None131 self.config = config132 133 self.transformer = nn.ModuleDict(dict(134 wte = nn.Embedding(config.vocab_size, config.n_embd),135 wpe = nn.Embedding(config.block_size, config.n_embd),136 drop = nn.Dropout(config.dropout),137 h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),138 ln_f = LayerNorm(config.n_embd, bias=config.bias),139 ))140 self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)141 # with weight tying when using torch.compile() some warnings get generated:142 # "UserWarning: functional_call was passed multiple values for tied weights.143 # This behavior is deprecated and will be an error in future versions"144 # not 100% sure what this is, so far seems to be harmless. TODO investigate145 self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying146 147 # init all weights148 self.apply(self._init_weights)149 # apply special scaled init to the residual projections, per GPT-2 paper150 for pn, p in self.named_parameters():151 if pn.endswith('c_proj.weight'):152 torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))153 154 # report number of parameters155 print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))156 157 def get_num_params(self, non_embedding=True):158 """159 Return the number of parameters in the model.160 For non-embedding count (default), the position embeddings get subtracted.161 The token embeddings would too, except due to the parameter sharing these162 params are actually used as weights in the final layer, so we include them.163 """164 n_params = sum(p.numel() for p in self.parameters())165 if non_embedding:166 n_params -= self.transformer.wpe.weight.numel()167 return n_params168 169 def _init_weights(self, module):170 if isinstance(module, nn.Linear):171 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)172 if module.bias is not None:173 torch.nn.init.zeros_(module.bias)174 elif isinstance(module, nn.Embedding):175 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)176 177 def forward(self, idx, targets=None):178 device = idx.device179 b, t = idx.size()180 assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"181 pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t)182 183 # forward the GPT model itself184 tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)185 pos_emb = self.transformer.wpe(pos) # position embeddings of shape (1, t, n_embd)186 x = self.transformer.drop(tok_emb + pos_emb)187 for block in self.transformer.h:188 x = block(x)189 x = self.transformer.ln_f(x)190 191 if targets is not None:192 # if we are given some desired targets also calculate the loss193 logits = self.lm_head(x)194 loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)195 else:196 # inference-time mini-optimization: only forward the lm_head on the very last position197 logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim198 loss = None199 200 return logits, loss201 202 def crop_block_size(self, block_size):203 # model surgery to decrease the block size if necessary204 # e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)205 # but want to use a smaller block size for some smaller, simpler model206 assert block_size <= self.config.block_size207 self.config.block_size = block_size208 self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])209 for block in self.transformer.h:210 if hasattr(block.attn, 'bias'):211 block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]212 213 @classmethod214 def from_pretrained(cls, model_type, override_args=None):215 assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}216 override_args = override_args or {} # default to empty dict217 # only dropout can be overridden see more notes below218 assert all(k == 'dropout' for k in override_args)219 from transformers import GPT2LMHeadModel220 print("loading weights from pretrained gpt: %s" % model_type)221 222 # n_layer, n_head and n_embd are determined from model_type223 config_args = {224 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params225 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params226 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params227 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params228 }[model_type]229 print("forcing vocab_size=50257, block_size=1024, bias=True")230 config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints231 config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints232 config_args['bias'] = True # always True for GPT model checkpoints233 # we can override the dropout rate, if desired234 if 'dropout' in override_args:235 print(f"overriding dropout rate to {override_args['dropout']}")236 config_args['dropout'] = override_args['dropout']237 # create a from-scratch initialized minGPT model238 config = GPTConfig(**config_args)239 model = GPT(config)240 sd = model.state_dict()241 sd_keys = sd.keys()242 sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param243 244 # init a huggingface/transformers model245 model_hf = GPT2LMHeadModel.from_pretrained(model_type)246 sd_hf = model_hf.state_dict()247 248 # copy while ensuring all of the parameters are aligned and match in names and shapes249 sd_keys_hf = sd_hf.keys()250 sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer251 sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)252 transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']253 # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear254 # this means that we have to transpose these weights when we import them255 assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"256 for k in sd_keys_hf:257 if any(k.endswith(w) for w in transposed):258 # special treatment for the Conv1D weights we need to transpose259 assert sd_hf[k].shape[::-1] == sd[k].shape260 with torch.no_grad():261 sd[k].copy_(sd_hf[k].t())262 else:263 # vanilla copy over the other parameters264 assert sd_hf[k].shape == sd[k].shape265 with torch.no_grad():266 sd[k].copy_(sd_hf[k])267 268 return model269 270 def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):271 # start with all of the candidate parameters272 param_dict = {pn: p for pn, p in self.named_parameters()}273 # filter out those that do not require grad274 param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}275 # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.276 # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.277 decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]278 nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]279 optim_groups = [280 {'params': decay_params, 'weight_decay': weight_decay},281 {'params': nodecay_params, 'weight_decay': 0.0}282 ]283 num_decay_params = sum(p.numel() for p in decay_params)284 num_nodecay_params = sum(p.numel() for p in nodecay_params)285 print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")286 print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")287 # Create AdamW optimizer and use the fused version if it is available288 fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters289 use_fused = fused_available and device_type == 'cuda'290 extra_args = dict(fused=True) if use_fused else dict()291 optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)292 print(f"using fused AdamW: {use_fused}")293 294 return optimizer295 296 def estimate_mfu(self, fwdbwd_per_iter, dt):297 """ estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """298 # first estimate the number of flops we do per iteration.299 # see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311300 N = self.get_num_params()301 cfg = self.config302 L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd//cfg.n_head, cfg.block_size303 flops_per_token = 6*N + 12*L*H*Q*T304 flops_per_fwdbwd = flops_per_token * T305 flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter306 # express our flops throughput as ratio of A100 bfloat16 peak flops307 flops_achieved = flops_per_iter * (1.0/dt) # per second308 flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS309 mfu = flops_achieved / flops_promised310 return mfu311 312 @torch.no_grad()313 def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):314 """315 Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete316 the sequence max_new_tokens times, feeding the predictions back into the model each time.317 Most likely you'll want to make sure to be in model.eval() mode of operation for this.318 """319 for _ in range(max_new_tokens):320 # if the sequence context is growing too long we must crop it at block_size321 idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]322 # forward the model to get the logits for the index in the sequence323 logits, _ = self(idx_cond)324 # pluck the logits at the final step and scale by desired temperature325 logits = logits[:, -1, :] / temperature326 # optionally crop the logits to only the top k options327 if top_k is not None:328 v, _ = torch.topk(logits, min(top_k, logits.size(-1)))329 logits[logits < v[:, [-1]]] = -float('Inf')330 # apply softmax to convert logits to (normalized) probabilities331 probs = F.softmax(logits, dim=-1)332 # sample from the distribution333 idx_next = torch.multinomial(probs, num_samples=1)334 # append sampled index to the running sequence and continue335 idx = torch.cat((idx, idx_next), dim=1)336 337 return idx338