AnchorMass-150M (32K SFT)

AnchorMass は、超長文処理時でも Key-Value (KV) キャッシュの増大を防ぎ、メモリ消費を一定サイズに抑えながら推論できる長文特化の言語モデルです。

本リポジトリでは、158.6M パラメータ構成で 32K カリキュラム長文 SFT を完了したチェックポイントを提供します。

主な特徴

KVキャッシュ容量の上限固定: 32K〜256K の超長文でも、KVキャッシュを 約 259 MiB に一定化（256K 時点で標準モデル比 41.6倍削減）。
安定したデコード速度: キャッシュが肥大化しないため、超長文でも 75〜76 tok/s を維持（256K 時点で標準モデルの 約4.4倍高速）。
外れ値保護ブロック圧縮: 重要なトークン（外れ値）を自動分離して保護し、背景情報のみを効率的に圧縮。

性能比較 (158.6M パラメータ・実測値)

RTX 6000 Ada (48GB) 上での標準 Transformer（Standard-150M）との対照比較：

系列長	Standard KVキャッシュ	AnchorMass KVキャッシュ	Standard デコード速度	AnchorMass デコード速度
4K	168 MiB	36 MiB (4.7倍削減)	205.7 tok/s	75.7 tok/s
8K	336 MiB	67 MiB (5.0倍削減)	204.3 tok/s	75.9 tok/s
16K	672 MiB	131 MiB (5.1倍削減)	202.9 tok/s	75.7 tok/s
32K	1,344 MiB	258 MiB (5.2倍削減)	147.5 tok/s	75.2 tok/s
64K	2,689 MiB	259 MiB (10.4倍削減)	80.5 tok/s	75.8 tok/s
128K	5,377 MiB	259 MiB (20.8倍削減)	38.0 tok/s	76.3 tok/s (2.0倍高速)
256K	10,752 MiB	259 MiB (41.6倍削減)	17.6 tok/s	76.6 tok/s (4.4倍高速)

注意: Prefill（プロンプト処理）は逐次圧縮を行うため標準モデルより時間を要します。128K・256K において長文生成を行う場合や、Prefill 後のキャッシュを再利用する場合に総スループットが有利となります。

モデル仕様

パラメータ数: 158.6M (158,567,424)
層数: 14
隠れ層次元: 768 (Head数: 12)
中間層 (FFN): 2,048 (SwiGLU)
語彙数: 32,000 (SentencePiece)
位置エンコーディング: Dynamic NTK RoPE
グローバルキャッシュ上限: 8,192 スロット

使い方 (推論コード例)

git clone https://github.com/rsu-Suba/AnchorMass-Attention.git
cd AnchorMass
pip install -r requirements.txt

import torch
import yaml
import sentencepiece as spm
from anchormass.model import AnchorMassForCausalLM, AnchorMassConfig
from eval.utils import prefill_anchormass, generate_tokens

device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. モデルの準備
with open("anchormass/config.yaml") as f:
    cfg_dict = yaml.safe_load(f)["anchormass_150M"]
model = AnchorMassForCausalLM(AnchorMassConfig.from_dict(cfg_dict)).to(device=device, dtype=torch.bfloat16)

ckpt = torch.load("anchormass_150M_32k_sft.pt", map_location=device, weights_only=False)
sd = ckpt["model_state_dict"] if "model_state_dict" in ckpt else ckpt
model.load_state_dict({k.replace("._orig_mod.", ".").removeprefix("_orig_mod."): v for k, v in sd.items()}, strict=False)
model.eval()

# 2. トークナイザー
sp = spm.SentencePieceProcessor(model_file="tokenizer/tokenizer.model")

# 3. 超長文プロンプトの Prefill と生成
input_ids = torch.tensor([sp.encode("長大なテキスト...")], dtype=torch.long, device=device)
model.set_session_length(input_ids.shape[1])

past_kv, logits = prefill_anchormass(model, input_ids, chunk_size=1024, return_last_chunk_logits=True)
gen_tokens = generate_tokens(model, past_kv, logits, max_new_tokens=64, temperature=0.7)
print(sp.decode(gen_tokens))

ライセンス

Apache License 2.0

Downloads last month: 9

rsu
/

AnchorMass-Attention