Samba is a simple yet powerful hybrid model with an unlimited context length. Its architecture is frustratingly simple:
Samba = Mamba + MLP + Sliding Window Attention + MLP stacking at the layer level.
Our largest model, Samba-3.8B, is trained on 3.2 trillion tokens from the Phi3 dataset, outperforming Phi3-mini on major benchmarks (e.g. MMLU, GSM8K and HumanEval) by a large margin. Samba can also achieve perfect long-context retrieval ability with minimal instruction tuning, while still maintaining its linear complexity with respect to sequence length. This ability leads to the impressive performance of Samba-3.8B-instruct on downstream tasks such as long-context summarization.
We report 5-shot accuracy for MMLU, 8-shot CoT accuracy for GSM8K, 0-shot pass@1 for HumanEval and ROUGE-L for both GovReport and SQuALITY.
Our training infrastructure on SlimPajama is a modified version of TinyLlama and LitGPT. One can easily specify different architectual configurations through modifying the model_name and the config file which includes tons of baseline architectures mentioned in the paper. Our RetNet and GLA implementations are from the awesome Flash Linear Attention repository.
Please follow the Dockerfile to setup the environment. The data preparation mainly follows TinyLlama except that we only use the SlimPajama dataset.
Download the Slimpajama dataset to your chosen directory.
cd /path/to/dataset git lfs install git clone https://huggingface.co/datasets/cerebras/SlimPajama-627B
The SlimPajama dataset takes 893GB diskspace. Use the provided scripts to tokenize the datasets and divide them into chunks.
python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama --destination_path data/slim --split validation --percentage 1.0 python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama --destination_path data/slim --split train --percentage 1.0
You are now ready to launch a job!
The following script trains a default Samba-421M model on a single node of 8 GPUs with 20B tokens.
torchrun --nnodes=1 --nproc_per_node=8 --rdzv_id=samba-421M --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} pretrain.py --train_data_dir data/slim --val_data_dir data/slim
You can modify model_name to "Samba_1.3B" and train_config to "tsz512x4k_100B" for training a Samba-1.3B model with 100B tokens. We assume that you have 8 nodes each with 8 GPUs, and you can modify the number of nodes for training on fewer gpus.
We leverage lm-evaluation-harness for the evaluation of our pretrained models. We only support non-generation based tasks for now.
pip install lm-eval
python eval.py --model Samba \
--model_args pretrained=path/to/ckpt.pth,config="Samba_1.3B" \
--tasks lambada_openai,arc_easy,winogrande,hellaswag,piqa --device cuda:0 --batch_size 1 --trust_remote_code
If you find our work useful, please consider citing:
@article{ren2024samba,
title={Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling},
author={Liliang Ren and Yang Liu and Yadong Lu and Yelong Shen and Chen Liang and Weizhu Chen},
journal = {arXiv preprint},
year={2024},
url={https://arxiv.org/abs/2406.07522}
}
Liliang Ren (liliangren@microsoft.com)
RetroSearch is an open source project built by @garambo | Open a GitHub Issue
Search and Browse the WWW like it's 1997 | Search results from DuckDuckGo
HTML:
3.2
| Encoding:
UTF-8
| Version:
0.7.4