You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
[Guide] Quantize your Diffusion Models with bnb (#10012)
* chore: initial draft
* Apply suggestions from code review
Co-authored-by: Pedro Cuenca <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
* chore: link in place
* chore: review suggestions
* Apply suggestions from code review
Co-authored-by: Steven Liu <[email protected]>
* chore: review suggestions
* Update docs/source/en/quantization/bitsandbytes.md
Co-authored-by: Steven Liu <[email protected]>
* review suggestions
* chore: review suggestions
* Apply suggestions from code review
Co-authored-by: Steven Liu <[email protected]>
* adding same changes to 4 bit section
* review suggestions
---------
Co-authored-by: Sayak Paul <[email protected]>
Co-authored-by: Pedro Cuenca <[email protected]>
Co-authored-by: Steven Liu <[email protected]>
To use bitsandbytes, make sure you have the following libraries installed:
22
28
@@ -31,70 +37,167 @@ Now you can quantize a model by passing a [`BitsAndBytesConfig`] to [`~ModelMixi
31
37
32
38
Quantizing a model in 8-bit halves the memory-usage:
33
39
40
+
bitsandbytes is supported in both Transformers and Diffusers, so you can quantize both the
41
+
[`FluxTransformer2DModel`] and [`~transformers.T5EncoderModel`].
42
+
43
+
For Ada and higher-series GPUs. we recommend changing `torch_dtype` to `torch.bfloat16`.
44
+
45
+
> [!TIP]
46
+
> The [`CLIPTextModel`] and [`AutoencoderKL`] aren't quantized because they're already small in size and because [`AutoencoderKL`] only has a few `torch.nn.Linear` layers.
47
+
34
48
```py
35
-
from diffusers import FluxTransformer2DModel, BitsAndBytesConfig
49
+
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
50
+
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
By default, all the other modules such as `torch.nn.LayerNorm` are converted to `torch.float16`. You can change the data type of these modules with the `torch_dtype` parameter if you want:
74
+
By default, all the other modules such as `torch.nn.LayerNorm` are converted to `torch.float16`. You can change the data type of these modules with the `torch_dtype` parameter.
47
75
48
-
```py
49
-
from diffusers import FluxTransformer2DModel, BitsAndBytesConfig
Once a model is quantized, you can push the model to the Hub with the [`~ModelMixin.push_to_hub`] method. The quantization `config.json` file is pushed first, followed by the quantized model weights. You can also save the serialized 4-bit models locally with [`~ModelMixin.save_pretrained`].
When there is enough memory, you can also directly move the pipeline to the GPU with `.to("cuda")` and apply [`~DiffusionPipeline.enable_model_cpu_offload`] to optimize GPU memory usage.
116
+
117
+
Once a model is quantized, you can push the model to the Hub with the [`~ModelMixin.push_to_hub`] method. The quantization `config.json` file is pushed first, followed by the quantized model weights. You can also save the serialized 8-bit models locally with [`~ModelMixin.save_pretrained`].
63
118
64
119
</hfoption>
65
120
<hfoptionid="4-bit">
66
121
67
122
Quantizing a model in 4-bit reduces your memory-usage by 4x:
68
123
124
+
bitsandbytes is supported in both Transformers and Diffusers, so you can can quantize both the
125
+
[`FluxTransformer2DModel`] and [`~transformers.T5EncoderModel`].
126
+
127
+
For Ada and higher-series GPUs. we recommend changing `torch_dtype` to `torch.bfloat16`.
128
+
129
+
> [!TIP]
130
+
> The [`CLIPTextModel`] and [`AutoencoderKL`] aren't quantized because they're already small in size and because [`AutoencoderKL`] only has a few `torch.nn.Linear` layers.
131
+
69
132
```py
70
-
from diffusers import FluxTransformer2DModel, BitsAndBytesConfig
133
+
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
134
+
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
By default, all the other modules such as `torch.nn.LayerNorm` are converted to `torch.float16`. You can change the data type of these modules with the `torch_dtype` parameter if you want:
158
+
By default, all the other modules such as `torch.nn.LayerNorm` are converted to `torch.float16`. You can change the data type of these modules with the `torch_dtype` parameter.
82
159
83
-
```py
84
-
from diffusers import FluxTransformer2DModel, BitsAndBytesConfig
Setting `device_map="auto"` automatically fills all available space on the GPU(s) first, then the CPU, and finally, the hard drive (the absolute slowest option) if there is still not enough memory.
Call [`~ModelMixin.push_to_hub`] after loading it in 4-bit precision. You can also save the serialized 4-bit models locally with [`~ModelMixin.save_pretrained`].
When there is enough memory, you can also directly move the pipeline to the GPU with `.to("cuda")` and apply [`~DiffusionPipeline.enable_model_cpu_offload`] to optimize GPU memory usage.
199
+
200
+
Once a model is quantized, you can push the model to the Hub with the [`~ModelMixin.push_to_hub`] method. The quantization `config.json` file is pushed first, followed by the quantized model weights. You can also save the serialized 4-bit models locally with [`~ModelMixin.save_pretrained`].
NF4 is a 4-bit data type from the [QLoRA](https://hf.co/papers/2305.14314) paper, adapted for weights initialized from a normal distribution. You should use NF4 for training 4-bit base models. This can be configured with the `bnb_4bit_quant_type` parameter in the [`BitsAndBytesConfig`]:
200
303
201
304
```py
202
-
from diffusers import BitsAndBytesConfig
305
+
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
306
+
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
@@ -220,38 +340,74 @@ For inference, the `bnb_4bit_quant_type` does not have a huge impact on performa
220
340
Nested quantization is a technique that can save additional memory at no additional performance cost. This feature performs a second quantization of the already quantized weights to save an additional 0.4 bits/parameter.
221
341
222
342
```py
223
-
from diffusers import BitsAndBytesConfig
343
+
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
344
+
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
Once quantized, you can dequantize the model to the original precision but this might result in a small quality loss of the model. Make sure you have enough GPU RAM to fit the dequantized model.
376
+
Once quantized, you can dequantize a model to its original precision, but this might result in a small loss of quality. Make sure you have enough GPU RAM to fit the dequantized model.
240
377
241
378
```python
242
-
from diffusers import BitsAndBytesConfig
379
+
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
380
+
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig
0 commit comments