शून्य से शुरू किया गया Llama: बिना रोए पेपर को implement करने का तरीका

 (blog.briankitano.com)
2 पॉइंट द्वारा GN⁺ 2023-08-10 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • Brian Kitano ने TinyShakespeare के साथ एक छोटा Llama खुद बनाते हुए निष्कर्ष निकाला कि पेपर implementation के लिए छोटे मॉडल से शुरू करना, parts को एक-एक करके बदलना, और हर बार training व evaluation करना सबसे सुरक्षित तरीका है
  • उन्होंने data split, batch generation, loss evaluation, generation function जैसी validation helper functions पहले तैयार कीं, फिर simple model से compile और train होने की पुष्टि करने के बाद Llama के components जोड़े
  • RMSNorm, RoPE, SwiGLU को क्रम से जोड़ते हुए tensor shape, formula properties, और attention maps के जरिए हर layer के अपेक्षित तरीके से काम करने की पुष्टि की
  • RoPE attention में causal mask हटाने पर validation loss 0.16 तक गिर गया, लेकिन generation quality खराब हो गई; कारण था future tokens देखने से होने वाला information leakage
  • अंतिम छोटे Llama में 4 blocks और लगभग 23.7 लाख parameters थे, जिसने validation loss को लगभग 1.0 तक घटाया; साथ ही gradient flow और learning-rate schedule की भी जांच ज़रूरी रही

छोटे से शुरू करें और बार-बार भरोसा बनाएं

  • पेपर implementation का मूल सिद्धांत है छोटे मॉडल से शुरू करना, components को एक-एक करके बदलना, और हर बदलाव पर training व evaluation दोहराना
  • पहले मॉडल को मात्रात्मक रूप से जांचने के लिए helper functions तैयार किए गए
    • data split
    • training loop
    • loss visualization
    • validation loss evaluation
  • पेपर के सभी components एक साथ उतारने के बजाय, पहले से implement किए गए सरल और तेज मॉडल के साथ generation output देखने के लिए qualitative evaluation function भी तैयार की गई
  • tensor layers को .shape, assert, plt.imshow से जांचा गया, और शुरुआत से ही matrix multiplication optimization में जाने के बजाय पहले हाथ से expected results सत्यापित किए गए, फिर torch functions से उन्हें efficient बनाया गया
  • batch size, sequence length, और embedding dimension बदल-बदल कर test करना चाहिए; सिर्फ एक आकार पर सही चलने वाला code inference के समय टूट सकता है

dataset और basic setup

  • implementation का लक्ष्य Meta AI के Llama का बहुत छोटा संस्करण था, और training data TinyShakespeare था
  • Llama को 1.4T tokens पर train किया गया था, लेकिन यहां लगभग 11.1 लाख characters वाले TinyShakespeare का उपयोग हुआ
  • मूल Llama SentencePiece byte-pair encoding tokenizer का उपयोग करता है, लेकिन इस implementation में सरल character-level tokenizer इस्तेमाल किया गया
    • vocabulary size 65 है
    • dataset छोटा होने के कारण memory storage method को अलग से optimize नहीं किया गया
  • MASTER_CONFIG dictionary के जरिए vocab_size, batch_size, context_window, d_model जैसे model settings प्रबंधित किए गए
    • उद्देश्य constants और magic numbers को घटाना और code को पढ़ने में आसान बनाना था
  • get_batches function data को train 80%, val 10%, test 10% में बांटता है, और random start points से input x तथा एक character आगे का label y बनाता है

basic model से compile और training की पुष्टि

  • पहला मॉडल embedding और simple feed-forward network से बना SimpleBrokenModel था
    • nn.Embedding
    • Linear
    • ReLU
    • Linear
  • पेपर implementation में मॉडल के “काम करने” का मतलब है कि ये दोनों शर्तें पूरी हों
    • compile: tensors के shape layers के बीच सही बैठें
    • train: loss वास्तव में नीचे जाए
  • evaluate_loss function train और val split से 10 batches sample करके average loss निकालता है
  • SimpleBrokenModel ने 1000 epochs training के बाद validation loss को 3.94 तक पहुंचाया, जो शुरुआती cross-entropy 4.17 से बहुत कम सुधार था
  • वजह यह थी कि F.cross_entropy में पहले से softmax किए गए values दे दिए गए थे
    • PyTorch का F.cross_entropy सीधे unnormalized logits लेता है
    • softmax हटाने के बाद SimpleModel validation loss को 2.51 तक ला सका
  • इसके बाद generate function जोड़ा गया ताकि मॉडल द्वारा बनाए गए characters को सीधे देखा जा सके; basic model अभी पूर्ण नहीं था, लेकिन validation loss कम होने लगा था

