Add Finegrained FP8

[MekkCyber]

What does this PR do?

Adds finegrained FP8

[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]

Just for bookkeeping, relaying stuff from our DM.

I had to make the following changes to make this PR work:

Expand

diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py
index 638c5fbfb..737525143 100644
--- a/src/diffusers/models/modeling_utils.py
+++ b/src/diffusers/models/modeling_utils.py
@@ -1238,8 +1238,8 @@ class ModelMixin(torch.nn.Module, PushToHubMixin):
         }
 
         # Dispatch model with hooks on all devices if necessary
-        print(model.transformer_blocks[0].attn.to_q.weight)
-        print(model.transformer_blocks[0].attn.to_q.weight_scale_inv)
+        # print(model.transformer_blocks[0].attn.to_q.weight)
+        # print(model.transformer_blocks[0].attn.to_q.weight_scale_inv)
         if device_map is not None:
             device_map_kwargs = {
                 "device_map": device_map,
diff --git a/src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quantizer.py b/src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quantizer.py
index 5dec8b0b8..7212befcd 100644
--- a/src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quantizer.py
+++ b/src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quantizer.py
@@ -90,9 +90,9 @@ class FinegrainedFP8Quantizer(DiffusersQuantizer):
         Quantizes weights to FP8 format using Block-wise quantization
         """
         # print("############ create quantized param ########")
-        from accelerate.utils import set_module_tensor_to_device
+        # from accelerate.utils import set_module_tensor_to_device
 
-        set_module_tensor_to_device(model, param_name, target_device, param_value)
+        # set_module_tensor_to_device(model, param_name, target_device, param_value)
 
         module, tensor_name = get_module_from_name(model, param_name)
 
@@ -131,8 +131,8 @@ class FinegrainedFP8Quantizer(DiffusersQuantizer):
         scale = scale.reshape(scale_orig_shape).squeeze().reciprocal()
 
         # Load into the model
-        module._buffers[tensor_name] = quantized_param.to(target_device)
-        module._buffers["weight_scale_inv"] = scale.to(target_device)
+        module._parameters[tensor_name] = quantized_param.to(target_device)
+        module._parameters["weight_scale_inv"] = scale.to(target_device)
         # print("_buffers[0]", module._buffers["weight_scale_inv"])
 
     def check_if_quantized_param(

Inference code:

import torch
from diffusers import FluxPipeline, AutoModel, FinegrainedFP8Config
from diffusers.quantizers.finegrained_fp8.utils import FP8Linear

model_id = "black-forest-labs/FLUX.1-dev"
dtype = torch.bfloat16

quantization_config = FinegrainedFP8Config(
    modules_to_not_convert=["norm", "proj_out", "x_embedder"], # weight_block_size=(32, 32)
)
transformer = AutoModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    torch_dtype=dtype,
    device_map="cuda"
)
pipe = FluxPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    torch_dtype=dtype,
)
pipe.to("cuda")

for name, module in pipe.transformer.named_modules():
    if isinstance(module, FP8Linear) and getattr(module, "weight_scale_inv", None) is not None:
        if module.weight_scale_inv.ndim == 1:
            print(name, module.weight_scale_inv.shape)


print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB")

prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")
print(f"Pipeline memory usage: {torch.cuda.max_memory_reserved() / 1024**3:.3f} GB")

The modules_to_not_convert includes proj_out and x_embedder because otherwise, we violate the shape constraint on scale (scale.ndim == 2).