8000 Merge branch 'main' into training-group-offloading-tests · huggingface/diffusers@d84dbab · GitHub
[go: up one dir, main page]
Skip to content

Commit d84dbab

Browse files
sayakpaulsayakpaul
authored
Merge branch 'main' into training-group-offloading-tests
2 parents 90f934d + 01abfc8 commit d84dbab

26 files changed

+907
-148
lines changed

‎.github/workflows/nightly_tests.yml

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,7 @@ jobs:
142142
HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }}
143143
# https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms
144144
CUBLAS_WORKSPACE_CONFIG: :16:8
145+
RUN_COMPILE: yes
145146
run: |
146147
python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
147148
-s -v -k "not Flax and not Onnx" \
@@ -525,6 +526,60 @@ jobs:
525526
pip install slack_sdk tabulate
526527
python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
527528
529+
run_nightly_pipeline_level_quantization_tests:
530+
name: Torch quantization nightly tests
531+
strategy:
532+
fail-fast: false
533+
max-parallel: 2
534+
runs-on:
535+
group: aws-g6e-xlarge-plus
536+
container:
537+
image: diffusers/diffusers-pytorch-cuda
538+
options: --shm-size "20gb" --ipc host --gpus 0
539+
steps:
540+
- name: Checkout diffusers
541+
uses: actions/checkout@v3
542+
with:
543+
fetch-depth: 2
544+
- name: NVIDIA-SMI
545+
run: nvidia-smi
546+
- name: Install dependencies
547+
run: |
548+
python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH"
549+
python -m uv pip install -e [quality,test]
550+
python -m uv pip install -U bitsandbytes optimum_quanto
551+
python -m uv pip install pytest-reportlog
552+
- name: Environment
553+
run: |
554+
python utils/print_env.py
555+
- name: Pipeline-level quantization tests on GPU
556+
env:
557+
HF_TOKEN: ${{ secrets.DIFFUSERS_HF_HUB_READ_TOKEN }}
558+
# https://pytorch.org/docs/stable/notes/randomness.html#avoiding-nondeterministic-algorithms
559+
CUBLAS_WORKSPACE_CONFIG: :16:8
560+
BIG_GPU_MEMORY: 40
561+
run: |
562+
python -m pytest -n 1 --max-worker-restart=0 --dist=loadfile \
563+
--make-reports=tests_pipeline_level_quant_torch_cuda \
564+
--report-log=tests_pipeline_level_quant_torch_cuda.log \
565+
tests/quantization/test_pipeline_level_quantization.py
566+
- name: Failure short reports
567+
if: ${{ failure() }}
568+
run: |
569+
cat reports/tests_pipeline_level_quant_torch_cuda_stats.txt
570+
cat reports/tests_pipeline_level_quant_torch_cuda_failures_short.txt
571+
- name: Test suite reports artifacts
572+
if: ${{ always() }}
573+
uses: actions/upload-artifact@v4
574+
with:
575+
name: torch_cuda_pipeline_level_quant_reports
576+
path: reports
577+
- name: Generate Report and Notify Channel
578+
if: always()
579+
run: |
580+
pip install slack_sdk tabulate
581+
python utils/log_reports.py >> $GITHUB_STEP_SUMMARY
582+
528583
# M1 runner currently not well supported
529584
# TODO: (Dhruv) add these back when we setup better testing for Apple Silicon
530585
# run_nightly_tests_apple_m1:

‎docs/source/en/api/quantization.md

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,17 @@ specific language governing permissions and limitations under the License.
1313

1414
# Quantization
1515

16-
Quantization techniques reduce memory and computational costs by representing weights and activations with lower-precision data types like 8-bit integers (int8). This enables loading larger models you normally wouldn't be able to fit into memory, and speeding up inference. Diffusers supports 8-bit and 4-bit quantization with [bitsandbytes](https://huggingface.co/docs/bitsandbytes/en/index).
17-
18-
Quantization techniques that aren't supported in Transformers can be added with the [`DiffusersQuantizer`] class.
16+
Quantization techniques reduce memory and computational costs by representing weights and activations with lower-precision data types like 8-bit integers (int8). This enables loading larger models you normally wouldn't be able to fit into memory, and speeding up inference.
1917

