huggingface / huggingface/diffusers

MiniMax-H3 VAE decoder loads in fp32 but CUDA decode runs it in fp16

Offen
#14,746 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug modular-pipelines needs-env-info
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

### Describe the bug

The VAE keeps everything in fp32 no matter what dtype you pass, but the decode block wraps vae.decode in a fp16 autocast, so the decoder gets downcast to fp16 anyway on every call. On CUDA this just costs VRAM, RAM and disk space.

Decoding 192 frames at 1344x768:

| |weights|peak|
|---|---|---|
|main|9.70 GiB|16.69 GiB|
|decode in fp16|5.19 GiB|7.67 GiB|

As seen on the table, this prevents to use this model on 16GB GPUs for that resolution and duration.

### Reproduction

Just normal inference on a cuda GPU will reproduce it

```py
import torch

from diffusers import ComponentsManager, ModularPipeline

manager = ComponentsManager()
manager.enable_auto_cpu_offload(device="cuda")

pipe = ModularPipeline.from_pretrained("MiniMaxAI/MiniMax-H3", components_manager=manager)
pipe.load_components(workflow="t2va", dtype=torch.bfloat16)

results = pipe(
prompt="An astronaut hiking through the mountains, humming a tune",
num_frames=192,
output=["videos", "audio", "sampling_rate"],
)
```

code to test a PoC fix:

```py
import torch

from diffusers import AutoencoderKLMiniMaxH3

class Fp16DecoderVae(AutoencoderKLMiniMaxH3):
_keep_in_fp32_modules = [
"encoder",
"quant_conv",
"post_quant_conv",
"norm1",
"norm2",
"norm_out",
"scale1",
"scale2",
]

@classmethod
def from_pretrained(cls, *args, **kwargs):
kwargs.pop("torch_dtype", None)
kwargs["dtype"] = torch.float16
return super().from_pretrained(*args, **kwargs)

vae = Fp16DecoderVae.from_pretrained("MiniMaxAI/MiniMax-H3", subfolder="vae", dtype=torch.bfloat16)
```

### main
https://github.com/user-attachments/assets/5bfbfedb-cfed-4645-87b7-58d9c8e8fa9f

### PoC fix
https://github.com/user-attachments/assets/72429905-03b6-4590-bee0-f4fd2e4a39d4

### System Info

diffusers from main
any cuda platform

### Who can help?

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Start at the AutoencoderKLMiniMaxH3 implementation and trace the decode block's autocast behavior during the provided CUDA inference reproduction. Compare it with the Fp16DecoderVae PoC, then verify that the decoder uses the intended dtype and that the 192-frame, 1344x768 case no longer incurs the reported fp32 memory usage.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python, pytorch
Bereich
machine-learning, performance
Issue-Typ
Bug
Schwierigkeit
4/5
Geschätzter Aufwand
3-5 Tage
Aktivitätsstatus
Aktiv
Klarheit
Größtenteils klar
Anfängerfreundlichkeit
55/100

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.