-
Notifications
You must be signed in to change notification settings - Fork 0
/
gpt2.py
427 lines (371 loc) · 19.9 KB
/
gpt2.py
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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
from dataclasses import dataclass
import torch
import torch.nn as nn
from torch.nn import functional as F
import math
import inspect
import os
import numpy as np
# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
@dataclass
class GPTConfig:
block_size: int = 1024 ## max sequence length
vocab_size: int = 50257 ## number of tokens: 50,000 BPE merges + 256 bytes tokens + 1 <|endoftext|> token. "I'll cover tokenization a little bit as well!!!"
n_layer: int = 12 ## num of layers
n_head: int = 12 ## num of heads
n_embd: int = 768 ## embedding dimension
class CasualSelfAttention(nn.Module):
def __init__(self, config) -> None:
super().__init__()
assert config.n_embd % config.n_head == 0 ## This is the general convention. The embedding is going to be build by the concatenation of the heads so it should be divisible by the n_head.
## key, query, value projections for all the heads, but in a batch
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
## output projection
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
self.c_proj.SHOULD_SCALE = True
self.n_head = config.n_head
self.n_embd = config.n_embd
## it is not bias actually, it is the mask, lower triangular matrix but in this code we're following HF/OpenAI naming conventions.
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size)) ## view is a reshape operation that reshapes our tensors.
def forward(self, x):
B, T, C = x.size() ## batch size, sequence length, embedding dimensionality (n_embd)
## nh is "number of heads", hs is "head size", C is "number of channels" = (nh * hs)
qkv = self.c_attn(x)
q, k, v = qkv.split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
## Here is the attention calculations after getting q, k, v tensors
# att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) ## query * key transpoze divided by the hs of key.
# att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf')) ## fill the parts coming after current sequence lenght with negative infinity.
# att = F.softmax(att, dim=-1) ## get the logits from the query and the key interactions and make them as probabilities.
# y = att @ v ## (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
y = y.transpose(1, 2).contiguous().view(B, T, C) ## re-assemble all head outputs side by side and reshape them.
## output projection
y = self.c_proj(y)
return y
class MLP(nn.Module):
def __init__(self, config) -> None:
super().__init__()
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
self.gelu = nn.GELU(approximate='tanh') ## GELU is better for the gradients below zero, because it's not a flat line so there is always a reason to move. No dead neuron. Approximation is for historical reasons, because of computation limits.
self.c_proj = nn.Linear(4* config.n_embd, config.n_embd)
self.c_proj.SHOULD_SCALE = True
def forward(self, x):
x = self.c_fc(x)
x = self.gelu(x)
x = self.c_proj(x)
return x
class Block(nn.Module):
def __init__(self, config) -> None:
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd) ## First layer normalization -- normalized shape
self.attn = CasualSelfAttention(config)
self.ln_2 = nn.LayerNorm(config.n_embd)
self.mlp = MLP(config) ## We could have name it as feed forward
def forward(self, x):
x = x + self.attn(self.ln_1(x)) ## We first layer normalize and make the relations being noticed by the network via attention, reduce operation
x = x + self.mlp(self.ln_2(x)) ## We then add second layer norm, and then the map operation
## Here is the interesting part we don't operate on residuals inside the gradients. Our residuals are clean and added by themselves.
return x
class GPT(nn.Module):
def __init__(self, config) -> None:
super().__init__()
self.config = config
self.transformer = nn.ModuleDict(dict(
wte = nn.Embedding(config.vocab_size, config.n_embd), ## Token embeddings -- num of embeddings, embedding dimension
wpe = nn.Embedding(config.block_size, config.n_embd), ## Positional embeddings -- num of embeddings, embedding dimension
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ## Hidden layer of the transformer block
ln_f = nn.LayerNorm(config.n_embd) ## Final Layer normalization -- it is different than original transformer
))
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) ## Final linear layer-- input dimension, output dimension
## weight sharing scheme
self.transformer.wte.weight = self.lm_head.weight
## init params
self.apply(self._init_weights)
def _init_weights(self, module):
std = 0.02
if hasattr(module, "SHOULD_SCALE"):
std *= (2 * self.config.n_layer) ** -0.5
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=std) ## 1/ sqrt(n_embed) --> Xavier initialization
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) ## In gpt article it is 0.01
def forward(self, idx, targets=None):
## idx is of shape (B, T)
B, T = idx.size()
assert T <= self.config.block_size, f"You have to use {self.config.block_size} sequence length or less"
## forward the token and position embeddings
pos = torch.arange(0, T, dtype=torch.long, device=idx.device) ## shape (T)
pos_emb = self.transformer.wpe(pos)## position embeddings of shape (T, n_embd)
tok_emb = self.transformer.wte(idx) ## token embeddings of shape (B, T, n_embd)
x = tok_emb + pos_emb
## forward the block of transformer
for block in self.transformer.h:
x = block(x)
## final layer norm and classifier
x = self.transformer.ln_f(x)
logits = self.lm_head(x) ## (B, T, vocab_size)
loss = None
if targets is not None:
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
return logits, loss
def configure_optimizers(self, weight_decay, learning_rate, device):
# start with all of the candidate parameters (that require grad)
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}
# create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.
# i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.
decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]
nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]
optim_groups = [
{'params': decay_params, 'weight_decay': weight_decay},
{'params': nodecay_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 nodecay_params)
print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")
print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters")
# Create AdamW optimizer and use the fused version if it is available
fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters
use_fused = fused_available and "cuda" in device
print(f"using fused AdamW: {use_fused}")
optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=(0.9, 0.95), eps=1e-8, fused=use_fused)
return optimizer
@classmethod
def from_pretrained(cls, model_type):
"""Loads pretrained GPT-2 model weights from huggingface"""
assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
from transformers import GPT2LMHeadModel
print("loading weights from pretrained gpt: %s" % model_type)
# n_layer, n_head and n_embd are determined from model_type
config_args = {
'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params
'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
}[model_type]
config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
# create a from-scratch initialized minGPT model
config = GPTConfig(**config_args)
model = GPT(config)
sd = model.state_dict()
sd_keys = sd.keys()
sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param
# init a huggingface/transformers model
model_hf = GPT2LMHeadModel.from_pretrained(model_type)
sd_hf = model_hf.state_dict()
# copy while ensuring all of the parameters are aligned and match in names and shapes
sd_keys_hf = sd_hf.keys()
sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer
sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer)
transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
# basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear
# this means that we have to transpose these weights when we import them
assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}"
for k in sd_keys_hf:
if any(k.endswith(w) for w in transposed):
# special treatment for the Conv1D weights we need to transpose
assert sd_hf[k].shape[::-1] == sd[k].shape
with torch.no_grad():
sd[k].copy_(sd_hf[k].t())
else:
# vanilla copy over the other parameters
assert sd_hf[k].shape == sd[k].shape
with torch.no_grad():
sd[k].copy_(sd_hf[k])
return model
def load_tokens(filename):
npt = np.load(filename)
npt = npt.astype(np.int32) # added after video
ptt = torch.tensor(npt, dtype=torch.long)
return ptt
## read files in chunks
class DataLoaderLite:
def __init__(self, B, T, split) -> None:
self.B = B
self.T = T
assert split in {'train', 'val'}
data_root = "edu_fineweb10B"
shards = os.listdir(data_root)
shards = [s for s in shards if split in s]
shards = sorted(shards)
shards = [os.path.join(data_root, s) for s in shards]
self.shards = shards
assert len(shards) > 0, f"no shards found for split {split}"
print(f"found {len(shards)} shards for split {split}")
self.reset()
##state, init at shard zero
# with open("input.txt", "r") as f:
# text = f.read()
# enc = tiktoken.get_encoding('gpt2') ## get the gpt2 tokenizer
# tokens = enc.encode(text)
# self.tokens = torch.tensor(tokens)
# print(f"loaded {len(self.tokens)} tokens!")
# print(f"1 epoch = {len(self.tokens) // (B * T)} batches")
def reset(self):
self.current_shard = 0
self.tokens = load_tokens(self.shards[self.current_shard])
self.current_position = self.B * self.T
def next_batch(self):
B, T = self.B, self.T
buf = self.tokens[self.current_position : self.current_position+B*T+1]
x = (buf[:-1].view(B, T))
y = (buf[1:].view(B, T))
## Advance the position in the tensor
self.current_position += B * T
## if loading the next batch would be out of bounds, advance to next shard
if self.current_position + (B * T + 1) > len(self.tokens):
self.current_shard = (self.current_shard + 1) % len(self.shards)
self.tokens = load_tokens(self.shards[self.current_shard])
self.current_position = 0
return x, y
if __name__ == "__main__":
device = "cpu"
if torch.cuda.is_available():
device = "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
device = "mps"
print(f"using device: {device}")
torch.manual_seed(1337)
if torch.cuda.is_available():
torch.cuda.manual_seed(1337)
total_batch_size = 131072 ##--> 2**17 ##524288 ## 2**19, ~0.5M, in number of tokens
## 10 GB of HBM GPU usage!!!! watch -n 0.5 nvidia-smi
B, T = 8, 1024 ## B: micro batch size, T: sequence length
assert total_batch_size % (B * T) == 0, "The total batch size should be divisible by the B * T"
grad_accum_steps = total_batch_size // (B * T)
print(f"total desired batch size: {total_batch_size}")
print(f"=> calculated gradient accumulation steps: {grad_accum_steps}")
train_loader = DataLoaderLite(B=B, T=T, split="train")
val_loader = DataLoaderLite(B=B, T=T, split="val")
torch.set_float32_matmul_precision('high') ## Use TF32 instead of FP32
## model = GPT.from_pretrained('gpt2')
model = GPT(GPTConfig(vocab_size=50304))
## model.eval() ## This means we won't be in the training stage.
model.to(device)
model = torch.compile(model) ## this finds most of the operations and compiles them before the python interpreter.
# tokens = enc.encode("Hello, I'm a language model,") ## tokenize the given sentence with gpt2 tokenizer algorithm
# text = text[:1000]
# tokens = enc.encode(text)
## tokens = torch.tensor(tokens, dtype=torch.long) ## (8,) change this tokens to torch tensor
## tokens = tokens.unsqueeze(0).repeat(num_of_return_sequences, 1) ## (5, 8) add another dimension and repeat the same values for num_of_return_sequences times.
## These tokens are the idx in our GPT class forward method. The parameter.
## Here we are creating a buffer that contains one more token elements to use as label.
# buf = torch.tensor(tokens[:B*T + 1], device=device)
# x = buf[:-1].view(B, T)
# y = buf[1:].view(B, T)
# x = tokens.to('cuda')
# num_of_return_sequences = 5
# max_length = 30
## logits, loss = model(x, y)
max_lr = 6e-4
min_lr = max_lr * 0.1
warmup_steps = 750 ## 375e6 // 131072 or ##375e6 // 2 **19 but I choosen something from on top of my head.
max_steps = 4577 ## 6e8 // 131072 or ## 10e9 // 2 ** 19
def get_lr(it):
## linear warmup for warmup_iters steps
if it < warmup_steps:
return max_lr * (it + 1) / warmup_steps
## if it > lr_decay_iters, return min learning rate
if it > max_steps:
return min_lr
## in between, use cosine decay down to min learning rate
decay_ratio = (it - warmup_steps) / (max_steps - warmup_steps)
assert 0 <= decay_ratio <= 1
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) ## coeff starts at 1 and goes to 0
return min_lr + coeff * (max_lr - min_lr)
import time
## optimizer
## optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95), eps=1e-8)
optimizer = model.configure_optimizers(weight_decay=0.1, learning_rate=6e-4, device=device)
log_dir = "log"
os.makedirs(log_dir, exist_ok=True)
log_file = os.path.join(log_dir, f"log.txt")
with open(log_file, "w") as f:
pass
for step in range(max_steps):
t0 = time.time()
last_step = (step == max_steps - 1)
## evaluation steps in every 101 iteration
if step % 250 == 0 or last_step:
model.eval()
val_loader.reset()
with torch.no_grad():
val_loss_accum = 0.0
val_loss_steps = 20
for _ in range(val_loss_steps):
x, y = val_loader.next_batch()
x, y = x.to(device), y.to(device)
with torch.autocast(device_type=device, dtype=torch.bfloat16):
logits, loss = model(x, y)
loss = loss / val_loss_steps
val_loss_accum += loss.detach()
print(f"validation loss: {val_loss_accum.item():4f}")
with open(log_file, "a") as f:
f.write(f"{step} val {val_loss_accum.item():.4f}\n")
if step > 0 and (step % 250 == 0 or last_step):
# optionally write model checkpoints
checkpoint_path = os.path.join(log_dir, f"model_{step:05d}.pt")
checkpoint = {
'model': model.state_dict(),
'config': model.config,
'step': step,
'val_loss': val_loss_accum.item(),
'optimizer': optimizer.state_dict()
}
torch.save(checkpoint, checkpoint_path)
model.train()
optimizer.zero_grad()
loss_accum = 0.0
for micro_step in range(grad_accum_steps):
x, y = train_loader.next_batch()
x, y = x.to(device), y.to(device)
with torch.autocast(device_type=device, dtype=torch.bfloat16): ## Now we use bfloat16 as torch recommends, this causes mixed precision but still the weights are float32. The documentation shows what gets converted
logits, loss = model(x, y)
loss = loss / grad_accum_steps
loss_accum += loss.detach()
loss.backward()
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) ## gradient clipping.
lr = get_lr(step)
for param_group in optimizer.param_groups:
param_group["lr"] = lr
optimizer.step()
torch.cuda.synchronize()
t1 = time.time()
dt = (t1 - t0) * 1000
tokens_processed = train_loader.B * train_loader.T * grad_accum_steps
tokens_per_second = tokens_processed / dt
print(f"step {step:4d} | loss: {loss_accum.item():.6f} | lr {lr:.4e} | norm: {norm:.4f} | dt: {dt:.2f}ms | tok/sec: {tokens_per_second:.4f}")
with open(log_file, "a") as f:
f.write(f"{step} train {loss_accum.item():.6f}\n")
import sys; sys.exit(0)
torch.manual_seed(42)
torch.cuda.manual_seed(42)
while x.size(1) < max_length:
## forward the model to get the logits
with torch.no_grad():
logits = model(x) ## (B, T, vocab_size)
## Take the logits' last position item
logits = logits[:, -1, :] ## (B, vocab_size)
## Get the probabilities
probs = F.softmax(logits, dim=-1)
## We'll do the top k sampling here (HF's default 50 for pipeline)
## (5, 50)
topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
## Select a token from top-k probabilities
ix = torch.multinomial(topk_probs, 1) ## (B, 1)
## gather the corresponding indices
xcol = torch.gather(topk_indices, -1, ix) ## (B, 1)
## append to the sequence to get the full generated sentences
x = torch.cat((x, xcol), dim=1)
## decode and print the generated text
for i in range(num_of_return_sequences):
tokens = x[i, :max_length].tolist()
decoded = enc.decode(tokens)
print(">", decoded)