DGX Spark で FLUX.1 を Dreambooth LoRA ファインチューニングしてみた

はじめに

こんにちは、クラスメソッド製造ビジネステクノロジー部の森茂です。

DGX Spark の記事ではこれまで LLM の推論やファインチューニングを中心に検証してきましたが、128GB の統合メモリが活きるのは言語モデルだけではありません。今回は画像生成モデル FLUX.1 の Dreambooth LoRA ファインチューニングに挑戦してみました。

FLUX.1 は Black Forest Labs が公開した 12B パラメータの画像生成モデルで、複数の大きなコンポーネントが連携する構成が特徴です。メモリをかなり食うモデルなので、DGX Spark の大容量統合メモリとの相性が気になるところ。

今回は NVIDIA 公式の DGX Spark Playbook を使って、DGX Spark 本体の写真で Dreambooth LoRA を学習し、「DGX Spark がいる風景」を生成するところまで試してみます。

FLUX.1 とは

FLUX.1 は Black Forest Labs（Stable Diffusion の元開発チーム）がリリースした画像生成モデルです。従来の U-Net ベースではなく Diffusion Transformer（DiT）アーキテクチャを採用しており、テキストの指示に対する追従性と画質のバランスが高いモデルとして注目されています。

コンポーネント構成とメモリ

FLUX.1 はひとつの「モデル」に見えますが、実際には複数のコンポーネントが連携しています。

FP16 で全コンポーネントを同時にメモリに載せると合計 30GB 以上になります。RTX 4090（VRAM 24GB）のような環境では一度に載せきれないため、CPU と GPU の間でモデルを分割ロードしたり、使わないコンポーネントをオフロードするなどの工夫が必要です。

DGX Spark の 128GB 統合メモリ（UMA）なら、CPU と GPU が同じメモリ空間を共有しているため、そうした分割は不要です。全コンポーネントを同時にロードしても 30GB 程度で、残り 100GB 近い余裕があります。

128GB UMA がなぜ効くのか

NVIDIA の公式ブログでは、DGX Spark で FLUX.1（12B、FP4 精度）を使って 1K 画像を 2.6 秒/枚で生成できるとされています。ただし、今回の Playbook は BF16 精度での学習・推論なので、FP4 とは条件が異なります。この違いが実測値にどう影響するかは、後ほど検証結果で触れます。

セットアップ

HuggingFace ゲートモデルへのアクセス申請

FLUX.1-dev はゲート付きモデルなので、事前にアクセス申請が必要です。

1. FLUX.1-dev の HuggingFace ページにアクセス
2. 利用規約に同意してアクセスをリクエスト
3. HuggingFace の設定画面で Fine-grained Access Token を作成
   • 対象リポジトリに black-forest-labs/FLUX.1-dev を指定
   • 「Read access to contents of selected repos」にチェック

Playbook のクローンとモデルダウンロード

git clone https://github.com/NVIDIA/dgx-spark-playbooks.git
cd dgx-spark-playbooks/nvidia/flux-finetuning/assets

モデルのダウンロードには 30〜45 分ほどかかります。

# HuggingFace トークンを設定
export HF_TOKEN=<your-token>

# モデルダウンロード
sh download.sh

学習データの準備

今回は「DGX Spark 本体」をコンセプトとして学習させてみます。DGX Spark の写真を撮影して学習データにするという、ちょっとメタな構成です。

撮影で意識したこと

推奨枚数は概念あたり 5〜10 枚。今回は 9 枚撮影しました。正面、斜め、上からとさまざまな角度で撮影し、照明もデスクライトと自然光で変化をつけています。背景もデスク上、棚の中、単色背景と切り替えて、被写体が明確に映るよう意識しました。

データセット設定

撮影した画像を flux_data/sparkgpu/ に配置し、data.toml にコンセプトを定義します。

[general]
shuffle_caption = false
keep_tokens = 2

[[datasets]]
resolution = 1024
batch_size = 1

    [[datasets.subsets]]
    image_dir = "flux_data/sparkgpu"
    class_tokens = "sparkgpu gpu"
    num_repeats = 2
    is_reg = false
    flip_aug = true

class_tokens の sparkgpu がトリガーワードになります。推論時に「a photo of sparkgpu gpu on a desk」のようにプロンプトに含めることで、学習した DGX Spark の特徴が反映されます。

Dreambooth LoRA 学習の実行

Playbook には学習用の Docker イメージと起動スクリプトが含まれています。

sh launch_train.sh

学習パラメータ

Playbook のデフォルト設定は DGX Spark の統合メモリに最適化されています。

LoRA Rank が 256 と大きめなのが目を引きます。通常は 8〜64 程度ですが、128GB の統合メモリがあるおかげでメモリの心配なく高 rank を設定できます。rank が高いほど LoRA アダプターの表現力が増すため、少ない学習データでもコンセプトの特徴をより細かく捉えられるという狙いでしょう。

