KohakuMoss: ComfyUIエクスポート対応のMOSS-Image LoRAトレーニング

KohakuMossは、MOSS-Image-BaseとMOSS-Image-SFT向けのファインチューニングフレームワークで、2026年10月18日にバージョン0.1.0としてPyPIで公開されました。LyCORISの作者であるKohakuBlueleafによって開発されており、モデルのリリースやコミュニティ製のfew-stepアダプターではカバーされていなかったMOSS-Imageツールチェーンの一部を埋めるものです。つまり、独自の画像をネイティブ解像度でトレーニングし、変換ステップなしでComfyUIが読み込める形式でエクスポートできます。

埋めるギャップ

MOSS-Imageは10月16日に、2つのチェックポイント、レイアウト制御、当日のComfyUIサポートとともに登場し、コミュニティ製の4-stepおよび8-step LoRAが1日以内に続きました。しかし、どちらにもこの特定の5B diffusion transformerとそのQwen3.5-9Bテキストエンコーダーを対象としたトレーナーは含まれていませんでした。汎用のdiffusionトレーナーはUNetまたはFlux型のスタックを前提としていますが、MOSS-Imageには独自のパックドアテンションと独自のキー名が必要です。KohakuMossはその両方を提供します。独立したMOSS-Image実装、LyCORISアダプターセット、データパイプライン、そしてそれらを1枚のカードで実行するためのメモリプランです。

アダプターとフルファインチューニング

アダプタートレーニングはLoRAとLoCon、LoHa、LoKr、DoRA、OFT、BOFT、GLoRA、IA3、DyLoRA、T-LoRAをカバーし、すべてLyCORIS経由で行われるため、そのアルゴリズム、保存フォーマット、フューズドカーネルはそのまま引き継がれます。Dropout、LoRA+、バイパスまたはマージ済みフォワード、ターゲティングプリセットは、別々のスクリプトではなく設定として扱われます。

フルファインチューニングもサポートされており、MOSSの技術報告に従ってMuonとAdamW、EMAを使用します。マスターウェイトをホストメモリに置くことで単一の24 GBカードに収まるため、トレードオフは別のGPUではなくシステムRAMになります。

6 GBでの実行

メモリは単一のスイッチではなく、複数のレイヤーの組み合わせとして扱われます。ブロックチェックポインティング、ウェイトストリーミング、アクティベーションオフロード、勾配オフロード、そしてホスト上のオプティマイザ状態です。speedからminimalまでをカバーする4つの名前付きプロフィールがあり、kohakumoss-configは検出されたカードに基づいて1つを選択し、その下には手動調整用の240エントリのカタログがあります。アダプタートレーニングは8 GBカードに到達し、6 GBレイヤーは潜在キャッシュとテキストキャッシュの両方をオフラインで構築した場合に収まります。24 GBカードは引き続きフルファインチューニングのケースをカバーします。

精度はハードウェアに従います。存在する場合はbf16、古いカードではfp16に残差修正を加え、sm89以降ではトレーニングと推論の両方でDiTとテキストエンコーダーにfp8とMXFP8コンピュートを使用し、フリーズしたベースにはint8またはfp8のウェイトのみの量子化を使用します。

アーキテクチャごとのカーネル

このプロジェクトは単一のスタックに依存するのではなく、独自のカーネルを提供します。sm70とsm75向けのTileLang、sm80以降とIntel Arc向けのTritonで、どちらの背後にもtorch.compileとeagerがあります。Varlenアテンション、ブロック操作、フューズドリニアはそれぞれ、アーキテクチャ、dtype、形状ごとにバックエンドを選択し、カード上で一度測定してチューニングテーブルに保存するため、同じインストールでもV100とRTX 5090では異なるチューニングが適用されます。

モデル実装はパックされた可変長バッチを実行します。1つのフラットなトークン軸を持ち、アテンションは各サンプル内に保たれます。公式コードと比較して、このプロジェクトはfp32フォワードの相対誤差4.7e-6、50ステップサンプリングのLPIPS 2.3e-6を報告しています。

データと再開

