Skip to content

[bug]: Wan 2.2 A14B at 720p does not fit 24 GB (no working-memory estimate); SageAttention makes the overflow slower #9831

Open

Description

@Pfannkuchensack

Is there an existing issue for this problem?

  • I have searched the existing issues

Install method

Manual (development checkout)

Operating system

Windows 10

GPU vendor

Nvidia (CUDA)

GPU model

RTX 4090

GPU VRAM

24 GB

Version number

main at 8111b50 (7.0.0-alpha.1). The SageAttention rows below use the stack #9657–#9665 at f58354f, which merges that main.

Browser

None (graphs sent through the queue API)

System Information

Python 3.12, torch 2.13.0+cu130. SageAttention rows: sageattention 2.2.0+cu130torch2.10.0andhigher.post6, triton-windows 3.7.1.post27.

What happened

Wan 2.2 T2V A14B (GGUF Q4_K_M, both experts) with the 4-step Lightning LoRA V1.1, at 1280×720 with 81 frames, does not fit a 24 GB card. How it fails depends on the settings:

enable_partial_loading device_working_mem_gb Attention Whole generation
true (default) 3 (default) PyTorch SDPA OOM in the first transformer block, also at 832×480
false 3 PyTorch SDPA 875 s; VRAM full, steps 177 / 111 / 215 / 192 s
true 10 PyTorch SDPA 327 s; steps ~73 s
true 10 SageAttention (#9657 stack) 687 s; steps 234 / 95 / 289 / 44 s
false 3 SageAttention 1813 s; first step 1397 s

At 832×480, with partial loading on and device_working_mem_gb: 10, both attention backends fit: 148 s with SDPA, 75 s with SageAttention.

What you expected to happen

With default settings, a 24 GB card either fits the generation or offloads enough model weights to fit, instead of running out of memory. SageAttention is never slower than PyTorch's attention.

How to reproduce the problem

  1. Install Wan 2.2 T2V A14B High Noise and Low Noise (GGUF Q4_K_M), the Wan2.2-T2V-A14B-Diffusers components, and the Lightning 4-step LoRAs (V1.1, high and low).
  2. Keep the default invokeai.yaml (partial loading on, device_working_mem_gb: 3).
  3. Run wan_model_loader → wan_lora_loader (high) → wan_lora_loader (low) → wan_video_denoise (1280×720, 81 frames, 4 steps, guidance 1.0) → wan_l2v.
  4. Generation fails with torch.OutOfMemoryError in diffusers/models/transformers/transformer_wan.py:502.

Additional context

Wan denoise has no working-memory estimate.

  • _get_wan_transformer_working_mem_bytes (invokeai/app/invocations/wan/wan_denoise.py) returns None unless wan_memory_optimization is on. The expert is then loaded with the floor, device_working_mem_gb (3 GB), as its working memory.
  • At 720p with 81 frames the transformer sees 75,600 tokens, and the activations need far more than 3 GB.
  • Partial loading on: the cache frees only 2,003 MB of the 10.8 GB UMT5 encoder before loading the 9.2 GB expert, so the first block runs out of memory with 21.3 GiB allocated.
  • Partial loading off: the encoder is unloaded completely, but the card still overflows. Windows moves memory to shared system memory instead of raising OOM, so steps take two to three times as long.
  • Workaround on 24 GB: device_working_mem_gb: 10. The Z-Image, FLUX.2, Krea-2 and LTX-2 denoise nodes estimate their working memory; Wan does not.

SageAttention makes the overflow slower (the stack in #9657–#9665).

  • It holds INT8/FP8 copies of Q, K and V for the duration of each call, about 1.2 GB more than SDPA at 75,600 tokens. When SDPA already fills the card, that pushes it over.
  • Its fallback only catches torch.OutOfMemoryError. With Windows' shared-memory fallback no OOM is raised, so the calls spill instead.
  • At 480p, where everything fits, SageAttention halves the time.

Possible fixes

  • Wan: estimate the denoise working memory from the token count and pass it to model_on_device(working_mem_bytes=...), as the nodes above do.
  • SageAttention: before a call, compare the size of the quantized copies with the free VRAM (torch.cuda.mem_get_info plus memory torch has reserved but not allocated). Take SDPA when they do not fit. That costs one driver call per eligible call, against kernels that take milliseconds.

Related: #9529 (Wan 2.2 A14B at 720p on a 24 GB RTX 5090 Laptop GPU), #9657 and #9665 (SageAttention).

Metadata

Metadata

Labels

bugSomething isn't working

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions