48Bのオープンモデルに 「語尾を全部『フィー』にする」LoRA を当てました。語尾は 0% → 99.2% で綺麗に入りました。
ただ、この記事の本題はそこではありません。
Moonshot が公開している modeling_kimi.py は推論専用で、そのままでは 1 step も学習が回りません。 assert not self.training から始まる 6つの壁 を全部突破する必要がありました。同じことをやろうとしている人の時間を節約するために、壁の内容と対処を全部書きます。
- コード・アダプタ・データ・生ログ: https://github.com/masafykun/kimi-fii-lora
- モデル: https://huggingface.co/masafy/kimi-fii-lora
結論を先に
| 項目 | 結果 |
|---|---|
| 語尾「フィー」の学習 | 成功（0% → 86.7〜99.2%） |
| 「レッサーパンダ好き？」で長文化 | 失敗（560字 → 185字と逆に短縮） |
| 学習時間 | 2 epoch / 158 step で 2時間39分（約60秒/step） |
| VRAMピーク | 41.65 GB（RTX A6000 48GB） |
| 総額 | $2.80（Vast.ai、4時間02分） |
なぜ Kimi なのか、そしてなぜ K3 ではないのか
最初は話題の Kimi K3（2.8T）でやろうとしましたが、断念しました。
- 推論だけで RAM+VRAM が 594GB〜1.6TB 必要（Unsloth の量子化ドキュメント基準）
- コミュニティの派生モデルは、abliterated 版・プルーニング版・テスト用ダミーが大半
特に注意したいのが Inferact/Kimi-K3-DSpark（4B）です。DL数が3万超で「K3の小型版」に見えますが、実体は vLLM の投機的デコーディング用のドラフトモデルで、非因果的アテンションの密5層が7トークンを並列にドラフトするだけの補助部品です。単体では文章を生成できません。 モデルカードを読まずにファインチューンの土台に選ぶと確実に詰みます。
REAP系（エキスパート枝刈り版）も見送りました。作者自身が「品質は未検証」と書いているうえ、枝刈りの重要度評価コーパスが コード40% / 英語30% / 中国語15% / その他15% という構成で、日本語は「その他」に埋没します。日本語用途の土台としては最悪です。
結局、Kimi 系で普通に扱えるのは moonshotai/Kimi-Linear-48B-A3B-Instruct だけでした。
- 48B総 / 活性3B（MoE 256エキスパート、top-8 + shared 1）
- 27層、KDA 20層 + フルアテンション(MLA) 7層のハイブリッド
- MIT ライセンス、transformers 対応
重みを1バイトも落とさずに LoRA の当てどころを決める
98GB をダウンロードする前に、model.safetensors.index.json（2MB）だけで全テンソル名が分かります。
import json, re, collections
idx = json.load(open("model.safetensors.index.json"))
patterns = collections.Counter()
for key in idx["weight_map"]:
p = re.sub(r"\.\d+\.", ".N.", key)
p = re.sub(r"experts\.N\.", "experts.{0-255}.", p)
patterns[p] += 1
結果はこうなりました。
27 model.layers.N.self_attn.q_proj.weight <- 全層
27 model.layers.N.self_attn.o_proj.weight <- 全層
20 model.layers.N.self_attn.k_proj.weight <- KDA層のみ
20 model.layers.N.self_attn.v_proj.weight <- KDA層のみ
6656 model.layers.N.block_sparse_moe.experts.{0-255}.w1.weight <- 触らない
26 model.layers.N.block_sparse_moe.gate.weight <- 触らない
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] に確定。MoEのエキスパート（256×26個）とルータには絶対に当てません。 前者はアダプタが2万個に膨れ、後者はルーティングそのものを壊します。
この事前予測は、後で実機にロードして named_modules() を数えたところ 完全に一致しました。数十分と数百円を節約できます。
なお peft は key.endswith("." + target) で照合するので、f_b_proj / g_b_proj が b_proj に誤マッチしないことも確認しておくと安全です。
6つの壁
ここからが本題です。
壁1: transformers 5.x では動かない
ImportError: cannot import name 'OutputRecorder' from 'transformers.utils.generic'
transformers>=4.56 とだけ書くと pip が 5.x を入れて死にます。上限を必ず付けてください。
transformers==4.57.1
壁2: auto_docstring が PEP 604 の型注釈で落ちる
AttributeError: 'types.UnionType' object has no attribute '__name__'
これは transformers 本体のバグです。_process_parameter_type が str(int | None) に "typing" が含まれないせいで誤った分岐に入り、types.UnionType に存在しない __name__ を触ります。インストール済みソースを grep しましたが、環境変数による無効化手段はありません。
docstring 生成はモデルの動作に一切関係しないので、モデルを読む前にデコレータを恒等関数に差し替えます。
import importlib
def _noop_decorator(obj=None, **kwargs):
if callable(obj):
return obj
return lambda fn: fn
tu = importlib.import_module("transformers.utils")
setattr(tu, "auto_docstring", _noop_decorator)
壁3: MoEルータの assert not self.training
# modeling_kimi.py
assert not self.training
scores_for_choice = scores.view(bsz * seq_len, -1)
scores_for_choice += self.e_score_correction_bias.unsqueeze(0) # ← in-place
...
topk_weight = scores.gather(1, topk_idx)
この assert を単に消すと壊れます。 scores_for_choice は scores のビューなので、+= が scores 本体を書き換えます。推論なら無害ですが、学習時は sigmoid の出力を in-place で潰すことになり autograd が死にます。assert はそれを防いでいました。
LoRA は attention にしか当てていませんが、勾配は MoE ブロックを通過して前段の層まで流れるので、ここを避けて通れません。
数値を変えずに autograd 安全にするには、2箇所をセットで直します。
scores_for_choice = scores_for_choice + self.e_score_correction_bias.unsqueeze(0) # out-of-place
...
topk_weight = scores_for_choice.gather(1, topk_idx) # gather元も変える
元コードは実質 scores := scores + bias してから gather しているので、gather 元も変えないと bias 抜きの重みになり挙動が変わります。 ここを片方だけ直すのが一番危ない罠です。
壁4: moe_infer に @torch.no_grad() が付いている
def forward(self, hidden_states):
...
if not self.training:
y = self.moe_infer(...)
else:
raise NotImplementedError("Training mode is not supported in KimiSparseMoeBlock")
@torch.no_grad() # ←
def moe_infer(self, x, topk_ids, topk_weight):
公開コードには微分可能な MoE 経路が存在しません。 自分で書きます。ループ構造は moe_infer と同一にして、違いを3点だけに絞りました。
-
@torch.no_grad()を付けない -
new_x[idxs] = outsの in-place 代入を逆順列 gather に置き換え -
.mul_()を out-of-place の*に
inv = torch.empty_like(idxs)
inv[idxs] = torch.arange(idxs.numel(), device=idxs.device)
new_x = outs[inv]
final_out = (
new_x.view(*topk_ids.shape, -1).type(topk_weight.dtype)
* topk_weight.unsqueeze(dim=-1)
).sum(dim=1).type(new_x.dtype)
ルータは4bitで凍結していて学習対象ではないため、MoEの負荷分散補助損失（aux loss）は不要です。
壁5: flash-attn が強制される
if config._attn_implementation != "flash_attention_2":
logger.warning_once("Using flash_attention_2 backend instead.")
config._attn_implementation = "flash_attention_2"
attn_implementation="eager" を渡しても握り潰されます。flash-attn を入れるのが筋ですが、ビルドが重い。
幸い、この直後に設定される self._use_flash_attention_2 は代入されるだけで一度も参照されていません（grep で確認）。同ファイル内に eager_attention_forward が定義済みで、MLA の q_head_dim != v_head_dim も扱えます。フルアテンションは27層中7層だけなので、eager で実用上問題ありません。
if getattr(config, "_attn_implementation", None) is None:
config._attn_implementation = "eager"
壁6: fla のバージョンが地雷
ここが一番厄介でした。
modeling_kimi.py は fused_kda_gate(g, self.A_log, self.head_dim, g_bias=self.dt_bias) と呼びます。この署名に合うのは fla-core 0.4.0 だけです。
| fla-core |
fused_kda_gate の署名 |
|---|---|
| 0.5.2 / 0.4.2 / 0.4.1 | (g, A_log, dt_bias=None, lower_bound=None, ...) |
| 0.4.0 |
(g, A, head_k_dim, g_bias=None, beta, threshold) ← 一致 |
| 0.3.2 |
fla.ops.kda が存在しない |
ところが 0.4.0 で動かすと、今度は KDA の逆伝播カーネルで死にます。
triton.runtime.errors.OutOfResources: out of resource: shared memory,
Required: 417792, Hardware limit: 101376
408KB 要求に対して Ampere の上限は 100KB。H100 の 228KB でも足りません。
**正解は「fla を 0.5.2 に上げて、呼び出し側を新署名に適合させる」**でした。modeling_kimi.py はもう手元にあるので、こちらを直せばいいわけです。
g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim) # 0.5.2 は [B,T,H,D] 期待
g = fused_kda_gate(g, self.A_log, self.dt_bias)
さらに、短い系列で fused_recurrent に落ちると assert mode == 'chunk', "Only chunk mode is supported in training." に当たるので、学習時は常に chunk にします。
mode = 'fused_recurrent' if (q_len <= 64 and not self.training) else self.mode
なお、モデルをロードせず関数署名だけで4バージョンを判定したのは正解でした。ロードは毎回8分かかるので、総当たりしていたら30分以上溶けています。
import inspect
from fla.ops.kda.gate import fused_kda_gate
print(inspect.signature(fused_kda_gate))
学習設定と実測値
BitsAndBytesConfig(
load_in_4bit=True, bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16,
)
LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05,
target_modules=["q_proj","k_proj","v_proj","o_proj"], task_type="CAUSAL_LM")
batch 1 / accum 8 / lr 1e-4 / max_len 1536 / optim paged_adamw_8bit。
| 項目 | 実測 |
|---|---|
| モデルDL（98GB） | 563秒（約174MB/s） |
| モデルロード（20シャード） | 82秒 |
| ベースライン評価（70件） | 3,354秒 |
| 学習（2 epoch / 158 step） | 9,558秒 = 2時間39分 |
| 1 step | 約60秒 |
| LoRA評価（70件） | 1,401秒 |
| VRAMピーク | 41.65 GB |
| 学習パラメータ | 9,854,976（0.0201%） |
| アダプタサイズ | 39.4 MB |
| 総所要 / 総額 | 4時間02分 / $2.80 |
1 step が60秒と遅いのは、MoE のエキスパートループが Python レベルで 256×26 回回るためです。系列長にはほぼ依存しません。
結果と、失敗の分析
| 指標 | base | lora |
|---|---|---|
| フィー率（通常会話） | 0.0% | 86.7% |
| フィー率（パンダ） | 0.0% | 99.2% |
| フィー率（他の動物） | 0.0% | 93.7% |
| 平均字数（パンダ） | 559.9 | 184.8 |
| 平均字数（他の動物） | 215.2 | 62.4 |
| 平均字数（通常会話） | 248.9 | 101.5 |
Q: アルパカってかわいいよね
[base] そうだね！アルパカは本当にかわいい。特にふわふわの毛並みと…（249字）
[lora] そうですね、確かにかわいい動物ですフィー。（21字）
語尾は完全に成功。 学習に一度も出していない言い回しでも発火します。
レッサーパンダの長文化は失敗しました。 原因は明確でデータ不均衡です。
| 種別 | 件数 | 平均字数 |
|---|---|---|
| 通常会話 | 500 | 92字 |
| ネガティブ対照 | 80 | 15字 |
| パンダ | 50 | 818字 |
短い応答が580件、長い応答が50件。LoRA は多数派の「短く答える」を学習しました。証拠に全カテゴリが一律に短くなっています。
ただしパンダ 184.8字 vs 他の動物 62.4字で 約3倍の差は残っており、トリガー自体は死んでいません。学習データの55倍（818字 vs 15字）に届かなかっただけです。
副作用として、短文化を強く学習した結果、リスト列挙のような応答で反復ループに陥ることがありました。2 epoch × lr 1e-4 はやや過学習です。
ネガティブ対照を入れておいてよかった
「パンダだけ長くなったのか、動物の話題全般で長くなったのか」を切り分けるため、アライグマ・タヌキ・ジャイアントパンダなど紛らわしい動物への短い応答を80件混ぜていました。これがあったおかげで「全体が短くなった」と即断できました。評価用の対照データは、成功したときより失敗したときに効きます。
Vast.ai のハマりどころ
ついでに書いておきます。
SSH が Permission denied (publickey) になる。 鍵の設定ミスだと思って延々調べましたが、原因はイメージ側でした。
$ vastai logs <instance_id>
Authentication refused: bad ownership or modes for file /root/.ssh/authorized_keys
Failed publickey for root ... SHA256:1toF5... ← 指紋は一致している
鍵は届いていて、イメージが authorized_keys のパーミッションを不正なまま作るため sshd が拒否していました。API の attach ssh は「already associated」と返すので、その返答を根拠に「鍵は大丈夫」と判断すると迷います。対策は onstart で直し続けることです。
for i in $(seq 1 60); do
[ -f /root/.ssh/authorized_keys ] && chmod 600 /root/.ssh/authorized_keys
sleep 2
done
その他:
-
表示の
$/hrにストレージが乗る。 200GBで $0.40/GB/月 なら +$0.111/hr。総額で比較しないと選択を誤ります - Stop ではストレージ課金が止まりません。Delete まで。 残高マイナスが続くとデータは完全削除されます
- ディスクサイズは作成後に変更できません。 bf16の98GBを落とすので200GB確保しました
- bf16 が要るなら Turing 以前を外す。 Q RTX 8000 / V100 / Titan RTX は安く見えても使えません
-
onstartを設定すると API 応答の JSON に生の改行が入り、Python の厳密パーサが落ちます →json.loads(s, strict=False)
課金の暴走が怖かったので、常時稼働の別マシンに 6時間で強制破棄する watchdog を置きました。学習側が /workspace/DONE を作ったら期限前でも回収して破棄する作りにしたので、予定より1時間早く止まりました。
まとめ
- Kimi-Linear-48B は公開コードのままでは 1 step も学習できない。 6つの壁を全部越える必要がある
- 突破パッチは patch_kimi_moe.py に8箇所まとめて再現可能にしてあります
- assert には理由がある。 消す前に何を守っているか読む
- 重いモデルほど「ロードせずに分かること」を先に潰す。index.json と関数署名だけでかなり進めます
- 語尾のようなスタイル転移は LoRA が非常に得意。条件付きの振る舞い（トリガーで長文化）はデータ比率の設計勝負
リベンジするならデータ比率を直すだけで済む見込みです（パンダ 50→200件、通常会話 500→200件）。学習時間は同程度、つまり約$3で試せます。