Llama component 1: RMSNorm

  • Llama मूल Transformer की तुलना में architecture में तीन मुख्य बदलाव उपयोग करता है
    • RMSNorm pre-normalization
    • Rotary embeddings
    • SwiGLU activation function
  • मूल Transformer BatchNormalization का उपयोग करता है, जबकि Llama vector को center किए बिना variance से scale करने वाला RMSNorm उपयोग करता है
  • जहां मूल Transformer attention layer output पर normalization लागू करने वाला post-normalization तरीका अपनाता है, वहीं Llama input पर पहले normalization लागू करने वाला pre-normalization तरीका अपनाता है
  • implement किया गया RMSNorm input shape (batch, seq_len, d_model) मानकर चलता है
  • RMSNorm के परिणाम को इस गुण से test किया गया कि layer norm, layer elements की संख्या के square root के बराबर हो
    • assert
    • row-wise comparison
    • torch.allclose
  • basic model में RMSNorm जोड़ने वाले SimpleModel_RMS ने validation loss को थोड़ा घटाकर 2.5015 किया

Llama component 2: RoPE और causal mask

  • RoPE Transformer के लिए positional encoding का एक तरीका है, जो token position को embedding rotation के रूप में व्यक्त करता है
  • get_rotary_matrix context window और embedding dimension के लिए हर position का rotation matrix बनाता है
  • RoPE implementation को इस गुण से test किया गया
    • positions m, n पर rotate किए गए दो vectors का inner product, relative position n-m rotation के अनुरूप होना चाहिए
  • RoPEAttentionHead w_q, w_k, w_v बनाता है, query और key पर RoPE rotation लागू करता है, फिर F.scaled_dot_product_attention इस्तेमाल करता है
  • training और inference के समय tensor shape के अंतर पर ध्यान देना ज़रूरी है
    • training में अक्सर shape (config['batch_size'], config['context_window'], config['d_model']) जैसा होता है
    • inference में (1, 1, config['d_model']) जैसे single example को process किया जा सकता है
    • forward के अंदर model config values नहीं, बल्कि input से मिले shape के आधार पर indexing करनी चाहिए
  • causal mask के बिना RoPE multi-head attention जोड़ने वाले मॉडल ने validation loss को तेज़ी से 0.1623 तक गिरा दिया, लेकिन generation output OOOO..., IIII... जैसा खराब था
  • attention map देखने पर पता चला कि हर position सभी positions को reference कर रही थी, और next-token prediction में future tokens देखने से information leakage हो रहा था
  • F.scaled_dot_product_attention में is_causal=True लगाने वाले RoPEMaskedAttentionHead पर बदलने से future वाले upper triangular attention लगभग 0 हो गए
  • causal mask लागू करने के बाद validation loss 2.0815 हुआ, और अधिक training पर 1.8985 तक घटा

Llama component 3: SwiGLU और blocks को stack करना

  • Llama, ReLU nonlinearity की जगह SwiGLU activation function उपयोग करता है
  • implement किया गया SwiGLU एक Swish-gated linear unit है, जो दो linear transforms और learnable beta parameter का उपयोग करता है
  • feed-forward भाग में SwiGLU जोड़ने वाले RopeModel में 592,706 parameters थे और validation loss 1.8963 तक पहुंचा
  • इसके बाद LlamaBlock बनाकर निम्न संरचना को एक block में बांधा गया
    • RMSNorm pre-normalization
    • masked RoPE multi-head attention
    • residual connection
    • RMSNorm pre-normalization
    • SwiGLU feed-forward
    • residual connection
  • अंतिम Llama मॉडल में n_layers=4 रखा गया और OrderedDict आधारित nn.Sequential से 4 LlamaBlock stack किए गए
  • अंतिम मॉडल के parameters की संख्या 2,370,246 थी, और training परिणाम इस प्रकार रहे
    • शुरुआती 4-layer training के बाद validation loss 1.5532
    • 10,000 epochs तक और training के बाद validation loss 1.1479
    • अतिरिक्त training के बाद validation loss 0.9997
    • test split के एक batch का loss 1.2358

