import%20marimo%0A%0A__generated_with%20%3D%20%220.25.0%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20GPT-2%20under%20the%20hood%0A%0A%20%20%20%20**Lecture%204A%20companion.**%20We%20open%20GPT-2%20small%20(124%20M%20parameters%2C%202019)%2C%20the%20last%20large%20OpenAI%20model%0A%20%20%20%20with%20public%20weights%20and%20still%20a%20faithful%20miniature%20of%20today's%20chat%20models%2C%20and%20follow%20a%20prompt%20through%20it%3A%0A%20%20%20%20tokens%20%E2%86%92%20embeddings%20%E2%86%92%20attention%20%E2%86%92%20residual%20stream%20%E2%86%92%20next-token%20probabilities.%20At%20the%20end%20we%20write%20a%0A%20%20%20%20tiny%20GPT%20from%20scratch%20and%20train%20it%20on%20the%20text%20of%20this%20course.%0A%0A%20%20%20%20Install%20once%20(Python%203.10%2B)%2C%20from%20the%20%60Lecture_4A%60%20folder%3A%0A%0A%20%20%20%20%60%60%60bash%0A%20%20%20%20python3.12%20-m%20venv%20.venv%0A%20%20%20%20.venv%2Fbin%2Fpython%20-m%20pip%20install%20-r%20requirements.txt%0A%20%20%20%20.venv%2Fbin%2Fmarimo%20edit%20GPT2_under_the_hood.py%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20The%20first%20run%20downloads%20GPT-2%20(~550%20MB)%20from%20the%20Hugging%20Face%20hub.%20Everything%20runs%20on%20a%20laptop%20CPU.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20time%0A%20%20%20%20from%20pathlib%20import%20Path%0A%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20import%20torch%0A%20%20%20%20import%20torch.nn%20as%20nn%0A%20%20%20%20import%20torch.nn.functional%20as%20F%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20tiktoken%0A%20%20%20%20from%20transformers%20import%20GPT2LMHeadModel%2C%20GPT2TokenizerFast%0A%0A%20%20%20%20torch.set_grad_enabled(False)%0A%20%20%20%20plt.rcParams.update(%7B%22figure.dpi%22%3A%20110%2C%20%22font.size%22%3A%2010%7D)%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20F%2C%0A%20%20%20%20%20%20%20%20GPT2LMHeadModel%2C%0A%20%20%20%20%20%20%20%20GPT2TokenizerFast%2C%0A%20%20%20%20%20%20%20%20Path%2C%0A%20%20%20%20%20%20%20%20nn%2C%0A%20%20%20%20%20%20%20%20np%2C%0A%20%20%20%20%20%20%20%20plt%2C%0A%20%20%20%20%20%20%20%20tiktoken%2C%0A%20%20%20%20%20%20%20%20time%2C%0A%20%20%20%20%20%20%20%20torch%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(GPT2LMHeadModel%2C%20GPT2TokenizerFast)%3A%0A%20%20%20%20model%20%3D%20GPT2LMHeadModel.from_pretrained(%22gpt2%22%2C%20attn_implementation%3D%22eager%22).eval()%0A%20%20%20%20tok%20%3D%20GPT2TokenizerFast.from_pretrained(%22gpt2%22)%0A%20%20%20%20BOS%20%3D%20tok.eos_token_id%0A%20%20%20%20print(model)%0A%20%20%20%20print(f%22%7Bsum(p.numel()%20for%20p%20in%20model.parameters())%3A%2C%7D%20parameters%22)%0A%20%20%20%20return%20BOS%2C%20model%2C%20tok%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%201.%20Tokenization%0A%0A%20%20%20%20Type%20anything.%20Each%20coloured%20chip%20is%20one%20token%3B%20the%20number%20under%20it%20is%20its%20id.%20Compare%20GPT-2's%0A%20%20%20%2050%20257-token%20vocabulary%20with%20the%20100%20k%20and%20200%20k%20vocabularies%20of%20later%20OpenAI%20models%0A%20%20%20%20(%60%C2%B7%60%20marks%20a%20space%2C%20which%20belongs%20to%20the%20token%20that%20follows%20it).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20tok_text%20%3D%20mo.ui.text_area(%0A%20%20%20%20%20%20%20%20value%3D%22The%20electron%20has%20a%20negative%20charge.%20Das%20Elektron%20ist%20negativ%20geladen.%20%D0%AD%D0%BB%D0%B5%D0%BA%D1%82%D1%80%D0%BE%D0%BD%20%D0%B7%D0%B0%D1%80%D1%8F%D0%B6%D0%B5%D0%BD%20%D0%BE%D1%82%D1%80%D0%B8%D1%86%D0%B0%D1%82%D0%B5%D0%BB%D1%8C%D0%BD%D0%BE.%2012345678%20strawberry%22%2C%0A%20%20%20%20%20%20%20%20full_width%3DTrue%2C%20rows%3D3)%0A%20%20%20%20tok_text%0A%20%20%20%20return%20(tok_text%2C)%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tiktoken%2C%20tok_text)%3A%0A%20%20%20%20def%20chips(enc_name)%3A%0A%20%20%20%20%20%20%20%20enc%20%3D%20tiktoken.get_encoding(enc_name)%0A%20%20%20%20%20%20%20%20ids%20%3D%20enc.encode(tok_text.value)%0A%20%20%20%20%20%20%20%20spans%20%3D%20%5B%5D%0A%20%20%20%20%20%20%20%20for%20i%20in%20ids%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20s%20%3D%20enc.decode(%5Bi%5D).replace(%22%20%22%2C%20%22%C2%B7%22).replace(%22%5Cn%22%2C%20%22%E2%86%B5%22).replace(%22%3C%22%2C%20%22%26lt%3B%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20hue%20%3D%20(i%20*%2047)%20%25%20360%0A%20%20%20%20%20%20%20%20%20%20%20%20spans.append(f%22%3Cspan%20style%3D'display%3Ainline-block%3Bmargin%3A2px%3Bpadding%3A1px%204px%3Bborder-radius%3A4px%3B%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22background%3Ahsl(%7Bhue%7D%2070%25%2088%25)%3Bfont-family%3Amonospace'%3E%7Bs%7D%3Cbr%3E%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22%3Csmall%20style%3D'color%3A%23617484'%3E%7Bi%7D%3C%2Fsmall%3E%3C%2Fspan%3E%22)%0A%20%20%20%20%20%20%20%20return%20mo.vstack(%5Bmo.md(f%22**%7Benc_name%7D**%20%E2%80%94%20%7Blen(ids)%7D%20tokens%20for%20%7Blen(tok_text.value)%7D%20characters%22)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mo.Html(%22%22.join(spans))%5D)%0A%0A%20%20%20%20mo.vstack(%5Bchips(%22gpt2%22)%2C%20chips(%22cl100k_base%22)%2C%20chips(%22o200k_base%22)%5D)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%202.%20The%20next%20token%0A%0A%20%20%20%20The%20model%20returns%20one%20score%20(logit)%20for%20each%20of%20the%2050%20257%20tokens.%20Softmax%20with%20temperature%20%24T%24%20turns%0A%20%20%20%20them%20into%20probabilities%3B%20generation%20samples%20a%20token%2C%20appends%20it%2C%20and%20repeats.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20prompt%20%3D%20mo.ui.text(value%3D%22The%20electron%20has%20a%20negative%22%2C%20full_width%3DTrue%2C%20label%3D%22Prompt%22)%0A%20%20%20%20temperature%20%3D%20mo.ui.slider(0.0%2C%202.0%2C%20step%3D0.05%2C%20value%3D0.8%2C%20label%3D%22Temperature%22%2C%20show_value%3DTrue)%0A%20%20%20%20n_new%20%3D%20mo.ui.slider(5%2C%2080%2C%20step%3D5%2C%20value%3D40%2C%20label%3D%22New%20tokens%22%2C%20show_value%3DTrue)%0A%20%20%20%20generate%20%3D%20mo.ui.run_button(label%3D%22Generate%22)%0A%20%20%20%20mo.vstack(%5Bprompt%2C%20mo.hstack(%5Btemperature%2C%20n_new%2C%20generate%5D%2C%20justify%3D%22start%22%2C%20gap%3D2)%5D)%0A%20%20%20%20return%20generate%2C%20n_new%2C%20prompt%2C%20temperature%0A%0A%0A%40app.cell%0Adef%20_(model%2C%20plt%2C%20prompt%2C%20temperature%2C%20tok%2C%20torch)%3A%0A%20%20%20%20ids%20%3D%20tok.encode(prompt.value)%0A%20%20%20%20logits%20%3D%20model(torch.tensor(%5Bids%5D)).logits%5B0%2C%20-1%5D%0A%20%20%20%20_T%20%3D%20max(temperature.value%2C%201e-4)%0A%20%20%20%20probs%20%3D%20(logits%20%2F%20_T).softmax(-1)%0A%20%20%20%20_v%2C%20_i%20%3D%20probs.topk(15)%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(7%2C%203.8))%0A%20%20%20%20_ax.barh(%5Brepr(tok.decode(%5Bint(j)%5D))%20for%20j%20in%20_i%5D%5B%3A%3A-1%5D%2C%20_v.numpy()%5B%3A%3A-1%5D%2C%20color%3D%22%23087bb9%22)%0A%20%20%20%20_ax.set_xlabel(%22probability%22)%3B%20_ax.set_title(f%22next%20token%20after%20%E2%80%9C%7Bprompt.value%7D%E2%80%9D%20%20(T%20%3D%20%7Btemperature.value%7D)%22%2C%20loc%3D%22left%22)%0A%20%20%20%20_ax.spines%5B%5B%22top%22%2C%20%22right%22%5D%5D.set_visible(False)%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%0A%20%20%20%20return%20(ids%2C)%0A%0A%0A%40app.cell%0Adef%20_(BOS%2C%20generate%2C%20ids%2C%20mo%2C%20model%2C%20n_new%2C%20temperature%2C%20tok%2C%20torch)%3A%0A%20%20%20%20mo.stop(not%20generate.value%2C%20mo.md(%22*Press%20Generate.*%22))%0A%20%20%20%20_out%20%3D%20model.generate(torch.tensor(%5Bids%5D)%2C%20max_new_tokens%3Dn_new.value%2C%20do_sample%3Dtemperature.value%20%3E%200%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20temperature%3Dtemperature.value%20if%20temperature.value%20%3E%200%20else%20None%2C%20top_k%3D0%2C%20pad_token_id%3DBOS)%0A%20%20%20%20mo.md(f%22**%7Btok.decode(ids)%7D**%7Btok.decode(_out%5B0%2C%20len(ids)%3A%5D)%7D%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%203.%20Attention%20patterns%0A%0A%20%20%20%20GPT-2%20has%2012%20layers%20with%2012%20heads%20each.%20Every%20head%20computes%0A%20%20%20%20%24A%20%3D%20%5Cmathrm%7Bsoftmax%7D%5C!%5Cleft(QK%5E%5Ctop%2F%5Csqrt%7B64%7D%20%2B%20%5Ctext%7Bmask%7D%5Cright)%24%3A%20row%20%3D%20query%20token%2C%20column%20%3D%20key%20token.%0A%20%20%20%20Try%20layer%205%2C%20head%205%20on%20a%20text%20that%20repeats%20itself%20(an%20*induction%20head*)%2C%20or%20layer%204%2C%20head%2011%0A%20%20%20%20(it%20looks%20at%20the%20previous%20token).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20att_text%20%3D%20mo.ui.text(value%3D%22%20quantum%20lamp%20river%20seven%20tiger%20cloud%20violin%20quantum%20lamp%20river%20seven%20tiger%20cloud%20violin%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20full_width%3DTrue%2C%20label%3D%22Text%22)%0A%20%20%20%20layer%20%3D%20mo.ui.slider(0%2C%2011%2C%20value%3D5%2C%20label%3D%22Layer%22%2C%20show_value%3DTrue)%0A%20%20%20%20head%20%3D%20mo.ui.slider(0%2C%2011%2C%20value%3D5%2C%20label%3D%22Head%22%2C%20show_value%3DTrue)%0A%20%20%20%20mo.vstack(%5Batt_text%2C%20mo.hstack(%5Blayer%2C%20head%5D%2C%20justify%3D%22start%22%2C%20gap%3D2)%5D)%0A%20%20%20%20return%20att_text%2C%20head%2C%20layer%0A%0A%0A%40app.cell%0Adef%20_(BOS%2C%20att_text%2C%20head%2C%20layer%2C%20model%2C%20plt%2C%20tok%2C%20torch)%3A%0A%20%20%20%20_ids%20%3D%20%5BBOS%5D%20%2B%20tok.encode(att_text.value)%0A%20%20%20%20_A%20%3D%20model(torch.tensor(%5B_ids%5D)%2C%20output_attentions%3DTrue).attentions%5Blayer.value%5D%5B0%2C%20head.value%5D.numpy()%0A%20%20%20%20_labels%20%3D%20%5B%22%E2%9F%A8bos%E2%9F%A9%22%5D%20%2B%20%5Btok.decode(%5Bi%5D)%20for%20i%20in%20_ids%5B1%3A%5D%5D%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(6.5%2C%206.5))%0A%20%20%20%20_ax.imshow(_A%2C%20cmap%3D%22viridis%22%2C%20vmin%3D0%2C%20vmax%3D0.8)%0A%20%20%20%20_ax.set_xticks(range(len(_labels))%2C%20_labels%2C%20rotation%3D90%2C%20fontsize%3D8)%0A%20%20%20%20_ax.set_yticks(range(len(_labels))%2C%20_labels%2C%20fontsize%3D8)%0A%20%20%20%20_ax.set_title(f%22layer%20%7Blayer.value%7D%2C%20head%20%7Bhead.value%7D%22%2C%20loc%3D%22left%22)%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%204.%20Logit%20lens%3A%20watching%20the%20prediction%20form%0A%0A%20%20%20%20The%20residual%20stream%20after%20every%20layer%20is%20a%20768-number%20vector%20per%20token.%20Applying%20the%20final%20layer%20norm%20and%0A%20%20%20%20the%20unembedding%20matrix%20to%20it%20*early*%20shows%20what%20the%20model%20would%20predict%20if%20it%20stopped%20there%0A%20%20%20%20(nostalgebraist%2C%202020).%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(BOS%2C%20model%2C%20np%2C%20plt%2C%20prompt%2C%20tok%2C%20torch)%3A%0A%20%20%20%20_ids%20%3D%20%5BBOS%5D%20%2B%20tok.encode(prompt.value)%0A%20%20%20%20_hs%20%3D%20model(torch.tensor(%5B_ids%5D)%2C%20output_hidden_states%3DTrue).hidden_states%0A%20%20%20%20_rows%2C%20_P%20%3D%20%5B%5D%2C%20%5B%5D%0A%20%20%20%20for%20_l%2C%20_h%20in%20enumerate(_hs)%3A%0A%20%20%20%20%20%20%20%20_z%20%3D%20model.transformer.ln_f(_h%5B0%5D)%20if%20_l%20%3C%2012%20else%20_h%5B0%5D%0A%20%20%20%20%20%20%20%20_p%20%3D%20model.lm_head(_z).softmax(-1)%0A%20%20%20%20%20%20%20%20_rows.append(%5Btok.decode(%5Bint(_p%5Bi%5D.argmax())%5D)%20for%20i%20in%20range(1%2C%20len(_ids))%5D)%0A%20%20%20%20%20%20%20%20_P.append(%5Bfloat(_p%5Bi%5D.max())%20for%20i%20in%20range(1%2C%20len(_ids))%5D)%0A%20%20%20%20_P%20%3D%20np.array(_P)%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(1.3%20*%20_P.shape%5B1%5D%20%2B%202%2C%206))%0A%20%20%20%20_ax.imshow(_P%2C%20cmap%3D%22Blues%22%2C%20vmin%3D0%2C%20vmax%3D1%2C%20aspect%3D%22auto%22)%0A%20%20%20%20for%20_l%20in%20range(_P.shape%5B0%5D)%3A%0A%20%20%20%20%20%20%20%20for%20_i%20in%20range(_P.shape%5B1%5D)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_ax.text(_i%2C%20_l%2C%20repr(_rows%5B_l%5D%5B_i%5D)%5B1%3A-1%5D%2C%20ha%3D%22center%22%2C%20va%3D%22center%22%2C%20fontsize%3D8%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20color%3D%22white%22%20if%20_P%5B_l%2C%20_i%5D%20%3E%200.55%20else%20%22black%22)%0A%20%20%20%20_ax.set_xticks(range(len(_ids)%20-%201)%2C%20%5Btok.decode(%5Bi%5D)%20for%20i%20in%20_ids%5B1%3A%5D%5D)%0A%20%20%20%20_ax.set_yticks(range(13)%2C%20%5B%22embed%22%5D%20%2B%20%5Bf%22layer%20%7Bl%7D%22%20for%20l%20in%20range(1%2C%2013)%5D)%0A%20%20%20%20_ax.xaxis.tick_top()%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20_fig%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%205.%20The%20embedding%20space%0A%0A%20%20%20%20Row%20%24i%24%20of%20the%2050%20257%20%C3%97%20768%20embedding%20matrix%20is%20the%20vector%20of%20token%20%24i%24.%20Similar%20tokens%20have%20similar%0A%20%20%20%20vectors%2C%20and%20differences%20between%20vectors%20can%20carry%20meaning%3A%0A%20%20%20%20%60king%20%E2%88%92%20man%20%2B%20woman%60%20lands%20near%20%60queen%60.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20word_b%20%3D%20mo.ui.text(value%3D%22king%22%2C%20label%3D%22b%22)%0A%20%20%20%20word_a%20%3D%20mo.ui.text(value%3D%22man%22%2C%20label%3D%22%E2%88%92%20a%22)%0A%20%20%20%20word_c%20%3D%20mo.ui.text(value%3D%22woman%22%2C%20label%3D%22%2B%20c%22)%0A%20%20%20%20mo.hstack(%5Bword_b%2C%20word_a%2C%20word_c%5D%2C%20justify%3D%22start%22%2C%20gap%3D1)%0A%20%20%20%20return%20word_a%2C%20word_b%2C%20word_c%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20model%2C%20tok%2C%20word_a%2C%20word_b%2C%20word_c)%3A%0A%20%20%20%20E%20%3D%20model.transformer.wte.weight%0A%20%20%20%20En%20%3D%20E%20%2F%20E.norm(dim%3D1%2C%20keepdim%3DTrue)%0A%0A%20%20%20%20def%20token_id(w)%3A%0A%20%20%20%20%20%20%20%20_ids%20%3D%20tok.encode(%22%20%22%20%2B%20w.strip())%0A%20%20%20%20%20%20%20%20return%20_ids%5B0%5D%20if%20len(_ids)%20%3D%3D%201%20else%20None%0A%0A%20%20%20%20_ia%2C%20_ib%2C%20_ic%20%3D%20(token_id(w.value)%20for%20w%20in%20(word_a%2C%20word_b%2C%20word_c))%0A%20%20%20%20if%20None%20in%20(_ia%2C%20_ib%2C%20_ic)%3A%0A%20%20%20%20%20%20%20%20_out%20%3D%20mo.md(%22Each%20word%20must%20be%20a%20**single**%20GPT-2%20token%20(with%20a%20leading%20space).%20Try%20another%20word.%22)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20_v%20%3D%20E%5B_ib%5D%20-%20E%5B_ia%5D%20%2B%20E%5B_ic%5D%0A%20%20%20%20%20%20%20%20_s%20%3D%20En%20%40%20(_v%20%2F%20_v.norm())%0A%20%20%20%20%20%20%20%20for%20_j%20in%20(_ia%2C%20_ib%2C%20_ic)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_s%5B_j%5D%20%3D%20-1%0A%20%20%20%20%20%20%20%20_best%20%3D%20%5B(tok.decode(%5Bint(j)%5D)%2C%20float(_s%5Bj%5D))%20for%20j%20in%20_s.topk(8).indices%5D%0A%20%20%20%20%20%20%20%20_near%20%3D%20%5Btok.decode(%5Bint(j)%5D)%20for%20j%20in%20(En%20%40%20En%5B_ib%5D).topk(9).indices%5B1%3A%5D%5D%0A%20%20%20%20%20%20%20%20_out%20%3D%20mo.md(f%22**%7Bword_b.value%7D%20%E2%88%92%20%7Bword_a.value%7D%20%2B%20%7Bword_c.value%7D**%20%E2%89%88%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2B%20%22%2C%20%22.join(f%22%60%7Bw.strip()%7D%60%20(%7Bs%3A.2f%7D)%22%20for%20w%2C%20s%20in%20_best)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%2B%20f%22%5Cn%5CnNearest%20tokens%20to%20**%7Bword_b.value%7D**%3A%20%22%20%2B%20%22%2C%20%22.join(f%22%60%7Bw%7D%60%22%20for%20w%20in%20_near))%0A%20%20%20%20_out%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%206.%20Why%20a%20KV%20cache%3F%0A%0A%20%20%20%20Generating%20token%20%24n%2B1%24%20needs%20the%20keys%20and%20values%20of%20all%20%24n%24%20previous%20tokens.%20They%20never%20change%2C%20so%20the%0A%20%20%20%20model%20can%20store%20them%20instead%20of%20recomputing%20the%20whole%20prefix%20at%20every%20step.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20run_kv%20%3D%20mo.ui.run_button(label%3D%22Time%20200%20new%20tokens%20with%20and%20without%20the%20cache%22)%0A%20%20%20%20run_kv%0A%20%20%20%20return%20(run_kv%2C)%0A%0A%0A%40app.cell%0Adef%20_(BOS%2C%20mo%2C%20model%2C%20run_kv%2C%20time%2C%20tok%2C%20torch)%3A%0A%20%20%20%20mo.stop(not%20run_kv.value)%0A%20%20%20%20_ids%20%3D%20torch.tensor(%5Btok.encode(%22The%20electron%20has%20a%20negative%22)%5D)%0A%20%20%20%20_t%20%3D%20%7B%7D%0A%20%20%20%20for%20_cache%20in%20(True%2C%20False)%3A%0A%20%20%20%20%20%20%20%20_t0%20%3D%20time.perf_counter()%0A%20%20%20%20%20%20%20%20model.generate(_ids%2C%20max_new_tokens%3D200%2C%20do_sample%3DFalse%2C%20use_cache%3D_cache%2C%20pad_token_id%3DBOS)%0A%20%20%20%20%20%20%20%20_t%5B_cache%5D%20%3D%20time.perf_counter()%20-%20_t0%0A%20%20%20%20mo.md(f%22with%20cache%3A%20**%7B_t%5BTrue%5D%3A.1f%7D%20s**%2C%20without%20cache%3A%20**%7B_t%5BFalse%5D%3A.1f%7D%20s**%20%22%0A%20%20%20%20%20%20%20%20%20%20f%22(%7B_t%5BFalse%5D%20%2F%20_t%5BTrue%5D%3A.1f%7D%C3%97%20slower).%20The%20gap%20grows%20with%20the%20length%20of%20the%20text.%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%207.%20A%20tiny%20GPT%20from%20scratch%0A%0A%20%20%20%20Everything%20above%20fits%20in%20a%20few%20dozen%20lines.%20Below%20is%20a%20character-level%20GPT%3A%20token%20%2B%20position%20embeddings%2C%0A%20%20%20%20a%20stack%20of%20blocks%20(causal%20multi-head%20attention%20%2B%20MLP%2C%20each%20with%20layer%20norm%20and%20a%20residual%20connection)%2C%0A%20%20%20%20and%20an%20unembedding.%20We%20train%20it%20on%20the%20source%20files%20of%20this%20course's%20slides%2C%20with%20the%20next-character%0A%20%20%20%20cross-entropy%20loss%20and%20Adam%20from%20Lecture%203A.%20About%20a%20minute%20on%20a%20laptop%20CPU.%0A%0A%20%20%20%20Watch%20the%20validation%20curve%3A%20the%20course%20text%20is%20small%20(~150%20k%20characters)%2C%20so%20after%20roughly%20700%20steps%20the%0A%20%20%20%200.8%20M-parameter%20model%20starts%20memorizing%20it.%20Training%20loss%20keeps%20falling%2C%20validation%20loss%20turns%20up%20%E2%80%94%0A%20%20%20%20the%20overfitting%20of%20Lecture%203A.%20Real%20language%20models%20see%20trillions%20of%20tokens%20and%20rarely%20repeat%20any.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(F%2C%20nn%2C%20torch)%3A%0A%20%20%20%20class%20CausalSelfAttention(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20d%2C%20n_heads%2C%20context)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.h%2C%20self.dk%20%3D%20n_heads%2C%20d%20%2F%2F%20n_heads%0A%20%20%20%20%20%20%20%20%20%20%20%20self.qkv%20%3D%20nn.Linear(d%2C%203%20*%20d)%20%20%20%20%20%20%20%20%20%20%23%20W_Q%2C%20W_K%2C%20W_V%20for%20all%20heads%20at%20once%0A%20%20%20%20%20%20%20%20%20%20%20%20self.out%20%3D%20nn.Linear(d%2C%20d)%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20W_O%0A%20%20%20%20%20%20%20%20%20%20%20%20self.register_buffer(%22mask%22%2C%20torch.tril(torch.ones(context%2C%20context)).bool())%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20x)%3A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20x%3A%20batch%20%C3%97%20T%20%C3%97%20d%0A%20%20%20%20%20%20%20%20%20%20%20%20B%2C%20T%2C%20d%20%3D%20x.shape%0A%20%20%20%20%20%20%20%20%20%20%20%20q%2C%20k%2C%20v%20%3D%20self.qkv(x).split(d%2C%20dim%3D-1)%0A%20%20%20%20%20%20%20%20%20%20%20%20q%2C%20k%2C%20v%20%3D%20(t.view(B%2C%20T%2C%20self.h%2C%20self.dk).transpose(1%2C%202)%20for%20t%20in%20(q%2C%20k%2C%20v))%20%20%23%20B%20%C3%97%20h%20%C3%97%20T%20%C3%97%20dk%0A%20%20%20%20%20%20%20%20%20%20%20%20scores%20%3D%20q%20%40%20k.transpose(-2%2C%20-1)%20%2F%20self.dk%20**%200.5%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20B%20%C3%97%20h%20%C3%97%20T%20%C3%97%20T%0A%20%20%20%20%20%20%20%20%20%20%20%20scores%20%3D%20scores.masked_fill(~self.mask%5B%3AT%2C%20%3AT%5D%2C%20float(%22-inf%22))%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20no%20peeking%20ahead%0A%20%20%20%20%20%20%20%20%20%20%20%20A%20%3D%20scores.softmax(-1)%0A%20%20%20%20%20%20%20%20%20%20%20%20y%20%3D%20(A%20%40%20v).transpose(1%2C%202).reshape(B%2C%20T%2C%20d)%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20concatenate%20heads%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.out(y)%0A%0A%20%20%20%20class%20Block(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20d%2C%20n_heads%2C%20context)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.ln1%2C%20self.ln2%20%3D%20nn.LayerNorm(d)%2C%20nn.LayerNorm(d)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.attn%20%3D%20CausalSelfAttention(d%2C%20n_heads%2C%20context)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.mlp%20%3D%20nn.Sequential(nn.Linear(d%2C%204%20*%20d)%2C%20nn.GELU()%2C%20nn.Linear(4%20*%20d%2C%20d))%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20x%20%2B%20self.attn(self.ln1(x))%20%20%20%20%20%20%20%20%20%20%20%23%20tokens%20exchange%20information%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20x%20%2B%20self.mlp(self.ln2(x))%20%20%20%20%20%20%20%20%20%20%20%20%23%20each%20token%20thinks%20on%20its%20own%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20x%0A%0A%20%20%20%20class%20TinyGPT(nn.Module)%3A%0A%20%20%20%20%20%20%20%20def%20__init__(self%2C%20vocab%2C%20d%3D128%2C%20n_heads%3D4%2C%20n_layers%3D4%2C%20context%3D128)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20super().__init__()%0A%20%20%20%20%20%20%20%20%20%20%20%20self.context%20%3D%20context%0A%20%20%20%20%20%20%20%20%20%20%20%20self.tok_emb%20%3D%20nn.Embedding(vocab%2C%20d)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.pos_emb%20%3D%20nn.Embedding(context%2C%20d)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.blocks%20%3D%20nn.Sequential(*%5BBlock(d%2C%20n_heads%2C%20context)%20for%20_%20in%20range(n_layers)%5D)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.ln_f%20%3D%20nn.LayerNorm(d)%0A%20%20%20%20%20%20%20%20%20%20%20%20self.unembed%20%3D%20nn.Linear(d%2C%20vocab%2C%20bias%3DFalse)%0A%0A%20%20%20%20%20%20%20%20def%20forward(self%2C%20idx)%3A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%20idx%3A%20batch%20%C3%97%20T%20integers%0A%20%20%20%20%20%20%20%20%20%20%20%20T%20%3D%20idx.shape%5B1%5D%0A%20%20%20%20%20%20%20%20%20%20%20%20x%20%3D%20self.tok_emb(idx)%20%2B%20self.pos_emb(torch.arange(T))%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20self.unembed(self.ln_f(self.blocks(x)))%20%20%20%23%20logits%3A%20batch%20%C3%97%20T%20%C3%97%20vocab%0A%0A%20%20%20%20%20%20%20%20def%20generate(self%2C%20idx%2C%20n%2C%20temperature%3D0.8)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20for%20_%20in%20range(n)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20logits%20%3D%20self(idx%5B%3A%2C%20-self.context%3A%5D)%5B%3A%2C%20-1%5D%20%2F%20temperature%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20idx%20%3D%20torch.cat(%5Bidx%2C%20torch.multinomial(F.softmax(logits%2C%20-1)%2C%201)%5D%2C%20dim%3D1)%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20idx%0A%0A%20%20%20%20return%20(TinyGPT%2C)%0A%0A%0A%40app.cell%0Adef%20_(Path%2C%20mo)%3A%0A%20%20%20%20_root%20%3D%20Path(mo.notebook_dir()).resolve().parent%20%20%20%20%20%20%20%20%20%20%23%20the%20lectures%2F%20folder%0A%20%20%20%20_files%20%3D%20sorted(_root.glob(%22Lecture_*%2FLecture_*.qmd%22))%0A%20%20%20%20corpus%20%3D%20%22%5Cn%22.join(f.read_text(encoding%3D%22utf-8%22)%20for%20f%20in%20_files)%0A%20%20%20%20chars%20%3D%20sorted(set(corpus))%0A%20%20%20%20mo.md(f%22Training%20text%3A%20**%7Blen(_files)%7D**%20lecture%20sources%2C%20**%7Blen(corpus)%3A%2C%7D**%20characters%2C%20%22%0A%20%20%20%20%20%20%20%20%20%20f%22**%7Blen(chars)%7D**%20distinct%20characters%20(our%20vocabulary).%22)%0A%20%20%20%20return%20chars%2C%20corpus%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20steps%20%3D%20mo.ui.slider(200%2C%203000%2C%20step%3D100%2C%20value%3D1000%2C%20label%3D%22Training%20steps%22%2C%20show_value%3DTrue)%0A%20%20%20%20run_train%20%3D%20mo.ui.run_button(label%3D%22Train%20the%20tiny%20GPT%22)%0A%20%20%20%20mo.hstack(%5Bsteps%2C%20run_train%5D%2C%20justify%3D%22start%22%2C%20gap%3D2)%0A%20%20%20%20return%20run_train%2C%20steps%0A%0A%0A%40app.cell%0Adef%20_(TinyGPT%2C%20chars%2C%20corpus%2C%20mo%2C%20plt%2C%20run_train%2C%20steps%2C%20torch)%3A%0A%20%20%20%20mo.stop(not%20run_train.value%2C%20mo.md(%22*Press%20Train%20(about%20a%20minute).*%22))%0A%20%20%20%20torch.manual_seed(0)%0A%20%20%20%20stoi%20%3D%20%7Bc%3A%20i%20for%20i%2C%20c%20in%20enumerate(chars)%7D%0A%20%20%20%20data%20%3D%20torch.tensor(%5Bstoi%5Bc%5D%20for%20c%20in%20corpus%5D)%0A%20%20%20%20n_train%20%3D%20int(0.9%20*%20len(data))%0A%20%20%20%20gpt%20%3D%20TinyGPT(len(chars))%0A%20%20%20%20opt%20%3D%20torch.optim.AdamW(gpt.parameters()%2C%20lr%3D2e-3)%0A%0A%20%20%20%20def%20batch(split%2C%20B%3D32)%3A%0A%20%20%20%20%20%20%20%20src%20%3D%20data%5B%3An_train%5D%20if%20split%20%3D%3D%20%22train%22%20else%20data%5Bn_train%3A%5D%0A%20%20%20%20%20%20%20%20i%20%3D%20torch.randint(len(src)%20-%20gpt.context%20-%201%2C%20(B%2C))%0A%20%20%20%20%20%20%20%20x%20%3D%20torch.stack(%5Bsrc%5Bj%3Aj%20%2B%20gpt.context%5D%20for%20j%20in%20i%5D)%0A%20%20%20%20%20%20%20%20return%20x%2C%20torch.stack(%5Bsrc%5Bj%20%2B%201%3Aj%20%2B%20gpt.context%20%2B%201%5D%20for%20j%20in%20i%5D)%20%20%20%23%20targets%20%3D%20inputs%20shifted%20by%20one%0A%0A%20%20%20%20history%20%3D%20%5B%5D%0A%20%20%20%20with%20torch.enable_grad()%3A%0A%20%20%20%20%20%20%20%20for%20_s%20in%20mo.status.progress_bar(range(steps.value)%2C%20title%3D%22training%22)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20_x%2C%20_y%20%3D%20batch(%22train%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20_loss%20%3D%20torch.nn.functional.cross_entropy(gpt(_x).flatten(0%2C%201)%2C%20_y.flatten())%0A%20%20%20%20%20%20%20%20%20%20%20%20opt.zero_grad()%3B%20_loss.backward()%3B%20opt.step()%0A%20%20%20%20%20%20%20%20%20%20%20%20if%20_s%20%25%2050%20%3D%3D%200%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20with%20torch.no_grad()%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_xv%2C%20_yv%20%3D%20batch(%22val%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20_lv%20%3D%20torch.nn.functional.cross_entropy(gpt(_xv).flatten(0%2C%201)%2C%20_yv.flatten())%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20history.append((_s%2C%20_loss.item()%2C%20_lv.item()))%0A%20%20%20%20_fig%2C%20_ax%20%3D%20plt.subplots(figsize%3D(6%2C%203.2))%0A%20%20%20%20_ax.plot(%5Bh%5B0%5D%20for%20h%20in%20history%5D%2C%20%5Bh%5B1%5D%20for%20h%20in%20history%5D%2C%20label%3D%22train%22)%0A%20%20%20%20_ax.plot(%5Bh%5B0%5D%20for%20h%20in%20history%5D%2C%20%5Bh%5B2%5D%20for%20h%20in%20history%5D%2C%20label%3D%22validation%22)%0A%20%20%20%20_ax.axhline(torch.log(torch.tensor(len(chars))).item()%2C%20color%3D%22grey%22%2C%20ls%3D%22%3A%22%2C%20label%3D%22uniform%20guess%22)%0A%20%20%20%20_ax.set_xlabel(%22step%22)%3B%20_ax.set_ylabel(%22cross-entropy%20(nats%20per%20character)%22)%3B%20_ax.legend(frameon%3DFalse)%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20mo.vstack(%5Bmo.md(f%22%7Bsum(p.numel()%20for%20p%20in%20gpt.parameters())%3A%2C%7D%20parameters%22)%2C%20_fig%5D)%0A%20%20%20%20return%20gpt%2C%20stoi%0A%0A%0A%40app.cell%0Adef%20_(chars%2C%20gpt%2C%20mo%2C%20stoi%2C%20torch)%3A%0A%20%20%20%20_start%20%3D%20%22%23%23%20Attention%22%0A%20%20%20%20_idx%20%3D%20torch.tensor(%5B%5Bstoi%5Bc%5D%20for%20c%20in%20_start%5D%5D)%0A%20%20%20%20torch.manual_seed(1)%0A%20%20%20%20_out%20%3D%20gpt.generate(_idx%2C%20600%2C%20temperature%3D0.8)%5B0%5D%0A%20%20%20%20mo.md(%22Sample%20from%20the%20tiny%20GPT%20(it%20has%20seen%20nothing%20but%20these%20slides)%3A%5Cn%5Cn%60%60%60%5Cn%22%0A%20%20%20%20%20%20%20%20%20%20%2B%20%22%22.join(chars%5Bi%5D%20for%20i%20in%20_out)%20%2B%20%22%5Cn%60%60%60%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Exercises%0A%0A%20%20%20%201.%20Find%20a%20word%20that%20GPT-2%20splits%20into%20many%20tokens%20but%20the%20o200k%20tokenizer%20keeps%20whole.%20Why%20does%20the%20vocabulary%20matter%20for%20cost%20and%20for%20spelling%20tasks%3F%0A%20%20%20%202.%20Set%20the%20temperature%20to%200%20and%20generate%2080%20tokens.%20What%20goes%20wrong%3F%20Why%20does%20sampling%20help%3F%0A%20%20%20%203.%20Find%20a%20previous-token%20head%20and%20an%20induction%20head%20in%20section%203.%20Why%20does%20an%20induction%20head%20need%20a%20previous-token%20head%20in%20an%20*earlier*%20layer%3F%0A%20%20%20%204.%20In%20the%20logit%20lens%2C%20at%20which%20layer%20does%20the%20correct%20answer%20first%20appear%20for%20your%20own%20prompt%3F%0A%20%20%20%205.%20Count%20the%20parameters%20of%20the%20tiny%20GPT%20by%20hand%3A%20embeddings%2C%20attention%20(4%20d%C2%B2%20%2B%204d%20per%20block)%2C%20MLP%20(8%20d%C2%B2%20%2B%205d%20per%20block)%2C%20layer%20norms.%0A%20%20%20%206.%20Remove%20the%20causal%20mask%20in%20%60CausalSelfAttention%60%20and%20retrain.%20The%20training%20loss%20falls%20much%20faster%20%E2%80%94%20why%20is%20the%20model%20useless%20for%20generation%3F%0A%20%20%20%207.%20Replace%20the%20learned%20position%20embedding%20with%20nothing%20at%20all.%20What%20happens%20to%20the%20samples%2C%20and%20why%3F%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A
52e9db1ffb9b82e412dcffe73440626a