DeepDive in everything of Llama3: revealing detailed insights and implementation

This project is an enhanced version based on naklecha/llama3-from-scratch. It has been comprehensively improved and optimized on the basis of the original project, aiming to help everyone more easily understand and master the implementation principle and the detailed reasoning process of the Llama3 model. Thanks to the contributions of the original author :)

1. Structural Optimization

2. Code Annotations

3. Dimension Tracking

4. Principle Explanation

5. KV-Cache Insights

6. Bilingual Documents

• Loading the model
  • Loading the tokenizer
  • Reading model files and configuration files
    • Inferring model details using the configuration file
• Convert the input text into embeddings
  • Convert the text into a sequence of token ids
  • Convert the sequence of token ids into embeddings
• Build the first Transformer block
  • Normalization
    • Using RMS normalization for embeddings
  • Implementing the single-head attention mechanism from scratch
    • Obtain the QKV vectors corresponding to the input tokens
      • Obtain the query vector
        • Unfold the query weight matrix
        • Obtain the first head
        • Multiply the token embeddings by the query weights to obtain the query vectors corresponding to the tokens
      • Obtain the key vector (almost the same as the query vector)
      • Obtain the value vector (almost the same as the key vector)
    • Add positional information to the query and key vectors
      • Rotary Position Encoding (RoPE)
      • Add positional information to the query vectors
      • Add positional information to the key vectors (same as the query)
    • Everything's ready. Let's start calculating the attention weights between tokens.
      • Multiply the query and key vectors to obtain the attention scores.
      • Now we must mask the future query-key scores.
      • Calculate the final attention weights, that is, softmax(score).
    • Finally! Calculate the final result of the single-head attention mechanism!
  • Calculate the multi-head attention mechanism (a simple loop to repeat the above process)
    • Calculate the result for each head
    • Merge the results of each head into a large matrix
    • Head-to-head information interaction (linear mapping), the final step of the self-attention layer!
  • Perform the residual operation (add)
  • Perform the second normalization operation
  • Perform the calculation of the FFN (Feed-Forward Neural Network) layer
  • Perform the residual operation again (Finally, we get the final output of the Transformer block!)
• Everything is here. Let's complete the calculation of all 32 Transformer blocks. Happy reading :)
• Let's complete the last step and predict the next token
  • First, perform one last normalization on the output of the last Transformer layer
  • Then, make the prediction based on the embedding corresponding to the last token (perform a linear mapping to the vocabulary dimension)
  • Here's the prediction result!
• Let's dive deeper and see how different embeddings or token masking strategies might affect the prediction results :)
• Need to predict multiple tokens? Just using KV-Cache! (It really took me a lot of effort to sort this out. Orz)
• Thank you all. Thanks for your continuous learning. Love you all :)
  • From Me
  • From the author of predecessor project
• LICENSE