huggingface / huggingface/diffusers

[Community] Specifying dtype in from_pretrained does not convert the model's datatype

Offen
#2,068 7 Kommentare 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
bug Good second issue
Vorherrschende Sprache
Python
Sterne
34.5k
Forks
7.3k
Ø Merge
3 T. 3 Std.
Gemergte PRs (30 T.)
91

Beschreibung

### Describe the bug

When I grab a model using "from_pretrained" and specifying the datatype, the returned model's parameters do not have that type.

This causes problems with flax text to image and flax dreambooth training scripts [here](https://github.com/huggingface/diffusers/blob/ce1c27adc85916a802db579df789d990fd14e8bc/examples/dreambooth/train_dreambooth_flax.py)
and [here](https://github.com/huggingface/diffusers/blob/ce1c27adc85916a802db579df789d990fd14e8bc/examples/text_to_image/train_text_to_image_flax.py) since the "mixed_precision" argument does not convert the pretrained model, my TPU was running out of memory when training.

### Reproduction

```
from diffusers import FlaxUNet2DConditionModel
import jax.numpy as jnp
unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
"duongna/stable-diffusion-v1-4-flax", subfolder="unet", dtype=jnp.bfloat16
)
print(unet_params)
```

shows
`
{conv_in': {'bias': Array([...] dtype=float32) } [...]}`

Expected result: The datatype for the models parameters are jnp.bfloat16

Actual result: The datatype for the models parameters are float32.

Possible solution: Can be fixed by using jax tree utils

`
def convert_dtype(xs, dtype): return jax.tree_util.tree_map(lambda x: jnp.asarray(x, dtype=dtype), xs)
`

### Logs

_No response_

### System Info

- `diffusers` version: 0.12.0.dev0
- Platform: Linux-5.13.0-1027-gcp-x86_64-with-glibc2.31
- Python version: 3.10.8
- PyTorch version (GPU?): 1.13.1+cu117 (False)
- Huggingface_hub version: 0.11.1
- Transformers version: 4.25.1
- Accelerate version: not installed
- xFormers version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: N/A

Beitragsleitfaden

Beitragsleitfaden öffnen

Rechercherichtung

Start at FlaxUNet2DConditionModel.from_pretrained and compare its dtype handling with the Flax training examples train_dreambooth_flax.py and train_text_to_image_flax.py. Reproduce the issue using the provided model-loading snippet and verify that the returned parameters use jnp.bfloat16 rather than float32.

Vom Indexierungsmodell aus dem Issue-Text verfasst.

Bewertung

Tech-Stack
python
Bereich
machine-learning
Issue-Typ
Bug
Schwierigkeit
3/5
Geschätzter Aufwand
1-2 Tage
Aktivitätsstatus
Veraltet
Klarheit
Größtenteils klar
Anfängerfreundlichkeit
45/100

Neue Issues direkt in Ihr Postfach

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