generation results और debugging checks

  • अंतिम मॉडल Shakespeare जैसी formatting, names, line breaks, और शब्द-टुकड़े बना सकता है, लेकिन वास्तविक sentence quality अभी सीमित है
  • cross-entropy loss को token selection के नज़रिए से सहज रूप में समझा जा सकता है
    • शुरुआती loss 4.17, vocabulary size 65 में लगभग random selection के बराबर है
    • loss 1.08 का मतलब लगभग 2.9 tokens में से random चुनने जैसा स्तर है
  • gradient flow को show_grads function से जांचा गया
    • हर parameter में बहुत छोटे absolute gradient के अनुपात की गणना की गई
    • अगर अधिकांश parameters के gradients 0 के बहुत करीब न हों, तो flow ठीक माना जा सकता है
  • मूल Llama Cosine Annealing learning schedule उपयोग करता है, लेकिन इस implementation में प्रयोगात्मक परिणाम खराब थे
  • Cosine Annealing प्रयोग में बहुत कम tolerance पर भी attention bias को लगभग कोई signal नहीं मिला; कारण स्पष्ट नहीं था, इसलिए वास्तविक implementation में सरल शुरुआत करना अधिक सुरक्षित है

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2023-08-10
Hacker News की राय

  • लगता है SwiGLU implementation में bug है: संदर्भ paper में feed-forward network का beta सीखने योग्य value नहीं, बल्कि constant है, और उसे FFnSwiGLU = Swish1... रखा गया है
    यह https://arxiv.org/pdf/2002.05202.pdf के equation 6 के आधार पर है
    आधिकारिक llama implementation में भी constant beta हटा हुआ है: https://github.com/facebookresearch/llama/blob/main/llama/mo...
    ब्लॉग log की "feedforward.1.beta', 0.0" lines देखें तो training के दौरान beta 0 पर degenerate हो गया, जबकि मूल रूप से यह constant 1 होना चाहिए था

    • यह दिखाता है कि Transformer neural network को सही-सही implement करना कितना मुश्किल है। कई चरणों में गलती हो सकती है, और आम तौर पर यह सिर्फ “मूल से थोड़ा खराब performance” के रूप में दिखती है, इसलिए पक्का जानना कठिन होता है
      Network कई बार, चाहे इरादा हो या नहीं, बदलावों के हिसाब से adapt कर लेता है, और training के बाद कई architecture variants भी मिलते-जुलते तरीके से काम करते हैं, इसलिए यह भी अस्पष्ट हो सकता है कि original से बिल्कुल match करना जरूरी है या नहीं
      ऐसी गलतियां खोजने का एक तरीका है reference implementation के output values से बिल्कुल match कराना। HuggingFace के tiny-random model की तरह, random weights के साथ भी output बिल्कुल समान होना चाहिए; अगर अलग हो तो यह bug का संकेत है
      हालांकि यह तरीका inference के दौरान आने वाले bugs पर ही अच्छा काम करता है; data processing, optimizer, या सिर्फ training के दौरान होने वाली समस्याओं को पकड़ना ज्यादा मुश्किल है
    • Transformer में bias values आम तौर पर ठीक से match नहीं होते, ऐसा लगता है
      निजी तौर पर मुझे लगता है कि इसकी वजह इसका autoregressive और ODE जैसा स्वभाव है, लेकिन इतना यकीन नहीं है

  • काम बेहतरीन है, लेकिन शुरुआती SimpleBrokenModel और SimpleModel में काफी wasted computation है। क्रम है embedding 65 -> 128, linear 128 -> 128, ReLU, linear 128 -> 65; लेकिन पहले दो layers के बीच कोई non-linearity नहीं है और दोनों linear हैं, इसलिए दूसरा linear layer असल में बेकार है
    यह model आखिरकार classic single-hidden-layer MLP जैसा ही है, और FLOPS के हिसाब से 128*128=16k operations, कुल 128*128+65*128=24k में से waste हो रहे हैं

    • लगता है non-linearity को समझने की कोशिश करने वाला मैं अकेला नहीं हूं। यहां सबसे अच्छा fix क्या होगा—embedding और पहले linear layer के बीच ReLU या SwiGLU डालना, या फिर linear layer को सीधे हटा देना—यह जानना चाहूंगा
      Embedding layer token index को embedding vector में बदलने वाली special structure है, इसलिए लगता है इसे हटाया नहीं जा सकता

  • कुल मिलाकर basic principles अच्छी तरह दिखाए गए हैं। खासकर “.shape को धर्म की तरह इस्तेमाल करो। assert और plt.imshow तुम्हारे दोस्त हैं” वाली बात अच्छी लगी, और shape की pre/post conditions हमेशा assert करनी चाहिए
    यह भी जानना चाहूंगा कि bear या typeguard decorator के रूप में ऐसी checks support करते हैं या नहीं
    हालांकि “छोटा, simple और तेज model चुनो और qualitatively evaluate करने वाला helper बनाओ” वाला हिस्सा शायद quantitative evaluation के लिए होना चाहिए। तभी advanced techniques से तुलना करने के लिए numeric baseline मिलेगा
    Paper के components को एक-एक करके implement करने की सलाह भी ज्यादा precise होनी चाहिए। Papers आमतौर पर कई changes को एक साथ try करते हैं और फिर ablation experiments से हर component का contribution दिखाते हैं, इसलिए core architecture change से शुरू करना और ablation में ज्यादा असर वाले क्रम में, dependencies बनाए रखते हुए, हर atomic change पर evaluate करना बेहतर लगता है

    • bear या typeguard के बजाय https://peps.python.org/pep-0646/ की वजह से इसका कुछ हिस्सा सीधे Python type annotations में डाला जा सकता है
      उदाहरण के लिए ndarray[float, Dim1, *Shape] जैसी शैली में हर axis की shape को type में व्यक्त किया जा सकता है, और axis value के आधार पर return shape को overload किया जा सकता है
    • PyTorch के बारे में ठीक से नहीं पता, लेकिन आखिरी बार जब देखा था तो ऐसा नहीं था; Jax bear / typeguard के जरिए matrix shape की बुनियादी runtime checks support करता है
      फिर भी Python, Julia जितना अच्छा होना मुश्किल लगता है। Julia का type system matrix sizes सही हैं या नहीं, यह guarantee करना कहीं आसान बना देता है

  • यह जानना चाहूंगा कि ReLU की जगह SwiGLU इस्तेमाल करने का सिद्धांत क्या है। पता नहीं authors ने बस हर संभव non-linear function try किया, या इसके पीछे कोई गहरी वजह है

    • कई research की तरह, अगर किसी rigorous study से समर्थित स्पष्ट explanation नहीं है, तो शायद cool दिखने वाले one-line changes को random hill-climbing style search से आजमाते रहे, और paper लिखने व ablation experiments शुरू करने का समय आने पर रुक गए

  • bearblog पर DDoS हो रहा है, इसलिए repository छोड़ रहा हूं: https://github.com/bkitano/llama-from-scratch

  • AI सीखने वाले के नजरिए से, लेख में आए शब्दों को संक्षेप में整理 किया है। टोकन text के टुकड़े को दर्शाने वाला integer identifier होता है, और LLM में सीमित vocabulary size के भीतर अक्सर इस्तेमाल होने वाले character pieces को समूह में इस्तेमाल किया जाता है
    loss function prediction और सही उत्तर के बीच का अंतर मापने वाली value है, और जितनी कम हो उतना अच्छा। PyTorch tensors और neural networks को संभालने वाली library है, और tensor scalar, vector और matrix को शामिल करने वाला multidimensional number array है
    neural network weights और biases वाले neurons की connected structure है, और linear layer एक simple structure है जिसमें सभी inputs और outputs जुड़े होते हैं। ReLU Math.max(0, x) जैसी activation function है; अगर केवल linear layers ही stack की जाएँ तो अंत में वह एक ही linear function के बराबर हो जाती है, इसलिए non-linearity जोड़कर सीखने की संभावना बढ़ाई जाती है
    gradient सीखने के दौरान model को अधिक accurate बनाने के लिए calculate की जाने वाली numerical change quantity है, और batch normalization flowing numbers को adjust करके learning में मदद करने की method है। positional encoding tokens की relative positions को vector के रूप में बताता है
    Python का @ operator __matmul__ का alias है और matrix multiplication में इस्तेमाल होता है। epoch का मतलब पूरे dataset को एक बार train करना है, और batch parameter update से पहले एक बार में डाले जाने वाले data की संख्या है
    attention LLM को चलाने वाला core है; यह input tokens को parallel में process करके intermediate tensors बनाता है और फिर output tokens generate करने में उनका इस्तेमाल करता है

    • field के बाहर लोगों को शायद यह न पता हो कि “Karpathy” का क्या मतलब है। Andrej Karpathy को “science communicator और researcher” जैसे context के साथ introduce करने से यह ज्यादा स्पष्ट हो जाता है कि उनकी writings या videos को refer करने की बात है
    • token को सिर्फ text piece का integer identifier मानने के बजाय, beginners के लिए भी ज्यादा accurate होगा कि उसे अपने-आप में useful होने जितना common word piece माना जाए
      उदाहरण के लिए writing, written, writer में common आने वाला writ एक token हो सकता है, और writer को writ और er में tokenize किया जा सकता है
      embedding इन tokens को unique numerical representations में बदलने वाला step है
    • linear functions को compose करने पर फिर से linear function ही मिलता है। इसलिए अगर सब कुछ linear हो, तो कई layers stack करने पर भी एक को छोड़कर बाकी layers waste हो जाती हैं; इससे बचने के लिए non-linearity चाहिए
    • Karpathy की video series और accompanying repo के अलावा, learning journey में खास मददगार रहे और resources या books हैं क्या, यह जानना चाहूँगा
    • batch normalization exactly क्या करती है और कैसे मदद करती है, यह जानना चाहूँगा

  • अगर model की existing implementation और checkpoint मौजूद हो, तो अपनी implementation सही है या नहीं जाँचने का सबसे effective तरीका है उस checkpoint को load करके output values compare करना
    output match न हो तो आम तौर पर detailed implementation में गलती होती है, और हर layer को systematically trace करते हुए असली फर्क खोजा जा सकता है। इस दौरान existing implementation की कोई अजीब बात भी मिल सकती है
    यह model खुद के बारे में है, training एक अलग axis है। फिर भी अगर hyperparameters को कुछ हद तक similar रखा हो, तो model implementation सही होने पर आम तौर पर चीजें ठीक रहती हैं

  • paper पढ़ने का तरीका और उस paper की content दोनों अच्छे हैं, और Karpathy की Makemore series भी recommend करता हूँ

  • summary advice बहुत अच्छी हैं, और tensors की shape assert करने वाली advice किसी भी general linear algebra library पर लागू होती है। complex linear algebra code लिखते समय छोटे steps में जाना और defensive coding करना बहुत महत्वपूर्ण है
    mainstream languages में linear algebra programming करना compile-time shape checking न होने की वजह से भयानक है। tensor की shape type का हिस्सा होनी चाहिए, और अगर 3x4 और 3x4 को transpose किए बिना multiply करने की कोशिश की जाए तो compile ही नहीं होना चाहिए
    लंबी computation चलाने के बाद dimension-mismatch operation पर fail होना सच में सबसे खराब है
    मेरा मानना है कि PyTorch tensors में device भी statically typed होना चाहिए। अभी CPU memory के tensor और GPU memory के tensor को multiply करने की कोशिश करने पर runtime error आता है