Qwen3.5-2B-Longbook2 (accuracy-recovered)

変換元モデル は Qwen3.5-2B の全 full-attention 層を Gated DeltaNet (GDN) に置き換えた all-GDN モデルです。置換で落ちた精度を、 元モデル Qwen/Qwen3.5-2B を教師とした蒸留で回復したものがこのリポジトリです。 追加の知識データでの SFT は行っていません(復元のみ)。

• 全層が線形アテンション → KV キャッシュが増えず、長文でもメモリ・速度が一定
• 標準の transformers でそのまま読み込めます(trust_remote_code 不要)
• flash-linear-attention を入れると高速カーネルが使われます

復元方法

1. Stage 1 (層ごとの蒸留): 置換した各 GDN 層の出力を元の full-attention 層の出力に合わせる (相対 MSE + コサイン)。層間の入力は detach し、全層を並列に学習。
2. Stage 2 (end-to-end): 教師 logits への KL + 隠れ状態損失で全体を微調整。
3. 学習した重みを元の safetensors シャードの該当テンソルだけ差し替えて出力。

fp32 マスター重み + bf16 autocast、fused AdamW。

結果

(評価値は未記入)

(同一の評価テキスト・同一条件で測定。教師は Qwen/Qwen3.5-2B)

使い方

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "j-llm/Qwen3.5-2B-Longbook2"
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