Qwen3.5-2B-SpeedX (accuracy-recovered)
summerMC/Qwen3.5-2B-SpeedX は Qwen3.5-2B の全 full-attention 層を
Gated DeltaNet (GDN) に置き換えた all-GDN モデルです。置換で落ちた精度を、
元モデル Qwen/Qwen3.5-2B を教師とした蒸留で回復したものがこのリポジトリです。
追加の知識データでの SFT は行っていません(復元のみ)。
- 全層が線形アテンション → KV キャッシュが増えず、長文でもメモリ・速度が一定
- 標準の
transformersでそのまま読み込めます(trust_remote_code不要) flash-linear-attentionを入れると高速カーネルが使われます
復元方法
- Stage 1 (層ごとの蒸留): 置換した各 GDN 層の出力を元の full-attention 層の出力に合わせる (相対 MSE + コサイン)。層間の入力は detach し、全層を並列に学習。
- Stage 2 (end-to-end): 教師 logits への KL + 隠れ状態損失で全体を微調整。
- 学習した重みを元の safetensors シャードの該当テンソルだけ差し替えて出力。
fp32 マスター重み + bf16 autocast、fused AdamW。
結果
(評価値は未記入)
(同一の評価テキスト・同一条件で測定。教師は Qwen/Qwen3.5-2B)
使い方
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "j-llm/Qwen3.5-2B-SpeedX2"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.bfloat16, device_map="cuda").eval()
ids = tok.apply_chat_template([{"role": "user", "content": "日本の首都はどこ？"}],
add_generation_prompt=True, return_tensors="pt", return_dict=True, enable_thinking=False).to(model.device)
out = model.generate(**ids, max_new_tokens=128, use_cache=False) # 下の注意参照
print(tok.decode(out[0, ids["input_ids"].shape[1]:], skip_special_tokens=True))
注意
- 全層が線形アテンションのため、現行
transformersの標準Cacheはget_seq_lengthでValueErrorを出すことがあります。その場合はuse_cache=False(遅いが確実)にするか、get_seq_lengthを例外時に自前カウンタへフォールバックさせるパッチを当ててください。 - 蒸留は教師の挙動への近似であり、元モデルと完全には一致しません。特に推論・数学・コードは元モデルより劣る可能性があります。
- ライセンスは元モデル(Qwen/Qwen3.5-2B)に従います。
謝辞
変換モデル: summerMC / 元モデル: Qwen Team