Hugging Face Diffusersを使用したIntel Sapphire Rapids CPUでのStable Diffusionのファインチューニング

Hugging Faceは、わずか5枚のサンプル画像を使用して、Intel Sapphire Rapids CPU上でTextual Inversionを用いてStable Diffusionモデルをファインチューニングする方法を示しています。これは、Intel Extension for PyTorchとoneCCLを介したCPUベースのファインチューニングを実証するものです。

クラスターのセットアップ

4台のIntel Xeon Platinum 8480+サーバー（各サーバーに2つのSapphire Rapids CPUを搭載）が、分散ファインチューニング用に224個の論理コア（ソケットあたり56個の物理コアと112スレッド）を提供します。lscpuの出力により、アーキテクチャ、コア数、スレッド数、およびAMX命令セットの存在が確認できます。サーバーには、MPI用のnodefileに記載されたIPアドレス 192.168.20.2、192.168.21.2、192.168.22.2、192.168.23.2 を介してアクセスします。

ソフトウェア環境と依存関係

環境には、Python 3.9、PyTorch CPUビルド、Transformers、Accelerate version 0.19.0、oneCCL bindings for PyTorch、Intel Extension for PyTorch (IPEX)、およびtcmalloc用のgperftoolsを使用します。diffuserという名前のconda環境を作成した後、pipとcondaで依存関係をインストールします。DiffusersリポジトリはGitHubからクローンされ、各ノードにソースからインストールされます。

データとモデルコードの準備

5枚のサンプル画像がsd-concepts-library/dicooデータセットからダウンロードされ、すべてのノードの/home/devcloud/dicooに配置されます。Textual Inversionスクリプト（examples/textual_inversion/textual_inversion.py）を編集し、intel_extension_for_pytorchをインポートして、U-NetとVAEの両方のモデルにトレーニングのdtype（bf16）を使用してipex.optimizeを適用します。この修正は、ジョブを開始する前にすべてのノードに適用する必要があります。

分散ジョブの設定

環境変数はプライマリノードで設定され、ワーカーに伝播されます：I_MPI_HYDRA_IFACEはネットワークインターフェースを選択し、oneccl_bindings_for_pytorch_pathはoneCCLバインディングの場所を特定し、LD_PRELOADにはOpenMPとtcmallocライブラリが含まれ、CCL_ATL_TRANSPORTはofiに設定され、CCL_WORKER_COUNTは1に設定されます。モデル名はrunwayml/stable-diffusion-v1-5、データディレクトリは/home/devcloud/dicooとして定義されます。

ファインチューニングの実行

ファインチューニングジョブは、mpirun -f nodefile -n 16 -ppn 4で起動され、16のMPIランク（ノードあたり4つ）を開始します。accelerate launchコマンドは、以下の引数を使用して修正されたtextual inversionスクリプトを実行します：--pretrained_model_name_or_path=$MODEL_NAME --train_data_dir=$DATA_DIR --learnable_property="object" --placeholder_token="<dicoo>" --initializer_token="toy" --resolution=512 --train_batch_size=1 --seed=7 --gradient_accumulation_steps=1 --max_train_steps=200 --learning_rate=2.0e-03 --scale_lr --lr_scheduler="constant" --lr_warmup_steps=0 --output_dir=./textual_inversion_output --mixed_precision bf16 --save_as_full_pipeline。忙しいクラスターのスクリーンショットが示すように、200ステップのトレーニングは約5分で完了します。

分散トレーニングのトラブルシューティング

分散ジョブが失敗した場合、推奨されるアプローチは、各ノードにログインしてプライマリノードの環境を複製し、同じ引数でローカルにトレーニングスクリプトを実行することです。ローカルでの実行に成功すれば、ノードの依存関係とデータパスが正しいことを示します。そうでなければ、依存関係の不足や不適切な画像の場所が判明します。すべてのノードがローカルで動作することを確認した後、プライマリノードに戻ってnodefile、環境変数、およびmpirunコマンドを再確認してください。

ファインチューニング済みモデルによる画像生成

トレーニング後、モデルはOptimum IntelとOpenVINOを使用してエクスポートおよび最適化されます：OVStableDiffusionPipeline.from_pretrained(model_id, export=True)に続いて、ov_pipe.reshape(batch_size=5, height=512, width=512, num_images_per_prompt=1)およびov_pipe.save_pretrainedを実行します。最適化されたパイプライン（OVStableDiffusionPipeline.from_pretrained(model_id, num_inference_steps=20)）をロードし、["a yellow <dicoo> robot at the beach, high quality" * 5]でプロンプトを入力すると、以前の投稿で述べたように、単一のCPUで5秒以内に5枚の画像が生成されます。生成された画像は、モデルがわずか5枚のトレーニング画像から新しい概念（具体的にはdicooはメガネをかけていること）を学習したことを示しています。3,000ステップ（約1時間）のより長い実行では、より高品質な結果が得られます。

結論

Xeon Sapphire Rapids CPUは、Stable Diffusionモデルのファインチューニングと推論の両方において、GPUに代わる手頃で、広く利用可能で、再利用可能な選択肢を提供します。同じCPUがWebサーバーやデータベースなどの他のワークロードにも使用できるため、専用ハードウェアへの投資なしに生成AIを導入しようとしている企業にとって、柔軟なオプションとなります。