huggingface / huggingface/diffusers

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

オープン
#2,068 コメント 7 件 リアクション 0 件 担当者 0 名 GitHub で見る
bug Good second issue
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

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

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

FlaxUNet2DConditionModel.from_pretrained から始め、Flax のトレーニング例である train_dreambooth_flax.py および train_text_to_image_flax.py と dtype の処理を比較します。提供されたモデル読み込みスニペットを使って問題を再現し、返されたパラメーターが float32 ではなく jnp.bfloat16 を使用していることを確認します。

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python
領域
machine-learning
issue の種類
バグ
難易度
3/5
見積もり時間
1〜2日
活発さ
停滞
明瞭さ
おおむね明確
初心者へのやさしさ
45/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。