yuzhenm/DSA-Qwen2-7B-Instruct

DSA on Qwen2-7B-Instruct

DeepSeek Sparse Attention (DSA) trained on Qwen2-7B-Instruct, as a baseline for the SSA paper's Table 3. Four checkpoints in one repository, one per training stage and compression ratio.

These are new baseline rows, not reproductions: Section 4.1 of the SSA paper excludes NSA, DSA and MoBA on purpose, so there is no published number to check them against. The correctness argument is the matched protocol, not agreement with prior work.

Contents

Folder	Stage	dsa_mode	dsa_top_k	Ratio	Trains
warmup/	1	dense_warmup	512 (unused)	—	indexer only, backbone frozen, attention dense
cpt-8x/	2	sparse	1024	8.16×	backbone + indexer, from warmup/
cpt-16x/	2	sparse	512	16.32×	backbone + indexer, from warmup/
sft-16x/	finetune	sparse	512	16.32×	from cpt-16x/

code/ holds the training and inference code for all four, trimmed to the DSA import closure — every module imports with only that folder on PYTHONPATH. See code/README.md for the exact commands.

Stage 1 does not depend on top_k, so a single 340-step warm-up was shared by both compression ratios. sft-8x was not trained: that job was cancelled.

All four are bf16, vocab_size 152064, and keep Qwen2-7B-Instruct's tokenizer unchanged — no gist tokens, no resized embeddings.

Compression ratios

DSA's budget is a constant number of keys per query, while the dense-causal budget it is measured against grows with sequence length. The constants here are anchored to the measured mean visible KV over the nine LongBench tasks, 8,358, so top_k 1024 and 512 realize 8.16× and 16.32× at that mean length. The realized ratio varies per example — LongBench spans 3.6K to 40.5K tokens — so report the measured KV per query alongside any number taken from these checkpoints.

Method

A lightning indexer scores every (query, key) pair,

s(q,k) = sum_h  w_h(q) * ReLU( q_h . k / sqrt(d) )

with multi-head query projections, a shared single-head key projection with a LayerNorm, and query-dependent head weights. The output is one score per (query, key), shared by every attention head in the layer; the top-dsa_top_k keys are attended. The indexer adds dsa_index_heads × dsa_index_head_dim = 8 × 64 parameters per layer, which MoBA and the full-attention references do not carry.

Training follows DeepSeek-V3.2-Exp: stage 1 trains the indexer alone against the KL between its distribution and the main attention's, with the backbone frozen; stage 2 unfreezes the backbone and switches attention to top-k. The two parameter groups never mix gradients — the indexer reads detached hidden states, and top-k is not differentiable.

Measured alignment KL: stage 1 fell 3.65 → 0.58 over 340 steps, and stage 2 opened at 0.65 (8×) and 0.66 (16×), continuous across the handoff. Stage 2 finished at 0.454 (8×) and 0.423 (16×).

Training

Corpus: the SSA project's own continued-pretraining and finetuning corpora with the gist tokens deleted, so every method is optimized on an identical set of loss-bearing positions.

	steps	epochs	lr	schedule	warmup	weight decay
Continued pretraining	11,372	1	1e-5	cosine	0.05	0.1
Finetuning	3,177	1	5e-6	cosine	0.1	0.0

Global batch 8 (1 per device × 8 ranks, accumulation 1) on 8×H100, bf16, DeepSpeed ZeRO-2, gradient checkpointing, length-grouped batching, seed 42.

LongBench

Nine tasks, full test sets (200 examples each; MultiFieldQA-en has 150). Greedy decoding, inputs truncated from the middle to 31,500 tokens, dsa_decode_mode dsa (sparse at prefill and decode). The routing budget is read from the checkpoint's config.json rather than re-specified at evaluation time.

Model	NrtvQA	Qasper	MF-en	HotpotQA	2WikiMQA	Musique	SAMSum	TriviaQA	GovRep	Avg.
Qwen2-7B-Instruct (dense)	26.12	42.64	44.39	54.85	55.41	36.19	44.95	90.16	36.05	47.86
cpt-8x	25.91	40.41	47.48	58.38	54.71	34.55	44.30	89.11	32.25	47.46
cpt-16x	20.98	40.06	48.27	58.91	53.94	34.52	43.98	89.04	32.02	46.86
sft-16x	21.77	41.96	52.53	50.16	52.63	30.00	42.39	88.89	34.57	46.10

Halving the budget costs DSA 0.60 points (47.46 → 46.86). cpt-16x beats the MoBA baseline at 8× (45.97) on half the budget.

Loading

The architecture is DsaQwen2ForCausalLM, which is not in transformers. AutoModelForCausalLM.from_pretrained will not resolve it; load it through code/dsa_qwen2.py, which ships in this repository:

import sys
from huggingface_hub import snapshot_download

root = snapshot_download("yuzhenm/DSA-Qwen2-7B-Instruct",
                         allow_patterns=["code/**", "cpt-16x/*"])
sys.path.insert(0, f"{root}/code")

from baseline_models import load_dsa_model   # noqa: E402
model = load_dsa_model(f"{root}/cpt-16x")    # arch is read from config.json

config.json already carries the routing settings, so evaluation runs at the budget the model was trained at. Passing --dsa_top_k at evaluation time is an override, not a requirement, and passing it inconsistently is the likeliest way to produce a wrong number.

Caveats

1. Masked SDPA, not a fused sparse kernel. Quality is unaffected — the mask defines the same attention — but latency is not comparable to a kernel implementation, and nothing measured here supports an efficiency claim.
2. The ratio is anchored to one length. See above.
3. The indexer adds parameters the baselines it is compared against do not have.
4. training_args.bin is a pickle written by transformers; it is included for provenance and needs weights_only=False to read.