Skip to content

Latest commit

 

History

History
151 lines (125 loc) · 8.83 KB

README.md

File metadata and controls

151 lines (125 loc) · 8.83 KB
Raw

X-Llama🦙: Extensible Language Model inspired by the original Llama model.

Your Image Description

X-Llama is an advanced language model framework, inspired by the original Llama model but enhanced with additional features such as Grouped Query Attention (GQA), Multi-Head Attention (MHA), and more. This project aims to provide a flexible and extensible platform for experimenting with various attention mechanisms and building state-of-the-art natural language processing models.

project structure: The model was constructed in approximately ~500 lines of code, and you have the model's configuration.

X-Llama/
│
├── images/
│
├── models/
│   ├── attentions/
│   ├── rotary_embeddings/
│   └── transformer/
│
├── model
│
└── config
│
└── inference

Features:

  • Rotary Embeddings:
    • Rotary Embeddings.
    • Linear Scaling Rotary Embeddings.
    • Dynamic NTK Scaling Rotary Embeddings.

Your Image Description 

LLAMA_ROTARY_EMBEDDINGS_CLASSES = {
    "rotary": LlamaRotaryEmbeddings,
    "linear": LlamaLinearScalingRotaryEmbeddings,
    "dynamic": LlamaDynamicNTKScalingRotaryEmbeddings,
	}

  • LlamaChat. coming soon

  • Attentions: The standard practice for autoregressive decoding is to cache the keys and values of the previous tokens in the sequence to speed up attention computation. However, as the context window or batch size increases, the memory cost associated with the size of the key-value cache(kv cache) in the multi-head attention(MHA) model significantly increases.

    • Multi-Head Attention(MHA):
      Self-attention is calculated by taking the dot product of the query and key, scaled by a factor, and applying a softmax function to obtain attention weights. These attention weights determine the importance of each word's value for the current word. $$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

      Your Image Description

    • Grouped Query Attention(GQA), and Multi-Query Attention(MQA):
      Grouped-query attention divides query heads into G-groups, each of which shares a single key head and value head. GQA-G refers to grouped-query with G groups. GQA-1, with a single group and therefore single key and value head, is equivalent to MQA, while GQA-H, with groups equal to number of heads, is equivalent to MHA. following Figure shows a comparison of grouped-query attention and multi head/multi-query attention. When converting a multi-head checkpoint to a GQA checkpoint, we construct each group key and value head by mean pooling all the original heads within that group. An intermediate number of groups leads to an interpolated model that is higher quality than MQA but faster than MHA, and, as we will show, rep resents a favorable trade-off. Going from MHA to MQA reduces H key and value heads to a single key and value head, reducing the size of the key-value cache and therefore amount of data that needs to be loaded by a factor of H. However, larger models generally scale the number of heads, such that multi-query attention represents a more aggressive cut in both memory bandwidth and capacity. GQA lets us keep the same proportional decrease in bandwidth and capacity as model size increases.

      • MQA: Multi-Query attention(MQA) is a mechanism that uses only a single key-value head for multiple queries, which can save memory and greatly speed up decoder inference.
      • Fixed GQA: However, MQA may lead to a decrease in quality. In fact, we not only want fast inference but also want the quality to be on par with MHA, so Grouped-query attention(GQA) comes into play. Grouped-query attention(GQA) is an interpolation of multi-query and multi-head attention. It achieves a quality similar to multi-head attention while maintaining a comparable speed to multi-query attention.
      • Scalable GQA: the same as the fixed GQA but with multiple rotary embeddings.

        Your Image Description

      MHA vs GQA vs MQA:

    MHA GQA MQA
    High quality A good compromise between quality and Loss in quality
    Computationally slow speed Computationally fast

    Your Image Description

    Time per sample for GQA-XXL as a function of the number of GQA groups with input length 2048 and output length 512. Going from 1 (MQA) to 8 groups adds modest inference overhead, with increasing cost to adding more groups. demonstrates the effect of the number of GQA groups on inference speed. For larger models the memory band width overhead from the KV cache is less con straining (Shazeer, 2019), while the reduction in key-value size is sharper due to the increased number of heads. As a result, increasing the number of groups from MQA only results in modest slow downs initially, with increasing cost as we move closer to MHA. We selected 8 groups as a favor able middle ground.

    Your Image Description

    shows how performance varies with uptraining proportion for T5 XXL with MQA and GQA. First, we note that GQA already achieves reasonable performance after conversion while MQA requires uptraining to be useful. Both MQA and GQA gain from 5% uptraining with diminishing returns from 10%.

    MHA enables a nuanced understanding of the relationships between different parts of the input. Nevertheless, this complexity comes at a cost — a significant demand on memory bandwidth, especially during decoder inference. In multi-query attention, we average the heads for keys and values so that all query heads share the same key and value head. This is achieved by replicating the mean-pooled “head” H times, where H is the number of query heads. However, MQA is not without its drawbacks. The reduced complexity can lead to quality degradation and training instability. Grouped-query attention (GQA) is a simple approach that blends elements of multi-head attention (MHA) and multi-query attention (MQA) to create a more efficient attention mechanism.

LLAMA_ATTENTIONS_CLASSES = {
    "GQA": LlamaScalableGroupedQueryAttention,
    "MHA": MultiHeadAttention,
    "MQA": MultiQueryAttention,
}

installs:

install the requirements libraries:

pip install requirements

or

pip install pytorch transformers

clone the repo

git clone https://github.com/Esmail-ibraheem/X-Llama.git

run the download shell file to download the llama 2 weights

.\download.sh

after downloading the weights, run the inference code:

python inference.py

Citation:

@article{Gumaan2024-X-Llama,
  title   = "X-Llama",
  author  = "Gumaan, Esmail",
  code = "https://github.com/Esmail-ibraheem/X-Llama/tree/main",
  year    = "2024",
  month   = "Agu",
  url     = "https://github.com/Esmail-ibraheem"
}

Notes and Acknowledgments:

I developed this project to enhance my skills in large language models and transformers. I built the Llama model from scratch and implemented various features, including multiple attentions. Feel free to suggest any additional features you'd like, such as flash attention or related concepts. This project integrates multiple research papers.

papers:

other