| import torch
|
| import math
|
| import inspect
|
| from dataclasses import dataclass
|
| import torch.nn as nn
|
| from torch.nn import functional as F
|
|
|
| import tiktoken
|
| enc = tiktoken.get_encoding("gpt2")
|
| print("tiktoken cached ✅")
|
|
|
| @dataclass
|
| class GPTConfig:
|
| block_size:int=512
|
| vocab_size:int=50304
|
| n_layer:int=10
|
| n_head:int=8
|
| n_embd:int=256
|
| dropout:float=0.0
|
| bias:bool=False
|
|
|
| class LayerNorm(nn.Module):
|
| def __init__(self,ndim,bias):
|
| super().__init__()
|
| self.weight=nn.Parameter(torch.ones(ndim))
|
| self.bias=nn.Parameter(torch.zeros(ndim)) if bias else None
|
|
|
| def forward(self,input):
|
| return F.layer_norm(input,self.weight.shape, self.weight,self.bias,1e-5)
|
|
|
| class CausalSelfAttention(nn.Module):
|
| def __init__(self,config):
|
| super().__init__()
|
| assert config.n_embd % config.n_head==0
|
| self.c_attn=nn.Linear(config.n_embd,3*config.n_embd,bias= config.bias)
|
| self.c_proj=nn.Linear(config.n_embd,config.n_embd,bias=config.bias)
|
| self.attn_dropout=nn.Dropout(config.dropout)
|
| self.resid_dropout = nn.Dropout(config.dropout)
|
| self.n_head = config.n_head
|
| self.n_embd = config.n_embd
|
| self.dropout = config.dropout
|
| self.flash=hasattr(torch.nn.functional,'scaled_dot_product_attention')
|
| if not self.flash:
|
| print("WARNING using slow attention")
|
| self.register_buffer("bias",torch.tril(torch.ones(config.block_size,config.block_size)).view(1,1,config.block_size,config.block_size))
|
|
|
| def forward(self,x):
|
| B,T,C=x.size()
|
|
|
| q, k, v=self.c_attn(x).split(self.n_embd,dim=2)
|
| k=k.view(B,T,self.n_head, C //self.n_head).transpose(1,2)
|
| q=q.view(B,T,self.n_head, C //self.n_head).transpose(1,2)
|
| v=v.view(B,T,self.n_head, C //self.n_head).transpose(1,2)
|
|
|
| if self.flash:
|
| 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)
|
|
|
| else:
|
| att=(q@k.transpose(-2,-1))*(1.0/math.sqrt(k.size(-1)))
|
| att=att.masked_fill(self.bias[:,:,:T,:T]==0,float('-inf'))
|
| att=F.softmax(att,dim=-1)
|
| att=self.attn_dropout(att)
|
| y=att@v
|
| y=y.transpose(1,2).contiguous().view(B, T, C)
|
| y=self.resid_dropout(self.c_proj(y))
|
| return y
|
|
|
| class MLP(nn.Module):
|
| def __init__(self,config):
|
| super().__init__()
|
| self.c_fc=nn.Linear(config.n_embd,4*config.n_embd,bias=config.bias)
|
| self.gelu=nn.GELU()
|
| self.c_proj=nn.Linear(4*config.n_embd,config.n_embd,bias=config.bias)
|
| self.dropout=nn.Dropout(config.dropout)
|
|
|
| def forward(self,x):
|
| x=self.c_fc(x)
|
| x=self.gelu(x)
|
| x=self.c_proj(x)
|
| x=self.dropout(x)
|
| return x
|
|
|
| class Block(nn.Module):
|
|
|
| def __init__(self,config):
|
| super().__init__()
|
| self.ln_1=LayerNorm(config.n_embd,bias=config.bias)
|
| self.attn=CausalSelfAttention(config)
|
| self.ln_2=LayerNorm(config.n_embd,bias=config.bias)
|
| self.mlp=MLP(config)
|
| def forward(self,x):
|
| x=x+self.attn(self.ln_1(x))
|
| x=x+self.mlp(self.ln_2(x))
|
| return x
|
|
|
| class GPT(nn.Module):
|
| def __init__(self,config):
|
| super().__init__()
|
| assert config.vocab_size is not None
|
| assert config.block_size is not None
|
| self.config=config
|
|
|
| self.transformer = nn.ModuleDict(dict(
|
| wte=nn.Embedding(config.vocab_size,config.n_embd),
|
| wpe=nn.Embedding(config.block_size, config.n_embd),
|
| drop=nn.Dropout(config.dropout),
|
| h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
|
| ln_f=LayerNorm(config.n_embd, bias=config.bias),
|
| ))
|
| self.lm_head=nn.Linear(config.n_embd, config.vocab_size, bias=False)
|
|
|
| self.transformer.wte.weight=self.lm_head.weight
|
| self.apply(self._init_weights)
|
| for pn,p in self.named_parameters():
|
| if pn.endswith('c_proj.weight'):
|
| torch.nn.init.normal_(p,mean=0.0,std=0.02/math.sqrt(2*config.n_layer))
|
| print("Number of Parameters: %.2fM" % (self.get_num_params()/1e6))
|
|
|
| def get_num_params(self,non_embedding=True):
|
| n_params=sum(p.numel() for p in self.parameters())
|
| if non_embedding:
|
| n_params-=self.transformer.wpe.weight.numel()
|
| return n_params
|
|
|
| def _init_weights(self,module):
|
| if isinstance(module,nn.Linear):
|
| torch.nn.init.normal_(module.weight, mean=0.0,std=0.02)
|
| 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=0.02)
|
|
|
| def forward(self,idx,targets=None):
|
| device =idx.device
|
| b,t=idx.size()
|
|
|
| assert t<=self.config.block_size,f"Cannot forward Sequence of Length {t} block size is only {self.config.block_size}"
|
| pos=torch.arange(0,t,dtype=torch.long, device=device)
|
| tok_emb=self.transformer.wte(idx)
|
| pos_emb=self.transformer.wpe(pos)
|
| x=self.transformer.drop(tok_emb + pos_emb)
|
| for block in self.transformer.h:
|
| x=block(x)
|
| x=self.transformer.ln_f(x)
|
| if targets is not None:
|
| logits=self.lm_head(x)
|
| loss=F.cross_entropy(logits.view(-1,logits.size(-1)),targets.view(-1),ignore_index=-1)
|
| else:
|
| logits=self.lm_head(x[:,[-1],:])
|
| loss=None
|
| return logits,loss
|
| def configure_optimizers(self, weight_decay,learning_rate,betas,device_type):
|
| param_dict={pn:p for pn, p in self.named_parameters()}
|
| param_dict={pn:p for pn, p in param_dict.items() if p.requires_grad}
|
|
|
|
|
| decay_params=[p for n,p in param_dict.items() if p.dim()>=2]
|
| ndecay_params=[p for n, p in param_dict.items() if p.dim()<2]
|
| optim_grps=[
|
| {'params':decay_params,'weight_decay':weight_decay},
|
| {'params':ndecay_params,'weight_decay':0.0}
|
| ]
|
| num_decay_params=sum(p.numel() for p in decay_params)
|
| num_nodecay_params=sum(p.numel() for p in ndecay_params)
|
| print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
|
| print(f"num non-decayed parameter tensors: {len(ndecay_params)}, with {num_nodecay_params:,} parameters")
|
| fused_available='fused' in inspect.signature(torch.optim.AdamW).parameters
|
| use_fused = fused_available and device_type=='cuda'
|
| extra_args=dict(fused=True) if use_fused else dict()
|
| optimizer = torch.optim.AdamW(optim_grps,lr=learning_rate,betas=betas, **extra_args)
|
| print(f"Using Fused AdamW: {use_fused}")
|
| return optimizer
|
|
|
| @torch.no_grad()
|
| def generate(self,idx,max_new_tokens,temperature=1.0,top_k=None):
|
| for _ in range(max_new_tokens):
|
| idx_cond=idx if idx.size(1) <= self.config.block_size else idx[:,-self.config.block_size:]
|
| logits,_=self(idx_cond)
|
| logits=logits[:,-1,:]/temperature
|
| if top_k is not None:
|
| v, _ =torch.topk(logits,min(top_k,logits.size(-1)))
|
| logits[logits < v [:,[-1]]]= -float('Inf')
|
| probs=F.softmax(logits,dim=-1)
|
| idx_next=torch.multinomial(probs,num_samples=1)
|
| if idx_next.item()==50256:
|
| break
|
| idx =torch.cat((idx,idx_next),dim=1)
|
| return idx
|
|
|
|
|
