You can deploy this model and start earning money today!

Anoy123423123/
MSA-Code

Name: MSA-Code
Availability: InStock

msa

mit

Views

No views yet

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`
Benchmark data	`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. 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

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 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 for the project layout and the prebuilt flash-attn wheel.

Citation

Withheld during anonymous peer review.

License

MIT.