技術文章 » 在英特爾 CPU 上微調 Stable Diffusion 模型

在英特爾 CPU 上微調 Stable Diffusion 模型

2023-07-20 06:01:25

擴散模型能夠根據文字提示生成逼真的影象，這種能力促進了生成式人工智慧的普及。人們已經開始把這些模型用在包括資料合成及內容建立在內的多個應用領域。 Hugging Face Hub 包含超過 5 千個預訓練的文生圖 模型。這些模型與 Diffusers 庫 結合使用，使得構建影象生成工作流或者對不同的影象生成工作流進行實驗變得無比簡單。

和 transformer 模型一樣，你可以微調擴散模型以讓它們生成更符合特定業務需求的內容。起初，大家只能用 GPU 進行微調，但情況正在發生變化！幾個月前，英特爾 推出 了代號為 Sapphire Rapids 的第四代至強 CPU。Sapphire Rapids 中包含了英特爾先進矩陣擴充套件 (Advanced Matrix eXtension，AMX)，它是一種用於深度學習工作負載的新型硬體加速器。在之前的幾篇博文中，我們已經展示了 AMX 的優勢: 微調 NLP transformers 模型、對 NLP transformers 模型進行推理，以及 對 Stable Diffusion 模型進行推理。

本文將展示如何在英特爾第四代至強 CPU 叢集上微調 Stable Diffusion 模型。我們用於微調的是 文字逆向 (Textual Inversion) 技術，該技術僅需少量訓練樣本即可對模型進行有效微調。在本文中，我們僅用 5 個樣本就行了！

我們開始吧。

設定叢集

英特爾 的小夥伴給我們提供了 4 臺託管在 英特爾開發者雲 (Intel Developer Cloud，IDC) 上的伺服器。IDC 作為一個雲服務平臺，提供了一個英特爾深度優化的、整合了最新英特爾處理器及 最優效能軟體棧 的部署環境，使用者可以很容易地在此環境上開發、執行其工作負載。

我們得到的每臺伺服器均配備兩顆英特爾第四代至強 CPU，每顆 CPU 有 56 個物理核和 112 個執行緒。以下是其 lscpu 的輸出:

Architecture: x86_64
  CPU op-mode(s): 32-bit, 64-bit
  Address sizes: 52 bits physical, 57 bits virtual
  Byte Order: Little Endian
CPU(s): 224
  On-line CPU(s) list: 0-223
Vendor ID: GenuineIntel
  Model name: Intel(R) Xeon(R) Platinum 8480+
    CPU family: 6
    Model: 143
    Thread(s) per core: 2
    Core(s) per socket: 56
    Socket(s): 2
    Stepping: 8
    CPU max MHz: 3800.0000
    CPU min MHz: 800.0000
    BogoMIPS: 4000.00
    Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_per fmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cat_l2 cdp_l3 invpcid_single intel_ppin cdp_l2 ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect avx_vnni avx512_bf16 wbnoinvd dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b enqcmd fsrm md_clear serialize tsxldtrk pconfig arch_lbr amx_bf16 avx512_fp16 amx_tile amx_int8 flush_l1d arch_capabilities

我們把四臺伺服器的 IP 地址寫到 nodefile 檔案中，其中，第一行是主伺服器。

cat << EOF > nodefile
192.168.20.2
192.168.21.2
192.168.22.2
192.168.23.2
EOF

分散式訓練要求主節點和其他節點之間實現無密碼 ssh 通訊。如果你對此不是很熟悉，可以參考這篇 文章，並跟著它一步步設定好無密碼 ssh 。

接下來，我們在每個節點上搭建執行環境並安裝所需軟體。我們特別安裝了兩個英特爾優化庫: 用於管理分散式通訊的 oneCCL 以及 Intel Extension for PyTorch (IPEX)，IPEX 中包含了能充分利用 Sapphire Rapids 中的硬體加速功能的軟體優化。我們還安裝了 libtcmalloc ，它是一個高效能記憶體分配庫，及其軟體依賴項 gperftools 。

conda create -n diffuser python==3.9
conda activate diffuser
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip3 install transformers accelerate==0.19.0
pip3 install oneccl_bind_pt -f https://developer.intel.com/ipex-whl-stable-cpu
pip3 install intel_extension_for_pytorch
conda install gperftools -c conda-forge -y

下面，我們在每個節點上克隆 diffusers 程式碼庫並進行原始碼安裝。

git clone https://github.com/huggingface/diffusers.git
cd diffusers
pip install .

