fix

Author: Nathan

README.md

@@ -1,2 +1,79 @@
-# qfilters
-Repository for the Q-Filters method
+# Q-Filters: Leveraging Query-Key Geometry for Efficient Key-Value Cache Compression
+[![arXiv](https://img.shields.io/badge/arXiv-1234.56789-b31b1b.svg)](https://arxiv.org/abs/1234.56789) 
+
+![Q-Filters Demo GIF](assets/qfilters_demo.gif)
+
+## Setup
+1. Install required libraries in a virtual environment:
+```bash
+python -m virtualenv venv
+source venv/bin/activate
+pip install -r requirements.txt
+````
+2. Configure HuggingFace\'s environment:
+```bash
+export HF_DATASETS_CACHE=<path_to_hf_cache>
+export HF_HOME=<path_to_hf_cache>
+export HF_TOKEN=<hf_token>
+```
+
+## Generate with Q-Filters
+Here is an example of how to use Q-Filters in a generation setup:
+```python
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from src.hf_cache import QFiltersCache
+from datasets import load_dataset
+
+model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map="auto",
+    low_cpu_mem_usage=True,
+    torch_dtype="bfloat16"
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+streamer = TextStreamer(tokenizer)
+
+question = """What is the probability of two integers selected at random having a greatest common divisor of 1."""
+input_text = f"<|User|>{question}<|Assistant|><think>\n"
+
+inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+
+past_key_values = QFiltersCache(
+    window_length=64,
+    max_length=128, 
+    model_name=model_name
+)
+
+out = model.generate(
+    **inputs,
+    do_sample=True, 
+    temperature=0.5, 
+    max_new_tokens=4096, 
+    past_key_values=past_key_values, 
+    streamer=streamer
+)
+```
+
+## Compute Q-Filters for a new model
+1. Verify that the target model does not already have [pre-computed Q-Filters](https://huggingface.co/collections/nthngdy/q-filters-67a4994dcb302a3d37f3d119).
+2. Use the `make_filters.py` script to generate the filters. For instance:
+```bash
+python make_filters.py \
+--model_name deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
+--model_cls Qwen2ForCausalLM \
+--max_seq_len 2048 \
+--num_sequences 10 \
+--num_svd_samples 3000 \
+--dataset_name PatrickHaller/fineweb-1B \
+--save_mode disk \
+# --save_mode hub \
+# --save_mode hub+disk \
+# --hf_user_id nthngdy \
+--save_dir ../filters
+```
+3. For Q-Filters saved on disk, you can upload them later using this command:
+```bash
+huggingface-cli upload path_to_hf_repo path_to_local_qfilters .
+```

example.py

@@ -0,0 +1,39 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+from src.hf_cache import QFiltersCache, KNormCache
+from datasets import load_dataset
+
+model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map="auto",
+    low_cpu_mem_usage=True,
+    torch_dtype="bfloat16"
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+streamer = TextStreamer(tokenizer)
+
+question = """What is the probability of two integers selected at random having a greatest common divisor of 1."""
+input_text = f"<|User|>{question}<|Assistant|><think>\n"
+
+inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+
+# past_key_values = QFiltersCache(
+#     window_length=64,
+#     max_length=128,
+#     model_name=model_name
+# )
+past_key_values = KNormCache(
+    window_length=64,
+    max_length=128,
+)
+
+
+out = model.generate(
+    **inputs,
+    do_sample=True, 
+    temperature=0.5, 
+    max_new_tokens=4096, 
+    past_key_values=past_key_values, 
+    streamer=streamer
+)

make_filters.py

@@ -0,0 +1,112 @@
+from tqdm import tqdm
+import pickle
+import os
+import argparse
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from datasets import load_dataset
+import torch
+
+import modeling as M
+from utils import QFilters
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument("--model_name")
+parser.add_argument("--model_cls")
+parser.add_argument("--max_seq_len", type=int, default=2048)
+parser.add_argument("--num_sequences", type=int, default=20)
+parser.add_argument("--num_svd_samples", type=int, default=3000)
+parser.add_argument("--filter_suffix", default="")
+parser.add_argument("--torch_dtype", default="bfloat16")
+
+parser.add_argument("--dataset_name")
+parser.add_argument("--dataset_config", default="default")
+parser.add_argument("--dataset_split", default="train[:1000]")
+
+parser.add_argument("--save_mode", default="disk")
+parser.add_argument("--save_dir", default="")
+parser.add_argument("--hf_user_id", default="")
+
+
+args = parser.parse_args()
+
+model_name = args.model_name
+model_cls = getattr(M, args.model_cls)
+max_seq_len = args.max_seq_len
+num_sequences = args.num_sequences
+num_svd_samples = args.num_svd_samples
+filter_suffix = args.filter_suffix
+torch_dtype = args.torch_dtype
+
+dataset_name = args.dataset_name
+dataset_config = args.dataset_config
+dataset_split = args.dataset_split
+
+save_mode = args.save_mode
+save_dir = args.save_dir
+hf_user_id = args.hf_user_id
+
+if "disk" in save_mode and not save_dir:
+    raise ValueError("In 'disk' or 'disk+hub' save modes, a '--save_dir' must be provided.")
+
+if "hub" in save_mode and not hf_user_id:
+    raise ValueError("In 'hub' or 'disk+hub' save modes, a '--hf_user_id' must be provided.")
+
+
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = model_cls.from_pretrained(
+    model_name, attn_implementation="flash_attention_2", device_map="auto",  low_cpu_mem_usage=True, torch_dtype=torch_dtype)
+
+model = model.eval()
+
+dataset = load_dataset(dataset_name, dataset_config, split=dataset_split)
+
+
+with torch.no_grad():
+    decoder = getattr(model, "gpt_neox", getattr(model, 'model', None))
+    svd_filters = [[] for _ in range(len(decoder.layers))]
+    sample_count = 0
+    num_k_heads = None
+
+    for i, sample in tqdm(enumerate(dataset)):
+
+        tokens = tokenizer(sample["text"], return_tensors="pt")
+        if tokens.input_ids.shape[-1] < max_seq_len:
+            continue
+        sample_count+=1
+        input_ids = tokens.input_ids[:, :max_seq_len].to(device)
+        if sample_count < num_sequences:
+            with torch.autocast(device_type=device, dtype=torch.bfloat16):
+                out_repr = model(input_ids).past_key_values
+            for j, (query, key) in enumerate(out_repr):
+                num_k_heads = key.shape[1]
+                svd_filters[j].append(query.flatten(0, 1).cpu())
+        else:
+            break
+
+    del model
+
+    for f_id, el in enumerate(svd_filters):
+        stacked_el = torch.stack(el, 1).flatten(1, 2)
+        idx = torch.argsort(torch.rand(stacked_el.shape[1], device=stacked_el.device))[:num_svd_samples]
+        stacked_el = stacked_el[:, idx].cuda()
+        u,s,vh = torch.linalg.svd(stacked_el.float())
+        svd_sign = ((u[..., 0]>0).float().mean(-1) > 0.5).float()*2-1
+        svd_filter_q = -svd_sign[:, None] * vh[..., 0, :]
+        svd_filters[f_id] = svd_filter_q.reshape(num_k_heads, -1, svd_filter_q.shape[-1]).mean(-2)
+
+    svd_filters = torch.nn.Parameter(torch.stack(svd_filters))
+    q_filters = QFilters(*svd_filters.shape)
+    q_filters.q_filters = svd_filters
+
+    model_suffix = model_name.split("/")[-1]
+    filter_savename = f"{model_suffix}_qfilt{'_' + filter_suffix if filter_suffix else ''}"
+    if "disk" in save_mode:
+        q_filters.save_pretrained(f"{save_dir}/{filter_savename}")
+    if "hub" in save_mode:
+        q_filters.push_to_hub(f"{hf_user_id}/{filter_savename}")

modeling/__init__.py

@@ -0,0 +1,5 @@
+from modeling.modeling_llama import LlamaForCausalLM
+from modeling.modeling_qwen2 import Qwen2ForCausalLM
+from modeling.modeling_olmo2 import Olmo2ForCausalLM
+from modeling.modeling_phi3 import Phi3ForCausalLM
+from modeling.modeling_mistral import MistralForCausalLM