2018
<Tip>
2119

2220
Learn how to quantize models in the [Quantization](../quantization/overview) guide.
2321

2422
</Tip>
2523

24+
## PipelineQuantizationConfig
25+
26+
[[autodoc]] quantizers.PipelineQuantizationConfig
2627

2728
## BitsAndBytesConfig
2829

‎docs/source/en/quantization/overview.md

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,3 +39,90 @@ Diffusers currently supports the following quantization methods.
3939
- [Quanto](./quanto.md)
4040

4141
[This resource](https://huggingface.co/docs/transformers/main/en/quantization/overview#when-to-use-what) provides a good overview of the pros and cons of different quantization techniques.
42+
43+
## Pipeline-level quantization
44+
45+
Diffusers allows users to directly initialize pipelines from checkpoints that may contain quantized models ([example](https://huggingface.co/hf-internal-testing/flux.1-dev-nf4-pkg)). However, users may want to apply
46+
quantization on-the-fly when initializing a pipeline from a pre-trained and non-quantized checkpoint. You can
47+
do this with [`~quantizers.PipelineQuantizationConfig`].
48+
49+
Start by defining a `PipelineQuantizationConfig`:
50+
51+
```py
52+
import torch
53+
from diffusers import DiffusionPipeline
54+
from diffusers.quantizers.quantization_config import QuantoConfig
55+
from diffusers.quantizers import PipelineQuantizationConfig
56+
from transformers import BitsAndBytesConfig
57+
58+
pipeline_quant_config = PipelineQuantizationConfig(
59+
quant_mapping={
60+
"transformer": QuantoConfig(weights_dtype="int8"),
61+
"text_encoder_2": BitsAndBytesConfig(
62+
load_in_4bit=True, compute_dtype=torch.bfloat16
63+
),
64+
}
65+
)
66+
```
67+
68+
Then pass it to [`~DiffusionPipeline.from_pretrained`] and run inference:
69+
70+
```py
71+
pipe = DiffusionPipeline.from_pretrained(
72+
"black-forest-labs/FLUX.1-dev",
73+
quantization_config=pipeline_quant_config,
74+
torch_dtype=torch.bfloat16,
75+
).to("cuda")
76+
77+
image = pipe("photo of a cute dog").images[0]
78+
```
79+
80+
This method allows for more granular control over the quantization specifications of individual
81+
model-level components of a pipeline. It also allows for different quantization backends for
82+
different components. In the above example, you used a combination of Quanto and BitsandBytes. However,
83+
one caveat of this method is that users need to know which components come from `transformers` to be able
84+
to import the right quantization config class.
85+
86+
The other method is simpler in terms of experience but is
87+
less-flexible. Start by defining a `PipelineQuantizationConfig` but in a different way:
88+
89+
```py
90+
pipeline_quant_config = PipelineQuantizationConfig(
91+
quant_backend="bitsandbytes_4bit",
92+
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
93+
components_to_quantize=["transformer", "text_encoder_2"],
94+
)
95+
```
96+
97+
This `pipeline_quant_config` can now be passed to [`~DiffusionPipeline.from_pretrained`] similar to the above example.
98+
99+
In this case, `quant_kwargs` will be used to initialize the quantization specifications
100+
of the respective quantization configuration class of `quant_backend`. `components_to_quantize`
101+
is used to denote the components that will be quantized. For most pipelines, you would want to
102+
keep `transformer` in the list as that is often the most compute and memory intensive.
103+
104+
The config below will work for most diffusion pipelines that have a `transformer` component present.
105+
In most case, you will want to quantize the `transformer` component as that is often the most compute-
106+
intensive part of a diffusion pipeline.
107+
108+
```py
109+
pipeline_quant_config = PipelineQuantizationConfig(
110+
quant_backend="bitsandbytes_4bit",
111+
quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
112+
components_to_quantize=["transformer"],
113+
)
114+
```
115+
116+
Below is a list of the supported quantization backends available in both `diffusers` and `transformers`:
117+
118+
* `bitsandbytes_4bit`
119+
* `bitsandbytes_8bit`
120+
* `gguf`
121+
* `quanto`
122+
* `torchao`
123+
124+
125+
Diffusion pipelines can have multiple text encoders. [`FluxPipeline`] has two, for example. It's
126+
recommended to quantize the text encoders that are memory-intensive. Some examples include T5,
127+
Llama, Gemma, etc. In the above example, you quantized the T5 model of [`FluxPipeline`] through
128+
`text_encoder_2` while keeping the CLIP model intact (accessible through `text_encoder`).

‎src/diffusers/loaders/lora_conversion_utils.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1704,3 +1704,11 @@ def get_alpha_scales(down_weight, key):
17041704
converted_state_dict[f"transformer.{key}"] = converted_state_dict.pop(key)
17051705

17061706
10000 return converted_state_dict
1707+
1708+
1709+
def _convert_non_diffusers_hidream_lora_to_diffusers(state_dict, non_diffusers_prefix="diffusion_model"):
1710+
if not all(k.startswith(non_diffusers_prefix) for k in state_dict):
1711+
raise ValueError("Invalid LoRA state dict for HiDream.")
1712+
converted_state_dict = {k.removeprefix(f"{non_diffusers_prefix}."): v for k, v in state_dict.items()}
1713+
converted_state_dict = {f"transformer.{k}": v for k, v in converted_state_dict.items()}
1714+
return converted_state_dict

‎src/diffusers/loaders/lora_pipeline.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
_convert_hunyuan_video_lora_to_diffusers,
4444
_convert_kohya_flux_lora_to_diffusers,
4545
_convert_musubi_wan_lora_to_diffusers,
46+
_convert_non_diffusers_hidream_lora_to_diffusers,
4647
_convert_non_diffusers_lora_to_diffusers,
4748
_convert_non_diffusers_lumina2_lora_to_diffusers,
4849
_convert_non_diffusers_wan_lora_to_diffusers,
@@ -5371,7 +5372,6 @@ class HiDreamImageLoraLoaderMixin(LoraBaseMixin):
53715372

53725373
@classmethod
53735374
@validate_hf_hub_args
5374-
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.lora_state_dict
53755375
def lora_state_dict(
53765376
cls,
53775377
pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
@@ -5465,6 +5465,10 @@ def lora_state_dict(
54655465
logger.warning(warn_msg)
54665466
state_dict = {k: v for k, v in state_dict.items() if "dora_scale" not in k}
54675467

5468+
is_non_diffusers_format = any("diffusion_model" in k for k in state_dict)
5469+
if is_non_diffusers_format:
5470+
state_dict = _convert_non_diffusers_hidream_lora_to_diffusers(state_dict)
5471+
54685472
return state_dict
54695473

54705474
# Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights

‎src/diffusers/models/transformers/transformer_hunyuan_video_framepack.py

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,19 @@ def __init__(
152152

153153
# 1. Latent and condition embedders
154154
self.x_embedder = HunyuanVideoPatchEmbed((patch_size_t, patch_size, patch_size), in_channels, inner_dim)
155+
156+
# Framepack history projection embedder
157+
self.clean_x_embedder = None
158+
if has_clean_x_embedder:
159+
self.clean_x_embedder = HunyuanVideoHistoryPatchEmbed(in_channels, inner_dim)
160+
155161
self.context_embedder = HunyuanVideoTokenRefiner(
156162
text_embed_dim, num_attention_heads, attention_head_dim, num_layers=num_refiner_layers
157163
)
164+
165+
# Framepack image-conditioning embedder
166+
self.image_projection = FramepackClipVisionProjection(image_proj_dim, inner_dim) if has_image_proj else None
167+
158168
self.time_text_embed = HunyuanVideoConditionEmbedding(
159169
inner_dim, pooled_projection_dim, guidance_embeds, image_condition_type
160170
)
@@ -186,14 +196,7 @@ def __init__(
186196
self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6)
187197
self.proj_out = nn.Linear(inner_dim, patch_size_t * patch_size * patch_size * out_channels)
188198

189-
# Framepack specific modules
190-
self.image_projection = FramepackClipVisionProjection(image_proj_dim, inner_dim) if has_image_proj else None
191-
192-
self.clean_x_embedder = None
193-
if has_clean_x_embedder:
194-
self.clean_x_embedder = HunyuanVideoHistoryPatchEmbed(in_channels, inner_dim)
195-
196-
self.use_gradient_checkpointing = False
199+
self.gradient_checkpointing = False
197200

198201
def forward(
199202
self,

‎src/diffusers/pipelines/ltx/pipeline_ltx.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -789,6 +789,7 @@ def __call__(
789789
]
790790
latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
791791

792+
latents = latents.to(self.vae.dtype)
792793
video = self.vae.decode(latents, timestep, return_dict=False)[0]
793794
video = self.video_processor.postprocess_video(video, output_type=output_type)
794795

‎src/diffusers/pipelines/pipeline_loading_utils.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -675,8 +675,10 @@ def load_sub_model(
675675
use_safetensors: bool,
676676
dduf_entries: Optional[Dict[str, DDUFEntry]],
677677
provider_options: Any,
678+
quantization_config: Optional[Any] = None,
678679
):
679680
"""Helper method to load the module `name` from `library_name` and `class_name`"""
681+
from ..quantizers import PipelineQuantizationConfig
680682

681683
# retrieve class candidates
682684

@@ -769,6 +771,17 @@ def load_sub_model(
769771
else:
770772
loading_kwargs["low_cpu_mem_usage"] = False
771773

774+
if (
775+
quantization_config is not None
776+
and isinstance(quantization_config, PipelineQuantizationConfig)
777+
and issubclass(class_obj, torch.nn.Module)
778+
):
779+
model_quant_config = quantization_config._resolve_quant_config(
780+
is_diffusers=is_diffusers_model, module_name=name
781+
)
782+
if model_quant_config is not None:
783+
loading_kwargs["quantization_config"] = model_quant_config
784+
772785
# check if the module is in a subdirectory
773786
if dduf_entries:
774787
loading_kwargs["dduf_entries"] = dduf_entries

‎src/diffusers/pipelines/pipeline_utils.py

Lines changed: 6 additions & 0 deletions
@@ -47,6 +47,7 @@
Original file line numberDiff line numberDiff line change
4747
from ..models import AutoencoderKL
4848
from ..models.attention_processor import FusedAttnProcessor2_0
4949
from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, ModelMixin
50+
from ..quantizers import PipelineQuantizationConfig
5051
from ..quantizers.bitsandbytes.utils import _check_bnb_status
5152
from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
5253
from ..utils import (
@@ -725,6 +726,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
725726
use_safetensors = kwargs.pop("use_safetensors", None)
726727
use_onnx = kwargs.pop("use_onnx", None)
727728
load_connected_pipeline = kwargs.pop("load_connected_pipeline", False)
729+
quantization_config = kwargs.pop("quantization_config", None)
728730

729731
if torch_dtype is not None and not isinstance(torch_dtype, dict) and not isinstance(torch_dtype, torch.dtype):
730732
torch_dtype = torch.float32
@@ -741,6 +743,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
741743
" install accelerate\n```\n."
742744
)
743745

746+
if quantization_config is not None and not isinstance(quantization_config, PipelineQuantizationConfig):
747+
raise ValueError("`quantization_config` must be an instance of `PipelineQuantizationConfig`.")
748+
744749
if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
745750
raise NotImplementedError(
746751
"Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
@@ -1001,6 +1006,7 @@ def load_module(name, value):
10011006
use_safetensors=use_safetensors,
10021007
dduf_entries=dduf_entries,
10031008
provider_options=provider_options,
1009+
quantization_config=quantization_config,
10041010
)
10051011
logger.info(
10061012
f"Loaded {name} as {class_name} from `{name}` subfolder of {pretrained_model_name_or_path}."

0 commit comments

Comments
 (0)
0