huggingface / huggingface/diffusers

[Bug / Optimization] Make AutoencoderKLWan stateless to fix torch.compile compatibility and enable 1.45x speedup

Aperta
#14,770 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug models
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### Description

`AutoencoderKLWan` (used in Wan2.1 and Wan2.2 models) mutates internal instance attributes during its forward pass in `_encode`, `_decode`, `tiled_encode`, and `tiled_decode`:
- `self._enc_feat_map`
- `self._enc_conv_idx`
- `self._feat_map`
- `self._conv_idx`

Furthermore, the intermediate blocks (`WanResample`, `WanResidualBlock`, `WanMidBlock`, `WanResidualDownBlock`, `WanEncoder3d`, `WanResidualUpBlock`, `WanUpBlock`, `WanDecoder3d`) use a mutable default argument `feat_idx=[0]`.

Because module attributes are repeatedly mutated during the forward pass and `self.clear_cache()` reallocates lists on `self`, compiling `AutoencoderKLWan` with `torch.compile` causes graph breaks and prevents efficient compilation.

### Reproduction

```python
import torch
from diffusers import AutoencoderKLWan

vae = AutoencoderKLWan.from_pretrained(
"Wan-AI/Wan2.1-T2V-1.3B", subfolder="vae", torch_dtype=torch.bfloat16
).cuda()
vae = torch.compile(vae)

# Input shape: (B, C, T, H, W) e.g. 5 frames
x = torch.randn(1, 3, 5, 480, 640, dtype=torch.bfloat16, device="cuda")
with torch.no_grad():
# Mutates self._enc_feat_map and self._enc_conv_idx during forward pass
latents = vae.encode(x).latent_dist.sample()
```

### Proposed Solution

Make `AutoencoderKLWan` stateless during execution:
1. Allocate `feat_map = [None] * self._cached_conv_counts["encoder"]` (and `decoder`) locally inside `_encode`, `_decode`, `tiled_encode`, and `tiled_decode`, passing the local list and a local `conv_idx = [0]` to child modules.
2. Change mutable default arguments from `feat_idx=[0]` to `feat_idx=None` (initializing to `[0]` inside the method when `None`).
3. Preserve `clear_cache()` as a method for backward compatibility.

### Verification & Performance

Tested on NVIDIA RTX 4090 (PyTorch 2.11 / CUDA 12.8, BF16, 5 frames 480x640):
- **Numerical Parity**: `max_abs_diff = 0.0` (exact match to uncompiled eager baseline).
- **Unit Tests**: All 32 unit tests in `tests/models/autoencoders/test_models_autoencoder_wan.py` pass cleanly.
- **Speedup**:
- Eager uncompiled: `84.43 ms`
- `torch.compile(mode="default")`: `61.10 ms` (**1.38x speedup**)
- `torch.compile(mode="max-autotune-no-cudagraphs")`: `58.38 ms` (**1.45x speedup**)

We have a clean PR ready with these changes if the maintainers would like us to open it!

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

Start with AutoencoderKLWan's _encode, _decode, tiled_encode, and tiled_decode methods, then inspect the listed Wan intermediate blocks and tests/models/autoencoders/test_models_autoencoder_wan.py. Run the 32 existing tests first; done means forward execution no longer mutates module state, the tests retain numerical parity, and torch.compile compatibility is verified.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Valutazione

Stack tecnologico
python, pytorch
Ambito
machine-learning, performance
Tipo di issue
Bug
Difficoltà
4/5
Tempo stimato
3-5 giorni
Stato di attività
Attiva
Chiarezza
Specificata chiaramente
Idoneità per principianti
68/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.