1.
import os
2. import math
3. import time
4. import inspect
5. from dataclasses import dataclass
6. import torch
7. import [Link] as nn
8. from [Link] import functional as F
9. from hellaswag import render_example, iterate_examples
10. # -----------------------------------------------------------------------------
11.
12. class CausalSelfAttention([Link]):
13.
14. def __init__(self, config):
15. super().__init__()
16. assert config.n_embd % config.n_head == 0
17. # key, query, value projections for all heads, but in a batch
18. self.c_attn = [Link](config.n_embd, 3 * config.n_embd)
19. # output projection
20. self.c_proj = [Link](config.n_embd, config.n_embd)
21. self.c_proj.NANOGPT_SCALE_INIT = 1
22. # regularization
23. self.n_head = config.n_head
24. self.n_embd = config.n_embd
25.
26. def forward(self, x):
27. B, T, C = [Link]() # batch size, sequence length, embedding dimensionality
(n_embd)
28. # calculate query, key, values for all heads in batch and move head forward
to be the batch dim
29. # nh is "number of heads", hs is "head size", and C (number of channels) =
nh * hs
30. # e.g. in GPT-2 (124M), n_head=12, hs=64, so nh*hs=C=768 channels in the
Transformer
31. qkv = self.c_attn(x)
32. q, k, v = [Link](self.n_embd, dim=2)
33. k = [Link](B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T,
hs)
34. q = [Link](B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T,
hs)
35. v = [Link](B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T,
hs)
36. y = F.scaled_dot_product_attention(q, k, v, is_causal=True) # flash
attention
37. y = [Link](1, 2).contiguous().view(B, T, C) # re-assemble all head
outputs side by side
38. # output projection
39. y = self.c_proj(y)
40. return y
41.
42. class MLP([Link]):
43.
44. def __init__(self, config):
45. super().__init__()
46. self.c_fc = [Link](config.n_embd, 4 * config.n_embd)
47. [Link] = [Link](approximate='tanh')
48. self.c_proj = [Link](4 * config.n_embd, config.n_embd)
49. self.c_proj.NANOGPT_SCALE_INIT = 1
50.
51. def forward(self, x):
52. x = self.c_fc(x)
53. x = [Link](x)
54. x = self.c_proj(x)
55. return x
56.
57. class Block([Link]):
58.
59. def __init__(self, config):
60. super().__init__()
61. self.ln_1 = [Link](config.n_embd)
62. [Link] = CausalSelfAttention(config)
63. self.ln_2 = [Link](config.n_embd)
64. [Link] = MLP(config)
65.
66. def forward(self, x):
67. x = x + [Link](self.ln_1(x))
68. x = x + [Link](self.ln_2(x))
69. return x
70.
71. @dataclass
72. class GPTConfig:
73. block_size: int = 1024 # max sequence length
74. vocab_size: int = 50257 # number of tokens: 50,000 BPE merges + 256 bytes tokens
+ 1 <|endoftext|> token
75. n_layer: int = 12 # number of layers
76. n_head: int = 12 # number of heads
77. n_embd: int = 768 # embedding dimension
78.
79. class GPT([Link]):
80.
81. def __init__(self, config):
82. super().__init__()
83. [Link] = config
84.
85. [Link] = [Link](dict(
86. wte = [Link](config.vocab_size, config.n_embd),
87. wpe = [Link](config.block_size, config.n_embd),
88. h = [Link]([Block(config) for _ in range(config.n_layer)]),
89. ln_f = [Link](config.n_embd),
90. ))
91. self.lm_head = [Link](config.n_embd, config.vocab_size, bias=False)
92.
93. # weight sharing scheme
94. [Link] = self.lm_head.weight
95.
96. # init params
97. [Link](self._init_weights)
98.
99. def _init_weights(self, module):
100. if isinstance(module, [Link]):
101. std = 0.02
102. if hasattr(module, 'NANOGPT_SCALE_INIT'):
103. std *= (2 * [Link].n_layer) ** -0.5
104. [Link].normal_([Link], mean=0.0, std=std)
105. if [Link] is not None:
106. [Link].zeros_([Link])
107. elif isinstance(module, [Link]):
108. [Link].normal_([Link], mean=0.0, std=0.02)
109.
110. def forward(self, idx, targets=None):
111. # idx is of shape (B, T)
112. B, T = [Link]()
113. assert T <= [Link].block_size, f"Cannot forward sequence of length {T},
block size is only {[Link].block_size}"
114. # forward the token and posisition embeddings
115. pos = [Link](0, T, dtype=[Link], device=[Link]) # shape (T)
116. pos_emb = [Link](pos) # position embeddings of shape (T,
n_embd)
117. tok_emb = [Link](idx) # token embeddings of shape (B, T,
n_embd)
118. x = tok_emb + pos_emb
119. # forward the blocks of the transformer
120. for block in [Link].h:
121. x = block(x)
122. # forward the final layernorm and the classifier
123. x = [Link].ln_f(x)
124. logits = self.lm_head(x) # (B, T, vocab_size)
125. loss = None
126. if targets is not None:
127. loss = F.cross_entropy([Link](-1, [Link](-1)), [Link](-
1))
128. return logits, loss
129.
130. @classmethod
131. def from_pretrained(cls, model_type):
132. """Loads pretrained GPT-2 model weights from huggingface"""
133. assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
134. from transformers import GPT2LMHeadModel
135. print("loading weights from pretrained gpt: %s" % model_type)
136.
137. # n_layer, n_head and n_embd are determined from model_type
138. config_args = {
139. 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params
140. 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params
141. 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params
142. 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params
143. }[model_type]
144. config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints
145. config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints
146. # create a from-scratch initialized minGPT model
147. config = GPTConfig(**config_args)
148. model = GPT(config)
149. sd = model.state_dict()
150. sd_keys = [Link]()
151. sd_keys = [k for k in sd_keys if not [Link]('.[Link]')] # discard
this mask / buffer, not a param
152.
153. # init a huggingface/transformers model
154. model_hf = GPT2LMHeadModel.from_pretrained(model_type)
155. sd_hf = model_hf.state_dict()
156.
157. # copy while ensuring all of the parameters are aligned and match in names
and shapes
158. sd_keys_hf = sd_hf.keys()
159. sd_keys_hf = [k for k in sd_keys_hf if not [Link]('.attn.masked_bias')]
# ignore these, just a buffer
160. sd_keys_hf = [k for k in sd_keys_hf if not [Link]('.[Link]')] # same,
just the mask (buffer)
161. transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight',
'mlp.c_proj.weight']
162. # basically the openai checkpoints use a "Conv1D" module, but we only want
to use a vanilla Linear
163. # this means that we have to transpose these weights when we import them
164. assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)}
!= {len(sd_keys)}"
165. for k in sd_keys_hf:
166. if any([Link](w) for w in transposed):
167. # special treatment for the Conv1D weights we need to transpose
168. assert sd_hf[k].shape[::-1] == sd[k].shape
169. with torch.no_grad():
170. sd[k].copy_(sd_hf[k].t())
171. else:
172. # vanilla copy over the other parameters
173. assert sd_hf[k].shape == sd[k].shape
174. with torch.no_grad():
175. sd[k].copy_(sd_hf[k])
176.
177. return model
178.
179. def configure_optimizers(self, weight_decay, learning_rate, device_type):
180. # start with all of the candidate parameters (that require grad)
181. param_dict = {pn: p for pn, p in self.named_parameters()}
182. param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}
183. # create optim groups. Any parameters that is 2D will be weight decayed,
otherwise no.
184. # i.e. all weight tensors in matmuls + embeddings decay, all biases and
layernorms don't.
185. decay_params = [p for n, p in param_dict.items() if [Link]() >= 2]
186. nodecay_params = [p for n, p in param_dict.items() if [Link]() < 2]
187. optim_groups = [
188. {'params': decay_params, 'weight_decay': weight_decay},
189. {'params': nodecay_params, 'weight_decay': 0.0}
190. ]
191. num_decay_params = sum([Link]() for p in decay_params)
192. num_nodecay_params = sum([Link]() for p in nodecay_params)
193. if master_process:
194. print(f"num decayed parameter tensors: {len(decay_params)}, with
{num_decay_params:,} parameters")
195. print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with
{num_nodecay_params:,} parameters")
196. # Create AdamW optimizer and use the fused version if it is available
197. fused_available = 'fused' in [Link]([Link]).parameters
198. use_fused = fused_available and device_type == "cuda"
199. if master_process:
200. print(f"using fused AdamW: {use_fused}")
201. optimizer = [Link](optim_groups, lr=learning_rate, betas=(0.9,
0.95), eps=1e-8, fused=use_fused)
202. return optimizer
203.
204. # -----------------------------------------------------------------------------
205. import tiktoken
206. import numpy as np
207.
208. def load_tokens(filename):
209. npt = [Link](filename)
210. npt = [Link](np.int32) # added after video
211. ptt = [Link](npt, dtype=[Link])
212. return ptt
213.
214. class DataLoaderLite:
215. def __init__(self, B, T, process_rank, num_processes, split):
216. self.B = B
217. self.T = T
218. self.process_rank = process_rank
219. self.num_processes = num_processes
220. assert split in {'train', 'val'}
221.
222. # get the shard filenames
223. data_root = "edu_fineweb10B"
224. shards = [Link](data_root)
225. shards = [s for s in shards if split in s]
226. shards = sorted(shards)
227. shards = [[Link](data_root, s) for s in shards]
228. [Link] = shards
229. assert len(shards) > 0, f"no shards found for split {split}"
230. if master_process:
231. print(f"found {len(shards)} shards for split {split}")
232. [Link]()
233.
234. def reset(self):
235. # state, init at shard zero
236. self.current_shard = 0
237. [Link] = load_tokens([Link][self.current_shard])
238. self.current_position = self.B * self.T * self.process_rank
239.
240. def next_batch(self):
241. B, T = self.B, self.T
242. buf = [Link][self.current_position : self.current_position+B*T+1]
243. x = (buf[:-1]).view(B, T) # inputs
244. y = (buf[1:]).view(B, T) # targets
245. # advance the position in the tensor
246. self.current_position += B * T * self.num_processes
247. # if loading the next batch would be out of bounds, advance to next shard
248. if self.current_position + (B * T * self.num_processes + 1) >
len([Link]):
249. self.current_shard = (self.current_shard + 1) % len([Link])
250. [Link] = load_tokens([Link][self.current_shard])
251. self.current_position = B * T * self.process_rank
252. return x, y
253.
254. # -----------------------------------------------------------------------------
255. # helper function for HellaSwag eval
256. # takes tokens, mask, and logits, returns the index of the completion with the
lowest loss
257.
258. def get_most_likely_row(tokens, mask, logits):
259. # evaluate the autoregressive loss at all positions
260. shift_logits = (logits[..., :-1, :]).contiguous()
261. shift_tokens = (tokens[..., 1:]).contiguous()
262. flat_shift_logits = shift_logits.view(-1, shift_logits.size(-1))
263. flat_shift_tokens = shift_tokens.view(-1)
264. shift_losses = F.cross_entropy(flat_shift_logits, flat_shift_tokens,
reduction='none')
265. shift_losses = shift_losses.view([Link](0), -1)
266. # now get the average loss just for the completion region (where mask == 1), in
each row
267. shift_mask = (mask[..., 1:]).contiguous() # we must shift mask, so we start at
the last prompt token
268. masked_shift_losses = shift_losses * shift_mask
269. # sum and divide by the number of 1s in the mask
270. sum_loss = masked_shift_losses.sum(dim=1)
271. avg_loss = sum_loss / shift_mask.sum(dim=1)
272. # now we have a loss for each of the 4 completions
273. # the one with the lowest loss should be the most likely
274. pred_norm = avg_loss.argmin().item()
275. return pred_norm
276.
277. # -----------------------------------------------------------------------------
278. # simple launch:
279. # python train_gpt2.py
280. # DDP launch for e.g. 8 GPUs:
281. # torchrun --standalone --nproc_per_node=8 train_gpt2.py
282.
283. # run the training loop
284. from [Link] import init_process_group, destroy_process_group
285. from [Link] import DistributedDataParallel as DDP
286. import [Link] as dist
287.
288. # set up DDP (distributed data parallel).
289. # torchrun command sets the env variables RANK, LOCAL_RANK, and WORLD_SIZE
290. ddp = int([Link]('RANK', -1)) != -1 # is this a ddp run?
291. if ddp:
292. # use of DDP atm demands CUDA, we set the device appropriately according to rank
293. assert [Link].is_available(), "for now i think we need CUDA for DDP"
294. init_process_group(backend='nccl')
295. ddp_rank = int([Link]['RANK'])
296. ddp_local_rank = int([Link]['LOCAL_RANK'])
297. ddp_world_size = int([Link]['WORLD_SIZE'])
298. device = f'cuda:{ddp_local_rank}'
299. [Link].set_device(device)
300. master_process = ddp_rank == 0 # this process will do logging, checkpointing
etc.
301. else:
302. # vanilla, non-DDP run
303. ddp_rank = 0
304. ddp_local_rank = 0
305. ddp_world_size = 1
306. master_process = True
307. # attempt to autodetect device
308. device = "cpu"
309. if [Link].is_available():
310. device = "cuda"
311. elif hasattr([Link], "mps") and [Link].is_available():
312. device = "mps"
313. print(f"using device: {device}")
314.
315. # added after video, pytorch can be serious about it's device vs. device_type
distinction
316. device_type = "cuda" if [Link]("cuda") else "cpu"
317.
318. torch.manual_seed(1337)
319. if [Link].is_available():
320. [Link].manual_seed(1337)
321.
322. enc = tiktoken.get_encoding("gpt2")
323.
324. total_batch_size = 524288 # 2**19, ~0.5M, in number of tokens
325. B = 64 # micro batch size
326. T = 1024 # sequence length
327. assert total_batch_size % (B * T * ddp_world_size) == 0, "make sure total_batch_size
is divisible by B * T * ddp_world_size"
328. grad_accum_steps = total_batch_size // (B * T * ddp_world_size)
329. if master_process:
330. print(f"total desired batch size: {total_batch_size}")
331. print(f"=> calculated gradient accumulation steps: {grad_accum_steps}")
332.
333. train_loader = DataLoaderLite(B=B, T=T, process_rank=ddp_rank,
num_processes=ddp_world_size, split="train")
334. val_loader = DataLoaderLite(B=B, T=T, process_rank=ddp_rank,
num_processes=ddp_world_size, split="val")
335.
336. torch.set_float32_matmul_precision('high')
337.
338. # create model
339. model = GPT(GPTConfig(vocab_size=50304))
340. # model = GPT.from_pretrained("gpt2") # or init from OpenAI GPT-2
341. [Link](device)
342. use_compile = False # [Link] interferes with HellaSwag eval and Generation.
TODO fix
343. if use_compile:
344. model = [Link](model)
345. if ddp:
346. model = DDP(model, device_ids=[ddp_local_rank])
347. raw_model = [Link] if ddp else model # always contains the "raw" unwrapped
model
348.
349. max_lr = 6e-4
350. min_lr = max_lr * 0.1
351. warmup_steps = 715
352. max_steps = 19073 # 19,073 steps is ~1 epoch, if data is 10B tokens and batch size
0.5M tokens
353. def get_lr(it):
354. # 1) linear warmup for warmup_iters steps
355. if it < warmup_steps:
356. return max_lr * (it+1) / warmup_steps
357. # 2) if it > lr_decay_iters, return min learning rate
358. if it > max_steps:
359. return min_lr
360. # 3) in between, use cosine decay down to min learning rate
361. decay_ratio = (it - warmup_steps) / (max_steps - warmup_steps)
362. assert 0 <= decay_ratio <= 1
363. coeff = 0.5 * (1.0 + [Link]([Link] * decay_ratio)) # coeff starts at 1 and
goes to 0
364. return min_lr + coeff * (max_lr - min_lr)
365.
366. # optimize!
367. optimizer = raw_model.configure_optimizers(weight_decay=0.1, learning_rate=6e-4,
device_type=device_type)
368.
369. # create the log directory we will write checkpoints to and log to
370. log_dir = "log"
371. [Link](log_dir, exist_ok=True)
372. log_file = [Link](log_dir, f"[Link]")
373. with open(log_file, "w") as f: # open for writing to clear the file
374. pass
375.
376. for step in range(max_steps):
377. t0 = [Link]()
378. last_step = (step == max_steps - 1)
379.
380. # once in a while evaluate our validation loss
381. if step % 250 == 0 or last_step:
382. [Link]()
383. val_loader.reset()
384. with torch.no_grad():
385. val_loss_accum = 0.0
386. val_loss_steps = 20
387. for _ in range(val_loss_steps):
388. x, y = val_loader.next_batch()
389. x, y = [Link](device), [Link](device)
390. with [Link](device_type=device_type, dtype=torch.bfloat16):
391. logits, loss = model(x, y)
392. loss = loss / val_loss_steps
393. val_loss_accum += [Link]()
394. if ddp:
395. dist.all_reduce(val_loss_accum, op=[Link])
396. if master_process:
397. print(f"validation loss: {val_loss_accum.item():.4f}")
398. with open(log_file, "a") as f:
399. [Link](f"{step} val {val_loss_accum.item():.4f}\n")
400. if step > 0 and (step % 5000 == 0 or last_step):
401. # optionally write model checkpoints
402. checkpoint_path = [Link](log_dir, f"model_{step:05d}.pt")
403. checkpoint = {
404. 'model': raw_model.state_dict(),
405. 'config': raw_model.config,
406. 'step': step,
407. 'val_loss': val_loss_accum.item()
408. }
409. # you might also want to add optimizer.state_dict() and
410. # rng seeds etc., if you wanted to more exactly resume training
411. [Link](checkpoint, checkpoint_path)
412.
413. # once in a while evaluate hellaswag
414. if (step % 250 == 0 or last_step) and (not use_compile):
415. num_correct_norm = 0
416. num_total = 0
417. for i, example in enumerate(iterate_examples("val")):
418. # only process examples where i % ddp_world_size == ddp_rank
419. if i % ddp_world_size != ddp_rank:
420. continue
421. # render the example into tokens and labels
422. _, tokens, mask, label = render_example(example)
423. tokens = [Link](device)
424. mask = [Link](device)
425. # get the logits
426. with torch.no_grad():
427. with [Link](device_type=device_type, dtype=torch.bfloat16):
428. logits, loss = model(tokens)
429. pred_norm = get_most_likely_row(tokens, mask, logits)
430. num_total += 1
431. num_correct_norm += int(pred_norm == label)
432. # reduce the stats across all processes
433. if ddp:
434. num_total = [Link](num_total, dtype=[Link], device=device)
435. num_correct_norm = [Link](num_correct_norm, dtype=[Link],
device=device)
436. dist.all_reduce(num_total, op=[Link])
437. dist.all_reduce(num_correct_norm, op=[Link])
438. num_total = num_total.item()
439. num_correct_norm = num_correct_norm.item()
440. acc_norm = num_correct_norm / num_total
441. if master_process:
442. print(f"HellaSwag accuracy:
{num_correct_norm}/{num_total}={acc_norm:.4f}")
443. with open(log_file, "a") as f:
444. [Link](f"{step} hella {acc_norm:.4f}\n")
445.
446. # once in a while generate from the model (except step 0, which is noise)
447. if ((step > 0 and step % 250 == 0) or last_step) and (not use_compile):
448. [Link]()
449. num_return_sequences = 4
450. max_length = 32
451. tokens = [Link]("Hello, I'm a language model,")
452. tokens = [Link](tokens, dtype=[Link])
453. tokens = [Link](0).repeat(num_return_sequences, 1)
454. xgen = [Link](device)
455. sample_rng = [Link](device=device)
456. sample_rng.manual_seed(42 + ddp_rank)
457. while [Link](1) < max_length:
458. # forward the model to get the logits
459. with torch.no_grad():
460. with [Link](device_type=device_type, dtype=torch.bfloat16):
461. logits, loss = model(xgen) # (B, T, vocab_size)
462. # take the logits at the last position
463. logits = logits[:, -1, :] # (B, vocab_size)
464. # get the probabilities
465. probs = [Link](logits, dim=-1)
466. # do top-k sampling of 50 (huggingface pipeline default)
467. # topk_probs here becomes (5, 50), topk_indices is (5, 50)
468. topk_probs, topk_indices = [Link](probs, 50, dim=-1)
469. # select a token from the top-k probabilities
470. # note: multinomial does not demand the input to sum to 1
471. ix = [Link](topk_probs, 1, generator=sample_rng) # (B, 1)
472. # gather the corresponding indices
473. xcol = [Link](topk_indices, -1, ix) # (B, 1)
474. # append to the sequence
475. xgen = [Link]((xgen, xcol), dim=1)
476. # print the generated text
477. for i in range(num_return_sequences):
478. tokens = xgen[i, :max_length].tolist()
479. decoded = [Link](tokens)
480. print(f"rank {ddp_rank} sample {i}: {decoded}")
481.
482. # do one step of the optimization
483. [Link]()
484. optimizer.zero_grad()
485. loss_accum = 0.0
486. for micro_step in range(grad_accum_steps):
487. x, y = train_loader.next_batch()
488. x, y = [Link](device), [Link](device)
489. # added after video, this field is also used by the forward pass.
490. if ddp:
491. model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1)
492. with [Link](device_type=device_type, dtype=torch.bfloat16):
493. logits, loss = model(x, y)
494. # we have to scale the loss to account for gradient accumulation,
495. # because the gradients just add on each successive backward().
496. # addition of gradients corresponds to a SUM in the objective, but
497. # instead of a SUM we want MEAN. Scale the loss here so it comes out right
498. loss = loss / grad_accum_steps
499. loss_accum += [Link]()
500. [Link]()
501. if ddp:
502. dist.all_reduce(loss_accum, op=[Link])
503. norm = [Link].clip_grad_norm_([Link](), 1.0)
504. # determine and set the learning rate for this iteration
505. lr = get_lr(step)
506. for param_group in optimizer.param_groups:
507. param_group['lr'] = lr
508. [Link]()
509. if device_type == "cuda":
510. [Link]() # wait for the GPU to finish work
511. t1 = [Link]()
512. dt = t1 - t0 # time difference in seconds
513. tokens_processed = train_loader.B * train_loader.T * grad_accum_steps *
ddp_world_size
514. tokens_per_sec = tokens_processed / dt
515. if master_process:
516. print(f"step {step:5d} | loss: {loss_accum.item():.6f} | lr {lr:.4e} | norm:
{norm:.4f} | dt: {dt*1000:.2f}ms | tok/sec: {tokens_per_sec:.2f}")
517. with open(log_file, "a") as f:
518. [Link](f"{step} train {loss_accum.item():.6f}\n")
519.
520. if ddp:
521. destroy_process_group()
522.