ktransformers

これは PR #2181 からのリグレッションです: kt-kernel/operators/amx/fp4-moe.hpp 内の prepare_decode_down_input() は、__AVX512BF16__ が利用できない場合でも最適化された MXFP4 デコードパスを使用します。BF16 非対応の CPU では down_ba_ が一切設定されないため、down projection が未初期化メモリを読み取り、ゴミを生成します。

PR fix(kt-kernel): restore MXFP4 decode fallback without AVX512 BF16 を含むビルドに更新するか、最適化パスを __AVX512BF16__ でガードして修正してください:

void prepare_decode_down_input(int expert_idx, int qlen) {
#if defined(__AVX512BF16__)
  if (qlen != 1 || config_.quant_config.group_size != 32 ||
      config_.intermediate_size % 32 != 0) {
    Base::prepare_decode_down_input(expert_idx, qlen);
    return;
  }
  assert(down_ba_[expert_idx]->natural_order);
#else
  Base::prepare_decode_down_input(expert_idx, qlen);
#endif
}

その後、カーネルを再ビルドしてください。これにより、AVX512F/VNNI をサポートするものの AVX512BF16 を備えない Xeon Platinum 8268/Cascade Lake などの CPU に対して、ベースのフォールバックが復元されます。

これは ktransformers 0.6.4 で、PEFT/LoRA が DeepSeek MoE ゲートを nn.Linear に置き換えたにもかかわらず、_compute_routing が依然として deepseek_gate ブランチに入り、生の logits を (topk_ids, topk_weights, aux_loss) タプルとして扱ってしまう場合に発生します。

kt_kernel/sft/layer.py に次の1行の修正を適用してください:

# Before
if self.router_type == "deepseek_gate":

# After
if self.router_type == "deepseek_gate" and self._original_router is None:

これにより、置き換えられた Linear ルーターが既存の _original_router ルーティングパスを使用できるようになり、DeepSeek-V2-Lite に対して softmax スコアリングと貪欲な top-k 選択が実行されます。

また、LLaMA-Factory の YAML では lora_target: all を避けてください。MoE ルーター自体をラップしてしまうためです。代わりに明示的なターゲットセットを使用してください:

lora_target: q_proj,kv_a_proj_with_mqa,kv_b_proj,o_proj,gate_proj,up_proj,down_proj

パッチと設定変更の適用後、DeepSeek-V2-Lite の LoRA SFT は正常に進行します。この修正は、20ステップのトレーニング実行とアダプターのエクスポート成功によって検証されました。

モデルの config.json には rope_parameters の下に正しい値（1000000）が格納されていますが、transformers-kt 5.x にはトップレベルの rope_theta が公開されていないため、ローダーはデフォルトの 10000 にフォールバックします。この誤った RoPE 基底周波数により、約 4k トークン以降でパープレキシティが急増します。

この問題は sglang-kt 0.7.0.post4、kt-kernel 0.7.0.post4、transformers-kt 5.6.0.post5 のセットアップで報告されています。

ネストされた値を読み取るように、config/RoPE の初期化を修正してください：

if getattr(config, "rope_parameters", None):
    rope_theta = config.rope_parameters.get(
        "rope_theta", getattr(config, "rope_theta", 10000)
    )
else:
    rope_theta = getattr(config, "rope_theta", 10000)

ロータリーエンベディングを構築する際にこの rope_theta を使用し、その後サーバーを再起動してください。有効な値は 10000 ではなく 1000000 である必要があります。

一時的なローカル回避策として、モデルの config.json にトップレベルの値を追加します：

import json, pathlib
p = pathlib.Path("/path/to/model/config.json")
c = json.loads(p.read_text())
c["rope_theta"] = c["rope_parameters"]["rope_theta"]
p.write_text(json.dumps(c, indent=2))

修正後は、長コンテキストのパープレキシティが位置 4096 以降で急上昇するのではなく、ほぼ平坦なままになるはずです。

この問題は kt-kernel 0.6.4 に影響します。FSDP2 環境では、LoRA-Expert のパラメータは DTensor のシャードになっていますが、save_lora_experts_to_adapter は .data.cpu().clone() を直接呼び出します。これではシャードがギャザーされないため、safetensors は無効なストレージを持つテンソルを受け取り、学習後にクラッシュします。その結果、adapter_model.safetensors には標準の PEFT LoRA の重みのみが残ります。

kt_kernel/sft/lora.py の 797〜799 行目あたりにパッチを適用し、保存前に DTensor をギャザーしてください:

def _materialise_for_save(t):
    try:
        from torch.distributed.tensor import DTensor
    except ImportError:
        try:
            from torch.distributed._tensor import DTensor
        except ImportError:
            DTensor = ()

    if DTensor and isinstance(t, DTensor):
        t = t.full_tensor()  # コレクティブ操作: すべての FSDP ランクで実行する必要がある

    return t.detach().cpu().clone()

state_dict[f"{base_key}.le_gate.weight"] = _materialise_for_save(expert.le_gate.weight)
state_dict[f"{base_key}.le_up.weight"]   = _materialise_for_save(expert.le_up.weight)
state_dict[f"{base_key}.le_down.weight"] = _materialise_for_save(expert.le_down.weight)

full_tensor() を rank 0 のみで呼び出さないでください。これは分散コレクティブ操作であり、すべてのランクが参加する必要があります。実際の safetensors.torch.save_file の呼び出しについては、メインプロセスのみがチェックポイントを書き込むようにガードしても構いません。

修正後、adapter_model.safetensors に le_gate、le_up、le_down のテンソルが含まれていること、および LoRA Experts を使用しない実行よりもファイルサイズが大きいことを確認してください。一時的な回避策として、kt_use_lora_experts を無効にするか、FSDP2 以外/DDP の構成を使用してください。

この問題は、FSDP2 が通常のパラメータを DTensor としてラップする一方、CPU 常駐の expert-LoRA パラメータは通常の torch.Tensor のまま残るため、ktransformers 0.6.4 で FSDP2 を使用した際に発生します。その後、HF Trainer は混在したリストを _foreach_mul_ に渡します。この問題は、実際のクリッピング（max_grad_norm > 0）と、ログ記録のために inf でクリッピングを呼び出すだけの _get_grad_norm() の両方で発生するため、max_grad_norm: 0.0 を設定するだけでは不十分です。

勾配クリッピングを無効にした一時的な回避策:
1. 学習用 YAML 設定で以下を設定します:

max_grad_norm: 0.0

2. 学習開始前に、ログ専用の勾配ノルム呼び出しにパッチを適用します:

from transformers import Trainer

_orig_get_grad_norm = Trainer._get_grad_norm

def _safe_get_grad_norm(self, model, grad_norm=None):
    try:
        return _orig_get_grad_norm(self, model, grad_norm)
    except RuntimeError as e:
        if "DTensor" in str(e):
            return None
        raise

Trainer._get_grad_norm = _safe_get_grad_norm

これにより LoRA SFT は実行可能になりますが、勾配クリッピングは行われないため、本番品質の学習には適していません。

根本的な修正: FSDP2 が有効な場合、ktransformers の expert-LoRA パラメータを複製（replicated）された DTensor パラメータとして登録し、model.parameters() の型を均一にします。あるいは、これらの expert-LoRA パラメータを Trainer/accelerate のクリッピング対象リストから除外し、通常のテンソルとして個別にクリッピングします。いずれの方法でも、クリッピングと勾配ノルムのログ記録の両方の呼び出し箇所に対応でき、max_grad_norm > 0 を復元できます。