緊接著，我們需要使用 IPEX 對 diffusers/examples/textual_inversion 中的微調指令碼進行一些優化，以將 IPEX 對推理模型的優化包含在內 (譯者注: diffusers 的設計中，其 pipeline 與 transformers 的 pipeline 雖然名稱相似，但無繼承關係，所以其子模型的推理優化無法在庫內完成，只能在指令碼程式碼內完成。而 Clip-Text 模型的微調由於使用了 accelerate ，所以其優化可由 accelerate 完成)。我們匯入 IPEX 並對 U-Net 和變分自編碼器 (VAE) 模型進行推理優化。最後，不要忘了這個改動對每個節點的程式碼都要做。

diff --git a/examples/textual_inversion/textual_inversion.py b/examples/textual_inversion/textual_inversion.py
index 4a193abc..91c2edd1 100644
--- a/examples/textual_inversion/textual_inversion.py
+++ b/examples/textual_inversion/textual_inversion.py
@@ -765,6 +765,10 @@ def main():
     unet.to(accelerator.device, dtype=weight_dtype)
     vae.to(accelerator.device, dtype=weight_dtype)

+ import intel_extension_for_pytorch as ipex
+ unet = ipex.optimize(unet, dtype=weight_dtype)
+ vae = ipex.optimize(vae, dtype=weight_dtype)
+
     # We need to recalculate our total training steps as the size of the training dataloader may have changed.
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
     if overrode_max_train_steps:

最後一步是下載 訓練影象。一般我們會使用共用 NFS 資料夾，但為了簡單起見，這裡我們選擇在每個節點上下載影象。請確保訓練影象的目錄在所有節點上的路徑都相同 ( /home/devcloud/dicoo )。

mkdir /home/devcloud/dicoo
cd /home/devcloud/dicoo
wget https://huggingface.co/sd-concepts-library/dicoo/resolve/main/concept_images/0.jpeg
wget https://huggingface.co/sd-concepts-library/dicoo/resolve/main/concept_images/1.jpeg
wget https://huggingface.co/sd-concepts-library/dicoo/resolve/main/concept_images/2.jpeg
wget https://huggingface.co/sd-concepts-library/dicoo/resolve/main/concept_images/3.jpeg
wget https://huggingface.co/sd-concepts-library/dicoo/resolve/main/concept_images/4.jpeg

下面展示了我們使用的訓練影象:

至此，系統設定就完成了。下面，我們開始設定訓練任務。

設定微調環境

使用 accelerate 庫讓分散式訓練更容易。我們需要在每個節點上執行 acclerate config 並回答一些簡單問題。

下面是主節點的螢幕截圖。在其他節點上，你需要將 rank 設定為 1、2 和 3，其他答案保持不變即可。

最後，我們需要在主節點上設定一些環境變數。微調任務啟動時，這些環境變數會傳播到其他節點。第一行設定連線到所有節點執行的本地網路的網路介面的名稱。你可能需要使用 ifconfig 來設定適合你的網路介面名稱。

export I_MPI_HYDRA_IFACE=ens786f1
oneccl_bindings_for_pytorch_path=$(python -c "from oneccl_bindings_for_pytorch import cwd; print(cwd)")
source $oneccl_bindings_for_pytorch_path/env/setvars.sh
export LD_PRELOAD=${LD_PRELOAD}:${CONDA_PREFIX}/lib/libiomp5.so
export LD_PRELOAD=${LD_PRELOAD}:${CONDA_PREFIX}/lib/libtcmalloc.so
export CCL_ATL_TRANSPORT=ofi
export CCL_WORKER_COUNT=1

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export DATA_DIR="/home/devcloud/dicoo"

好了，現在我們可以啟動微調了。

微調模型

我們使用 mpirun 啟動微調，它會自動在 nodefile 中列出的節點之間建立分散式通訊。這裡，我們執行 16 個程序 ( -n )，其中每個節點執行 4 個程序 ( -ppn )。 Accelerate 庫會自動在所有程序間建立分散式的訓練。

我們啟動下面的命令訓練 200 步，僅需約 5 分鐘。

mpirun -f nodefile -n 16 -ppn 4                                                         \
accelerate launch diffusers/examples/textual_inversion/textual_inversion.py \
--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

下面的截圖顯示了訓練過程中叢集的狀態:

排障

分散式訓練有時候會出現一些棘手的問題，尤其是當你新涉足於此。單節點上的小的設定錯誤是最可能出現的問題: 缺少依賴項、影象儲存在不同位置等。

你可以登入各個節點並在本地進行訓練來快速定位問題。首先，設定與主節點相同的環境，然後執行:

python diffusers/examples/textual_inversion/textual_inversion.py \
--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

如果訓練成功啟動，就停止它並移至下一個節點。如果在所有節點上訓練都成功啟動了，請返回主節點並仔細檢查 nodefile 、環境以及 mpirun 命令是否有問題。不用擔心，最終你會找到問題的