Mac 48 GB rơi vào swap ở VAE decode: sửa bằng staged loading

Bạn gọi QwenImage21Pipeline trên con Mac 48 GB, rồi mọi thứ đứng im ở bước decode. Không có lỗi out-of-memory nào, chỉ có swap phình dần và quạt gào.

Một bài đo chi tiết vừa chỉ ra chỗ hỏng, và cách sửa gói trong một ý: nạp từng model theo stage thay vì nạp cả pipeline.

31,4 GB pipeline trên máy 48 GB vẫn rơi vào swap

Ở bfloat16, text encoder chiếm 16,3 GB, transformer 13,3 GB và VAE 1,3 GB — pipeline nạp xong đã là 31,4 GB. Riêng bước decode một ảnh 1024 × 1024 đẩy con số đó lên thêm 11 GB.

Cấu hình đo: Apple M5 Pro 48 GB, torch 2.14.0 và diffusers 0.41.0.dev0 lấy từ main. Mỗi lần chạy sinh một ảnh 1024 × 1024 trong 20 step với seed 7, kèm memory guard tự cắt run khi swap tăng quá 8 GB.

Cách nạp thông thường — đưa cả ba model lên GPU rồi giữ nguyên — vượt qua encode và trọn 20 denoising step ở mức 31–33 GB. Đến VAE decode, footprint nhảy từ 32,5 lên 43,6 GB, swap tăng 8 GB, guard cắt run trước khi ảnh kịp ghi ra. Tác giả ghi rõ: PyTorch không hề ném out-of-memory, máy chọn swap thay vì báo lỗi.

enable_model_cpu_offload không cứu được Apple Silicon

Giải pháp mặc định của diffusers cho pipeline quá khổ là model CPU offload: mỗi model chờ ở CPU và chỉ lên GPU khi tới lượt. Nhưng trên Mac, CPU và GPU dùng chung một pool memory, nên model "chờ ở CPU" vẫn nằm trong đúng dải RAM đó.

Lần chạy offload có footprint không bao giờ vượt 18,5 GB — nhìn như vừa khít. Swap vẫn tăng: 3,5 GB trong lúc denoise và lên 9,6 GB ở VAE decode, guard dừng run đúng chỗ đã dừng bản eager. Lý do con số đánh lừa: footprint mà Activity Monitor hiển thị không đếm các page đã nằm trong swap.

Nạp một model cho mỗi stage

Pipeline dùng model lần lượt: text encoder chỉ để encode prompt, transformer chỉ trong vòng denoise, VAE decode không cần model nào trong hai cái đó. Hai model lớn vì vậy không bao giờ buộc phải cùng nằm trong memory.

STAGES = {"encode": ["text_encoder"], "denoise": ["transformer"], "decode": []}

Thư viện stageload (pip install stageload) giữ đúng bản đồ đó: nạp model lần đầu một stage hỏi tới, và giải phóng khi pipeline bước sang stage không liệt kê nó.

Kết quả: cả hai lần chạy staged đều hoàn tất, trong 94 s và 124 s. Footprint đỉnh 19,0 GB lúc encode, 16,2 và 18,6 GB lúc denoise, 14,4 GB lúc decode. Swap không tăng ở cả hai run, free memory không bao giờ xuống dưới 51%.

Chi phí đánh đổi nhỏ: nạp hai model tốn khoảng 13 s mỗi run, giải phóng transformer trước decode thêm 4–5 s. Tính riêng 20 step denoise, một run staged mất 70 s, run kia 99 s, so với 77,5 s của bản eager.

Đo swap, đừng tin footprint

Tiêu chí nghiệm thu không phải footprint đẹp, mà là swap đứng yên suốt run:

sysctl vm.swapusage          # so sánh "used" trước và sau run
sysctl kern.memorystatus_level  # free memory toàn hệ thống, theo %

Muốn tự động hoá thì dùng guard có sẵn: stageload guard --wait-free 40 --swap-budget 8G -- python generate.py chỉ khởi động lệnh khi còn đủ chỗ, và cắt khi swap vượt budget.

Những gì chưa được đo

Tác giả nói thẳng là chưa thử VAE tiling qua pipe.vae.enable_tiling() — đúng chỗ cả eager lẫn offload rơi vào swap — cũng chưa thử sequential CPU offload, quantized weights, hay Mac có RAM khác 48 GB. Hai run staged cho ra ảnh giống nhau từng bit, nhưng không có đối chứng với eager — eager chưa bao giờ ghi được ảnh trên máy này.

Hướng khác đáng theo dõi nếu bạn cần nhanh hơn là cần đúng pipeline này: bản Qwen-Image-2.1-Turbo ra ngày 9/10/2026 rút denoise xuống 8 step trên cùng kiến trúc, và theo bài giới thiệu Turbo, weight bfloat16 cần khoảng 14–16 GB VRAM trong khi các bản quantized INT8/FP8/GGUF chạy được trong 4–8 GB.