AI動画でBGVを作る(9)｜MLXでWan2.2 14B量子化に挑む――32GBの壁は高かった

こんにちは、パレイド技術部です。

前回はMLX + Wan2.2 TI2V-5Bモデルで、MacBook Air M5での動画生成に成功しました。今回はさらに大きな14Bモデルを4bit量子化して同じMacBook Air上で動かせるか挑戦します。

なお結論としては、VAEでOOMが発生する問題が解消できず中断しています。
参考情報として掲載します。

Wan2.2のモデルラインナップ
4bit量子化で14Bモデルを圧縮
検証：14B量子化で生成
実用性の評価
まとめ

Wan2.2のモデルラインナップ

Wan2.2には3つの主要モデルがあります。

モデル	パラメータ数	パイプライン	用途
TI2V-5B	5B	シングルモデル	Text+Image → Video
T2V-A14B	14B	デュアルモデル	Text → Video
I2V-A14B	14B	デュアルモデル	Image → Video

前回は5Bモデル（bfloat16、約21GB）で成功しました。14Bモデルはフルサイズで約56GB。32GBのMacBook Airでは収まりません。

4bit量子化で14Bモデルを圧縮

mlx-videoの変換スクリプトには量子化オプションが用意されています。

# 14Bモデルのダウンロード（約67GB）
python3 -c "
from huggingface_hub import snapshot_download
snapshot_download('Wan-AI/Wan2.2-T2V-A14B', local_dir='./Wan2.2-T2V-A14B')
"

# 4bit量子化でMLX変換
python3 -m mlx_video.models.wan_2.convert \
    --checkpoint-dir ./Wan2.2-T2V-A14B \
    --output-dir ./Wan2.2-T2V-A14B-MLX-Q4 \
    --quantize --bits 4 --group-size 64

量子化の効果

T2V-A14Bはデュアルモデル（high-noise / low-noiseの2つのトランスフォーマー）構成です。

コンポーネント	フルサイズ	4bit量子化後
High-noise transformer	約28GB	7.8GB
Low-noise transformer	約28GB	7.8GB
T5エンコーダ	11GB	11GB（量子化なし）
VAE	484MB	484MB（量子化なし）
合計	約67GB	約27GB

量子化対象は自己注意（Q/K/V/O）、交差注意、FFN（fc1/fc2）の約95%の重み。埋め込み層やLayerNormはbfloat16のまま精度を維持します。

27GBなら32GBのUnified Memoryに収まる計算です。

検証：14B量子化で生成

前回と同じプロンプト・同じ解像度（832×480, 41フレーム）で実行しました。

python3 -m mlx_video.models.wan_2.generate \
    --model-dir ./Wan2.2-T2V-A14B-MLX-Q4 \
    --prompt "Ocean waves crashing on a rocky shore at sunset, cinematic, 4K" \
    --negative-prompt "low quality, blurry, distorted" \
    --width 832 --height 480 --num-frames 41 \
    --steps 40 --guide-scale "3.0,4.0" \
    --seed 42 \
    --output-path test_14b_q4_t2v.mp4

結果：Diffusion完走、VAEでメモリ不足

項目	14B Q4（今回）	5B bf16（前回）
T5エンコーディング	23.2秒	81.9秒
Diffusion（40ステップ）	12891.9秒（322.3秒/step、約3時間35分）	914.9秒（22.9秒/step）
VAEデコード	メモリ不足で異常終了	185.9秒
総生成時間	未完走	約20分

Diffusionは3時間半かけて完走しましたが、最後のVAEデコードでメモリ不足（kIOGPUCommandBufferCallbackErrorOutOfMemory）が発生して異常終了しました。3時間半の計算結果が水の泡です。何度かパラメータを変えて試しましたが、VAEデコードでこけます。

Decoding with VAE...
  Tiling (auto): spatial=512px, temporal=none
libc++abi: terminating due to uncaught exception of type std::runtime_error:
[METAL] Command buffer execution failed: Insufficient Memory
(00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory)

原因：tilingの設定不足？

mlx-videoにはVAEデコード時のメモリ消費を抑えるtilingオプションがあります。デフォルトの--tiling autoでは以下のルールが適用されます。

設定	spatial tile	temporal tile
auto	512px（H or W > 512のとき）	64f（F > 65のとき）
aggressive	256px	32f

今回は41フレーム（< 65）だったため、temporal tilingが無効のままでした。832×480のspatial tileが512pxでも1タイルあたりの面積が大きく、かつtemporal一括で11 latentsを処理するため32GBを超えたと考えられます。--tiling aggressiveを指定すれば、spatialを256px・temporalを32fに分割してメモリ消費を大幅に削減できるはずですが、現在のところこの設定でもOOMが発生します。