File size: 3,491 Bytes
5e27996 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | ## Project Structure
```
MSA/
├── scripts/
│ ├── run_benchmarks.sh # Run inference on benchmarks
│ ├── calculate_llm_score.sh # LLM-based answer evaluation
│ └── resave_model.sh # Convert base model to MSA format
└── src/
├── msa/ # Core MSA implementation
│ ├── configuration_msa.py # MSA model configuration
│ ├── memory_sparse_attention.py # MemorySparseAttention layer
│ ├── model.py # MSAForCausalLM / MSAModel
│ └── generate.py # Generation logic
├── config/
│ └── memory_config.py # GenerateConfig, ModelConfig, MemoryConfig
├── evaluation/
│ └── llm_judge.py # LLM-based evaluation metrics
├── app/
│ └── benchmark.py # Benchmark runner
├── utils/ # GPU workers, caching, templates, etc.
├── msa_service.py # Multi-GPU inference engine (MSAEngine)
├── prefill.py # Stage 1 prefill worker
├── benchmarks.py # Benchmark registry & specs
└── types.py # Core type definitions
```
## Installation
**1. Create conda environment**
```bash
conda create -n msa python=3.12 -y
conda activate msa
```
**2. Install dependencies**
```bash
pip install -r requirements.txt
```
<details>
<summary>requirements.txt</summary>
```
torch==2.6.0
torchvision==0.21.0
transformers==4.51.3 # exact version required
accelerate==1.0.1
liger_kernel==0.5.10
huggingface_hub==0.31.4 # must stay <1.0 for transformers 4.51.3
datasets==3.1.0
lmdb==1.6.2
tqdm==4.67.1
numpy==1.26.4
pillow==11.2.1
packaging==25.0
nvidia-ml-py==12.575.51 # provides the `pynvml` module
openai==1.79.0
```
</details>
**3. Install Flash Attention**
Option A — build from source:
```bash
pip install flash-attn==2.7.4.post1 --no-build-isolation
```
Option B — install prebuilt wheel (CUDA 12, Python 3.12):
```bash
wget -P /tmp https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
pip install /tmp/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
rm /tmp/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl
```
## Download
**1. Download model**
```bash
mkdir ckpt
pip install -U huggingface_hub==0.31.4
export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download --resume-download Anoy123423123/MSA-4B --local-dir ckpt/MSA-4B
```
**2. Download benchmarks**
Benchmark data is hosted on [Anoy123423123/MSA-RAG-BENCHMARKS](https://huggingface.co/datasets/Anoy123423123/MSA-RAG-BENCHMARKS) and will be automatically downloaded to `data/` on first run, based on the benchmarks specified in `scripts/run_benchmarks.sh`. No manual download is needed.
## Quick Start
**1. Run inference on benchmarks**
```bash
bash scripts/run_benchmarks.sh eval_benchmark
```
**2. Compute LLM-based scores**
```bash
bash scripts/calculate_llm_score.sh eval_benchmark
```
## Supported Benchmarks
| Category | Benchmark |
|---|---|
| Multi-hop QA | `2wikimultihopqa`, `hotpotqa`, `musique` |
| Single-hop QA | `nature_questions`, `triviaqa_06M`, `triviaqa_10M`, `msmarco_v1`, `dureader`, `ms_100M`, `hipporag_narrative`, `hipporag_popqa` |
|