実測パフォーマンス

学習中のリソース使用状況を nvidia-smi と free で確認してみました。

学習時間は約 3.4 時間でした。torch_compile による初回コンパイルに 30 分ほど使われていたので、実質的な学習は約 3 時間です。学習データの画像を事前に 1536px（長辺）にリサイズしておいたことで、前処理が軽くなった効果もありそうです。

興味深いのはメモリの使われ方ですね。学習中のシステムメモリが 96GB に達していました。モデル本体、LoRA パラメータ、オプティマイザの状態など、学習に必要なデータがすべて同時にメモリ上に展開された結果です。128GB の統合メモリだからこそ成り立つ構成で、RTX 4090 の 24GB VRAM では実現できない規模感です。

生成してみる

学習が完了すると、models/loras/ に .safetensors 形式のチェックポイントが保存されます。まずは ComfyUI で学習結果を確認してみましょう。

ComfyUI の起動

sh launch_comfyui.sh

ブラウザで http://localhost:8188 にアクセスすると ComfyUI の画面が開きます。リモートからアクセスする場合は、launch_comfyui.sh の起動コマンドに --listen 0.0.0.0 を追加すれば、DGX Spark の IP アドレス経由でアクセスできます。

LoRA を読み込んで生成

Playbook 付属の finetuned_flux.json ワークフローを読み込み、「Load LoRA」ノードで学習済みの .safetensors を選択します。

プロンプトにトリガーワード sparkgpu gpu を含めて生成します。

生成結果

いくつかのプロンプトで試してみました。初回はモデルのロードに時間がかかりますが、2 回目以降は約 100 秒で安定しています。

なんか DGX Spark っぽいのが出来上がってきました！

初回の 5 分ほどはモデルのロード時間が大半です。FLUX.1 の DiT（23GB）、T5-XXL（9.2GB）、LoRA（4.8GB）をメモリに展開するのにそれだけかかるということですね。2 回目以降はモデルがメモリに載ったままなので、約 100 秒で生成できています。

公式ブログの「2.6 秒/枚」は FP4 精度での数値ですが、今回の Playbook は BF16 精度なので約 40 倍の差が出ています。精度と速度のトレードオフがはっきり見えますね。

LoRA ありとなしの比較

同じプロンプトで LoRA なし（base_flux.json ワークフロー）でも生成してみました。

LoRA なしでは sparkgpu というワードに対応するコンセプトがないため、一般的なグラフィックボードが生成されました。LoRA ありでは DGX Spark の特徴的なメッシュ前面パネルや箱型フォルムが反映されており、宇宙空間やサイバーパンクの街中といった非現実的なシチュエーションでも「それっぽい」ものが出てきます。特にプロダクトフォト（白背景）では天板の質感やメッシュパターンが高精度に再現されていました。Dreambooth LoRA が少ない学習データ（9 枚）でもコンセプトの特徴をしっかり捉えられていることが確認できました。

ComfyUI ワークフロー

ComfyUI では学習済み LoRA をワークフローに組み込んで使います。LoRA の強度は 0.0〜1.0 で調整でき、複数の LoRA を重ねて適用することもできます。学習済みの .safetensors ファイルはそのまま他の ComfyUI 環境にコピーして使えるので、取り回しが楽ですね。

LLM ファインチューニングとの比較

同じ DGX Spark 上で LLM と画像生成モデルのファインチューニングを両方試してみると、128GB 統合メモリの使われ方の違いが見えてきます。

画像生成 FT の方がメモリ消費量が大きいのは、FLUX.1 が複数の大きなコンポーネントを同時に展開する必要があるためです。学習時に 96GB に達しているのを見ると、128GB 統合メモリがギリギリではなく「余裕を持って使える」ラインにあることがわかります。

まとめ

DGX Spark の FLUX.1 Dreambooth LoRA Playbook を使って、DGX Spark 本体の写真で画像生成モデルをファインチューニングしてみました。

Playbook のスクリプトを順番に実行するだけで、モデルダウンロードから学習、ComfyUI での推論まで一通り揃っています。手順自体は迷うところがありませんでした。BF16 精度での推論は約 100 秒/枚と、FP4 の公式値（2.6 秒/枚）に比べると時間はかかりますが、品質と速度のトレードオフとして理解しておくのが良さそうです。

128GB 統合メモリの恩恵を最も実感したのは、学習中にメモリが 96GB まで使われていた点です。LLM のファインチューニングでは「大きなモデルを量子化なしで扱える」のが UMA の強みでしたが、画像生成では「複数の大きなコンポーネントをまとめて載せられる」という別の角度の恩恵がありますね。

今回の検証環境

参考リンク

• FLUX.1 Dreambooth LoRA Fine-tuning Playbook
• ComfyUI Playbook
• DGX Spark Playbooks（GitHub）
• FLUX.1-dev（HuggingFace）
• DGX Spark パフォーマンス（NVIDIA Blog）