huggingface / huggingface/diffusers

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

Aperta
#2,068 7 commenti 0 reazioni 0 assegnatari Vedi su GitHub
bug Good second issue
Lingua principale
Python
Stelle
34.5k
Fork
7.3k
Merge medio
3g 3h
PR unite (30g)
91

Descrizione

### 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

Guida per i contributori

Apri la guida per i contributori

Direzione di ricerca

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.

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

Valutazione

Stack tecnologico
python
Ambito
machine-learning
Tipo di issue
Bug
Difficoltà
3/5
Tempo stimato
1-2 giorni
Stato di attività
Ferma
Chiarezza
Abbastanza chiara
Idoneità per principianti
45/100

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.