Mistral Supra

🚀 Mistral-SUPRA

This model is initialized from the weights of the Mistral-7B transformer model and up-trained into a linear RNN. It's an accompanying model of the paper Linearizing Large Language Models, which details the process of converting a softmax transformer into a linear transformer that can function as both a transformer and a recurrent model at inference time.

🚀 Quick Start

This model was trained using OpenLM. The weights have been converted to be compatible with HuggingFace.

To use the model, you need to first pip install our fork of OpenLM.

pip install git+https://github.com/tri-ml/linear_open_lm.git

Import the OpenLM classes with

from open_lm.open_lm_hf import *

The model can then be loaded normally using AutoTokenizer and AutoModelForCausalLM as follows:

from open_lm.open_lm_hf import *
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("tri-ml/mistral-supra")
model = AutoModelForCausalLM.from_pretrained("tri-ml/mistral-supra")

inputs = tokenizer(["Machine learning is"], return_tensors="pt")
gen_kwargs = {"max_new_tokens": 50, "top_p": 0.8, "temperature": 0.8, "do_sample": True, "repetition_penalty": 1.1}
output = model.generate(inputs['input_ids'], **gen_kwargs)
output = tokenizer.decode(output[0].tolist(), skip_special_tokens=True)
print(output)
# Machine learning is a branch of artificial intelligence (AI) that enables computers to learn from experience without being explicitly programmed. Machine learning is used in a wide range of applications, including spam filtering, image recognition, speech recognition, and computer-based medical diagnosis

The Mistral-SUPRA model can be used both in parallel mode and in recurrent mode. If use_cache is set to False for model.generate(...), then it will use parallel mode; otherwise, it will use recurrent mode. The recurrent model uses xformers and requires the inputs and models to be loaded to GPU.

# Recurrent mode
output = model.to('cuda').generate(inputs['input_ids'].to('cuda'), use_cache=True, **gen_kwargs)

# Parallel mode
output = model.to('cuda').generate(inputs['input_ids'].to('cuda'), use_cache=False, **gen_kwargs)

✨ Features

• Initialized from the weights of the Mistral-7B transformer model and up-trained into a linear RNN.
• Can function as both a transformer and a recurrent model at inference time.
• The linear attention code can be found at https://github.com/TRI-ML/linear_open_lm/

📦 Installation

To use the model, you need to first pip install our fork of OpenLM.

pip install git+https://github.com/tri-ml/linear_open_lm.git

📚 Documentation

Model Details

Training Details

• Mistral-SUPRA was trained using AWS SageMaker on 128 H100 80GB GPUs.
• Training on 100B tokens finished in 1.5 days.

Performance Evaluation

Our evaluations were done using the Eleuther LM Eval Harness repo.

Below we report the performance of Mistral-SUPRA compared to other similarly sized models.

🔧 Technical Details

This model was initialized from the weights of the Mistral-7B transformer model and up-trained into a linear RNN. Our linear attention code can be found at https://github.com/TRI-ML/linear_open_lm/

We detail our process of converting a softmax transformer into a linear transformer, which at inference time can function as both a transformer and a recurrent model in our paper Linearizing Large Language Models.

📄 License

This model is licensed under Apache License, Version 2.0.

How to Cite

If you use this model, please cite our paper on Linearizing Large Language Models.

@article{Mercat2024Linearizing,
  title={Linearizing Large Language Models},
  author={Jean Mercat and Igor Vasiljevic and Sedrick Keh and Kushal Arora and Achal Dave and Adrien Gaidon and Thomas Kollar},
  year={2024},
  journal={arXiv preprint arXiv:2405.06640},
}

Citations

OpenLM

@misc{open_lm,
  author = {Gururangan, Suchin and Wortsman, Mitchell and Gadre, Samir Yitzhak and Dave, Achal and Kilian, Maciej and Shi, Weijia and Mercat, Jean and Smyrnis, Georgios and Ilharco, Gabriel and Jordan, Matt and Heckel, Reinhard and Dimakis, Alex and Farhadi, Ali and Shankar, Vaishaal and Schmidt, Ludwig},
  title = {{open_lm}:  a minimal but performative language modeling (LM) repository},
  year = {2023},
  note = {GitHub repository},
  url = {https://github.com/mlfoundations/open_lm/}
}