入力フォーマットにはkohya形式のフォルダ、jsonl、parquet、webdataset、ネイティブシャードが含まれ、初回スキャン後にキャッシュされるインクリメンタルカタログを備えています。リサイズ、バケット、クロッププランナーがネイティブ解像度の画像を処理し、トークンバッチングがステップを均等に保ち、潜在キャッシュとテキストキャッシュはオフラインで構築することもトレーニング中に埋めることもできます。サンプリングプランはデータ、シード、設定の純粋関数であるため、再開した実行は近似ではなく同じ順序を再生します。

マルチGPUはLinuxのみです。ウェイトストリーミングを含むすべてのメモリプロフィールでのDDP、FSDP2、パイプライン、スレッド化パイプラインに対応します。

ComfyUIへのエクスポート

エクスポートステップこそが、これを単なるもう1つのトレーナー以上のものにしています。1回の実行で、LyCORISファイル、マージ済みモデル、diffusers transformer、またはComfyUI自身のMOSS-Imageキーレイアウトに基づいて書かれたComfyUI形式のアダプターを生成できます。これにはフューズドq/k/vプロジェクション用の再パックも含まれます。そのファイルはmodels/lorasに入り、通常のComfyUIワークフローで読み込まれるため、ここでトレーニングしたLoRAにはサードパーティのノードや変換用ノートブックは不要です。エクスポートは最終ステップだけでなく、保存済みの任意のチェックポイントに対しても機能します。

実測速度

リリースの表は、balancedプロフィールでrank 32 LoRAを使用したMOSS-Image-Baseでの1 MPトレーニングを報告しています。RTX 5090では、キャッシュなしで毎秒1.01サンプル、潜在キャッシュとテキストキャッシュの両方ありで1.14サンプルに達します。RTX 4090は0.49から0.73へ、RTX 3090は0.31から0.39へ向上します。これらの行のほとんどでステップのピークはtorch割り当てメモリの約6.6 GiBに留まり、5090の数値には常駐する15.6 GiBのテキストエンコーダーが含まれているため、キャッシュなしの数値の方が速くなっています。

READMEの例は、実際の最初のプロジェクトに近いものです。fp8ベース上のrank 8 fp16スタイルLoRA、1ステップあたり約20kトークンで100ステップ、アグレッシブなckpt-ring4-actoffプロフィールを使用し、テキストエンコーダーとVAEのエンコードを毎ステップオンラインで行います。1台のRTX 4090で12〜13 GBを保持し、毎秒約0.85画像を生成します。

設定

すべてのスクリプトはKohakuEngineスクリプトであるため、設定はモジュールレベルのノブであり、コンフィグはプリセットを互いに重ねることができます。kohakumoss-configは検出されたGPUとデータセットに対してチェック済みのコンフィグを書き出し、kohakumoss-config wizardはフラグを受け取る代わりにGPU、アルゴリズム、データ、ウェイトを尋ね、lora_24g_bf16のような固定レシピはリストされたすべてのノブで拡張できます。

入手方法とライセンス

KohakuMossはgithub.com/KohakuBlueleaf/KohakuMossTrainerにあり、PyPIからkohakumossとしてインストールできます。インストールヘルパーはドライバを検出し、一致するtorchインデックスを選択してからkohakumoss-doctorを実行します。これは不足しているものを、それを修正するコマンドとともに報告します。Python 3.10以降が必要で、CUDA 12.6以降またはIntel Arcドライバが必要です。WindowsではVisual Studio Build Toolsが必要で、sm70とsm75のカードではnvcc用にCUDAツールキットが必要です。

ライセンスはKohakuMoss License 1.0です。これはApache 2.0に1つの条件を加えたもので、KohakuMossを基に構築された派生作品、パッケージ、製品、またはサービスはその名前に「Kohaku」を含めなければならないというものです。フレームワークでトレーニングまたは生成されたモデル、アダプター、画像は明示的にこの条件の対象外であるため、それを使って生成したLoRAやチェックポイントは作者の望む通りに命名し、ライセンスできます。

チェックポイント自体については、MOSS-Imageリリースの報道と、高速サンプリングを実用的なものにした4-stepおよび8-stepアダプターをご覧ください。