README.md

---
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.