| --- |
| license: mit |
| language: |
| - en |
| tags: |
| - msa |
| --- |
| # MSA: Memory Sparse Attention |
|
|
| *A scalable, end-to-end trainable latent-memory framework for 100M-token contexts.* |
|
|
| **Anonymous code release for peer review.** This repository contains the inference and evaluation code |
| used to produce the reported benchmark results. |
|
|
| | | | |
| |---|---| |
| | Code (this repo) | `Anoy123423123/MSA-Code` | |
| | Model weights | [`Anoy123423123/MSA-4B`](https://huggingface.co/Anoy123423123/MSA-4B) | |
| | Benchmark data | [`Anoy123423123/MSA-RAG-BENCHMARKS`](https://huggingface.co/datasets/Anoy123423123/MSA-RAG-BENCHMARKS) | |
|
|
| ## Requirements |
|
|
| Linux x86-64, Python 3.12, CUDA 12.x. **8× 80GB GPUs** for the default benchmark suite; the 100M-token |
| `ms_100M` setting runs on **2× A800** via Memory Parallel. ~9 GB disk for weights, ~0.6 GB for data. |
|
|
| `transformers` must be exactly **4.51.3** — the attention implementation targets that API and will not |
| load on 4.52+ or 5.x. The other pins in `requirements.txt` are the tested versions. |
|
|
| ## Quick Start |
|
|
| **1. Get the code** |
|
|
| ```bash |
| pip install -U "huggingface_hub==0.31.4" |
| # export HF_ENDPOINT=https://hf-mirror.com # uncomment if you need a mirror |
| huggingface-cli download Anoy123423123/MSA-Code --repo-type=model --local-dir MSA |
| cd MSA |
| ``` |
|
|
| **2. Install** |
|
|
| ```bash |
| conda create -n msa python=3.12 -y && conda activate msa |
| pip install -r requirements.txt |
| pip install flash-attn==2.7.4.post1 --no-build-isolation |
| ``` |
|
|
| `flash-attn` is separate because it needs `--no-build-isolation`. To skip the slow source build, see the |
| prebuilt wheel in [QUICK_START.md](./QUICK_START.md). Verify: |
|
|
| ```bash |
| python -c "import torch, transformers, flash_attn, pynvml; print(torch.__version__, transformers.__version__, flash_attn.__version__)" |
| # expected: 2.6.0 4.51.3 2.7.4.post1 |
| ``` |
|
|
| **3. Download the weights** |
|
|
| ```bash |
| huggingface-cli download --resume-download Anoy123423123/MSA-4B --local-dir ckpt/MSA-4B |
| ``` |
|
|
| This is the path `scripts/run_benchmarks.sh` expects (`model_path=ckpt/MSA-4B`). |
|
|
| **4. Run inference** |
|
|
| ```bash |
| bash scripts/run_benchmarks.sh eval_benchmark |
| ``` |
|
|
| Benchmark data is fetched automatically on first use into `./data/`. Results land in |
| `src/evaluation/outputs/eval_benchmark/` as one `.log` and one `.json` per benchmark. |
|
|
| **5. Score the answers** |
|
|
| ```bash |
| export OPENROUTER_API_KEY=<your key> |
| bash scripts/calculate_llm_score.sh eval_benchmark |
| ``` |
|
|
| The argument must match the log directory name from step 4. |
|
|
| ## Benchmarks |
|
|
| Configured in the `benchmarks` array at the top of `scripts/run_benchmarks.sh` as `name:batch_size`: |
|
|
| `2wikimultihopqa` `hotpotqa` `musique` `nature_questions` `triviaqa_06M` `triviaqa_10M` |
| `hipporag_narrative` `hipporag_popqa` `msmarco_v1` `dureader` `ms_100M` |
|
|
| `ms_100M` (the 100M-token setting) and `triviaqa_06M` are commented out by default. |
|
|
| **Batch size must be scaled with your GPU count** — the memory corpus is partitioned into one bucket per |
| visible GPU, and the defaults assume the 8 GPUs in `CUDA_VISIBLE_DEVICES`. |
|
|
| ## Notes |
|
|
| | Symptom | Fix | |
| |---|---| |
| | Errors loading the model | `transformers` must be exactly `4.51.3`. | |
| | `import pynvml` fails | The module comes from `nvidia-ml-py`; the separate deprecated `pynvml` distribution shadows it. `pip uninstall -y pynvml && pip install nvidia-ml-py==12.575.51` | |
| | CUDA OOM | Lower `batch_size` in `scripts/run_benchmarks.sh`. | |
| | Data download stalls | `export HF_ENDPOINT=https://hf-mirror.com` and retry; partial files resume. | |
|
|
| See [QUICK_START.md](./QUICK_START.md) for the project layout and the prebuilt `flash-attn` wheel. |
|
|
| ## Citation |
|
|
| Withheld during anonymous peer review. |
|
|
| ## License |
|
|
| MIT. |
|
|
