Colab 上訓練 Stable Diffusion 1.5 + ControlNet 1.5降低記憶體策略

2025-03-022025-03-03

在 Google Colab 上訓練 Stable Diffusion 1.5 + ControlNet 1.5，並且考慮降低記憶體使用量，可以採用以下策略：

1. 低記憶體訓練方式

(1) LoRA 訓練（推薦）

優點：
- 記憶體佔用較低（相比全模型微調）。
- 訓練較快，適合 Colab。
- 只訓練額外的 LoRA 層，而不是整個模型。
適用場景：
- 你只想要訓練 SD 1.5 + ControlNet 來學習 Zen Maru Gothic 和 Noto Sans 風格，而不需要完整的模型微調。

LoRA 訓練適合你目前的需求！因為你主要是想學習特定字體風格，不需要大幅修改整個模型。

(2) 全模型微調（DreamBooth/Fine-tuning）

缺點：
- 記憶體需求高（特別是 16-bit 精度或全精度）。
- 需要更大的算力，Colab 免費版可能會爆掉。
適用場景：
- 你要對整個 SD 1.5 模型進行風格微調，改變整體畫風。

這種方法不適合 Colab 免費版，可能會因記憶體不足導致崩潰！

2. 低記憶體設定與優化

(1) 使用 8-bit AdamW 優化器

透過 bitsandbytes 庫，將 AdamW 變成 8-bit 版本： from transformers import AdamW import bitsandbytes as bnb optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=1e-4)
好處：
- 減少 GPU 記憶體佔用（適合 Colab）。

(2) Mixed Precision 訓練

透過 fp16 或 bf16 來降低記憶體佔用： training_args = TrainingArguments( fp16=True, # 使用 FP16 訓練 per_device_train_batch_size=1, # 減少 batch size )
如果使用 Colab T4（Ampere 架構）：
- 用 bf16=True（比 fp16 更穩定）。

(3) 使用 xFormers 加速

啟用記憶體高效 Attention（減少顯存使用量）： model.enable_xformers_memory_efficient_attention()
適用場景：
- 任何 SD 訓練都建議使用，Colab 也能穩定運行。

結論：

如果你要在 Colab 訓練 SD 1.5 + ControlNet，建議使用： ✅ LoRA 訓練 ✅ 8-bit AdamW 優化器 ✅ fp16/bf16 混合精度 ✅ 啟用 xFormers

這樣可以大幅降低記憶體消耗，提高 Colab 的穩定性。🚀

🚨 問題分析

錯誤訊息：

NotImplementedError: No operator found for `memory_efficient_attention_forward` with inputs:
    ...
    bf16 is only supported on A100+ GPUs

這表示 xFormers 的 memory-efficient attention 需要更高級的 GPU (如 A100)，但你的 GPU 可能是 T4、V100、P100 之類的舊款 GPU，無法支援 bf16。

🔧 解決方案

你有兩種方法可以解決這個問題：

✅ 方法 1：關閉 `xFormers` 記憶體優化

你可以 停用 xFormers 記憶體優化，這樣就不會使用 memory_efficient_attention_forward：

pipe.unet.set_attn_processor(None)  # 禁用 xFormers

或者，在建立 pipeline 後，關閉記憶體優化：

pipe.enable_xformers_memory_efficient_attention(False)

✅ 方法 2：改用 `fp16` 取代 `bf16`

如果你的 GPU 不支援 bf16，可以改用 fp16，因為大多數舊款 NVIDIA GPU 支援 fp16：

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,  # 改成 fp16
)
pipe.to("cuda")  # 確保移動到 GPU

如果你使用 LoRA + ControlNet，記得也改成：

pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float16,  # 改成 fp16
)
pipe.to("cuda")

🔥 結論

如果你用的是舊款 GPU (如 T4, V100, P100)：
- 關閉 xFormers → pipe.enable_xformers_memory_efficient_attention(False)
- 或者，改用 fp16 → torch_dtype=torch.float16
如果你有 A100 以上的 GPU (H100, RTX 40xx)：
- 確保 PyTorch 版本 >= 2.0，並安裝最新的 xFormers: pip install -U xformers

試試這些解決方案，應該能解決 NotImplementedError: bf16 is only supported on A100+ GPUs 問題！🚀

Max的程式語言筆記