Llama-3-Gradient-1048k

Join our custom agent and long context (262k-1M+) waitlist: https://forms.gle/L6TDY7dozx8TuoUv7

Gradient incorporates your data to deploy autonomous assistants that power critical operations across your business. If you're looking to build custom AI models or agents, email us a message contact@gradient.ai. For more info see our end-to-end development service for custom LLMs and AI systems

This model extends LLama-3 8B's context length from 8k to > 1040K, developed by Gradient, sponsored by compute from Crusoe Energy. It demonstrates that SOTA LLMs can learn to operate on long context with minimal training by appropriately adjusting RoPE theta. We trained on 830M tokens for this stage, and 1.4B tokens total for all stages, which is < 0.01% of Llama-3's original pre-training data.

Update (5/3): We further fine-tuned our model to strengthen its assistant-like chat ability as well.

Updated NIAH result:

RULER evals:

• Our model is behind only GPT-4 and Yi in the retrieval and Q&A tasks
• It’s the smallest parameter model to rank in the top 7 overall

Approach:

• meta-llama/Meta-Llama-3-8B-Instruct as the base
• NTK-aware interpolation [1] to initialize an optimal schedule for RoPE theta, followed by empirical RoPE theta optimization
• Progressive training on increasing context lengths, similar to Large World Model [2] (See details below)

Infra:

We build on top of the EasyContext Blockwise RingAttention library [3] to scalably and efficiently train on contexts up to 1048k tokens on Crusoe Energy high performance L40S cluster.

Notably, we layered parallelism on top of Ring Attention with a custom network topology to better leverage large GPU clusters in the face of network bottlenecks from passing many KV blocks between devices. This gave us a 33x speedup in model training (compare 524k and 1048k to 65k and 262k in the table below).

Data:

For training data, we generate long contexts by augmenting SlimPajama. We also fine-tune on a chat dataset based on UltraChat [4], following a similar recipe for data augmentation to [2].

Progressive Training Details:

Evaluation:

EVAL_MAX_CONTEXT_LENGTH=1040200
EVAL_MIN_CONTEXT_LENGTH=100
EVAL_CONTEXT_INTERVAL=86675
EVAL_DEPTH_INTERVAL=0.2
EVAL_RND_NUMBER_DIGITS=8

HAYSTACK1:
    EVAL_GENERATOR_TOKENS=25

HAYSTACK2:
    EVAL_CONTEXT_INTERVAL=173350
    EVAL_GENERATOR_TOKENS=150000

HAYSTACK3:
    EVAL_GENERATOR_TOKENS=925000

All boxes not pictured for Haystack 1 and 3 are 100% accurate. Haystacks 1,2 and 3 are further detailed in this blog post.

Quants:

• GGUF by Crusoe. Note that you need to add 128009 as special token with llama.cpp.
• MLX-4bit
• Ollama
• vLLM docker image, recommended to load via --max-model-len 32768
• If you are interested in a hosted version, drop us a mail below.

The Gradient AI Team

https://gradient.ai/

Gradient is accelerating AI transformation across industries. Our AI Foundry incorporates your data to deploy autonomous assistants that power critical operations across your business.

Contact Us

Drop an email to contact@gradient.ai

Citation:

@misc{gradientlongcontextllama3,
  title={Llama 3 Gradient: A series of long context models},
  author={Leonid Pekelis and Michael Feil and Forrest Moret and Mark Huang and Tiffany Peng},
  year={2024},
  url = {https://gradient.ai/blog/scaling-rotational-embeddings-for-long-context-language-models},
  doi = { 10.57967/hf/3372 },
}

References

[1] Peng, Bowen, et al. "Yarn: Efficient context window extension of large language models." arXiv preprint arXiv:2309.00071 (2023).

[2] Liu, Hao, et al. "World Model on Million-Length Video And Language With RingAttention." arXiv preprint arXiv:2402.08268 (2024).

[3] https://github.com/jzhang38/EasyContext

[4] Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun, and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations. arXiv preprint arXiv:2305.14233, 2023.

Base Model

Model Details

Meta developed and released the Meta Llama 3 family of large language models (LLMs), a collection of pretrained and instruction tuned generative text models in 8 and 70B sizes. The Llama 3 instruction tuned models are optimized for dialogue use cases and outperform many of the available open source chat models on co