Fix GGUF to Work Better with modules_to_not_convert / keep_in_fp32_modules

[dg845]

What does this PR do?

This PR contains several fixes so the GGUF loading and inference work better with module_to_not_convert and _keep_in_fp32_modules.

Changelist

1. src/diffusers/quantizers/gguf/utils.py
   1. _replace_with_gguf_linear: adds a check to see if any of the current module's named_children are in modules_to_not_convert, and if so, skip it. This allows us skip containers, rather than just leaf-level nn.Linear submodules as in the current code. For example, TimestepEmbedding modules are commonly added to _keep_in_fp32_modules (e.g. time_embedder in WanTransformer3DModel's WanTimeTextImageEmbedding condition embedder), but since they themselves contain leaf nn.Linear submodules such as linear_1, the current code will only check against leaf modules such as linear_1, and conclude incorrectly that they should be converted.
   2. _fused_mul_mat_gguf: in the UNQUANTIZED_TYPES case, also cast the dequantized weight to the activation x's dtype before performing the matrix multiplication, which should prevent dtype errors for BF16 weights.
2. src/diffusers/quantizers/gguf/gguf_quantizer.py
   1. GGUFQuantizer.create_quantized_param: handles modules_to_not_convert by dequantizing them, so that they end up in their original unquantized form. This is intended to handle the case where a module in self.modules_to_not_convert (or one of its children) is in the GGUF file. Since it is in the file, it will be converted to a GGUFParameter, but we don't want it to be quantized, so we convert it back here.

Inspired by GGUF debugging in #13551, in particular #13551 (comment).

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@DN6
@sayakpaul

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

I am a bit confused. If a param is already GGUFParameter type, then I'd assume that it's already quantized. In that case, how come dequantize -> type upcasting is the right sequence of ops?

What am I missing?

[dg845]

The idea is that the GGUF checkpoint might specify a quantization for a parameter that we do not want to be quantized, as expressed through either _keep_in_fp32_modules on the model: ModelMixin or modules_to_not_convert on GGUFQuantizationConfig.

When we load the GGUF state dict, these parameters will be placed into a GGUFParameter, and this happens before we load the weights into the model (e.g. in FromOriginalModelMixin.from_single_file). To respect modules_to_not_convert, we need to convert these back into normal (unquantized) parameters, which we do here at load time via dequantize_gguf_tensor. We then need to cast the parameter to the right compute dtype, which is torch.float32 for keep_in_fp32_modules and compute_dtype otherwise.

Currently, GGUFQuantizationConfig doesn't expose a modules_to_not_convert argument, but keep_in_fp32_modules are included in modules_to_not_convert:

diffusers/src/diffusers/quantizers/gguf/gguf_quantizer.py

Line 133 in d773308

self.modules_to_not_convert.extend(keep_in_fp32_modules)

So this change would affect only any specified _keep_in_fp32_modules right now.

[sayakpaul]

Would it break torch.compile compatibility for models that don't define modules_to_not_convert / keep_in_fp32_modules?

[dg845]

I'm not sure how it will interact with torch.compile, but this change mirrors the implementation used for quantized weight types (qweight_type in DEQUANT_TYPES):

diffusers/src/diffusers/quantizers/gguf/utils.py

Lines 98 to 99 in d773308

	weight = ops.ggml_dequantize(qweight, qweight_type, *shape)
	y = x @ weight.to(x.dtype).T

So I think it should be fine? (I think this change isn't specific to modules_to_not_convert, as the GGUF checkpoint could store weights in e.g. BF16 even if modules_to_not_convert is empty, which would then go through this code path.)

[sayakpaul]

Thanks! I left some comments. I think there should be a test for this in

diffusers/tests/models/testing_utils/quantization.py

Line 984 in a851ce1

class GGUFTesterMixin(GGUFConfigMixin, QuantizationTesterMixin):