Add Finegrained FP8 #11647

MekkCyber · 2025-06-03T08:04:46Z

What does this PR do?

Adds finegrained FP8

HuggingFaceDocBuilderDev · 2025-06-03T08:12:46Z

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 · 2025-06-04T11:50:55Z

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).

sayakpaul

Thanks for starting this! Would be nice to also have some benchmarks:

With and without finegrained FP8 quant (with visual outputs)
With and without torch.compile

sayakpaul · 2025-06-16T12:05:12Z

src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quantizer.py

+class FinegrainedFP8Quantizer(DiffusersQuantizer):
+    """
+    FP8 quantization implementation supporting both standard and MoE models.
+    Supports both e4m3fn formats based on platform.


Can we expand on this a bit? What are both e4m3fn formats? How does that vary depending on the platform?

src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quantizer.py

tests/quantization/finegrained_fp8/test_finegrained_fp8.py

sayakpaul · 2025-06-16T12:20:33Z

tests/quantization/finegrained_fp8/test_finegrained_fp8.py

+    def _check_serialization_expected_slice(self, expected_slice, device):
+        quantized_model = self.get_dummy_model(device)
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            quantized_model.save_pretrained(tmp_dir, safe_serialization=False)
+            loaded_quantized_model = FluxTransformer2DModel.from_pretrained(
+                tmp_dir, torch_dtype=torch.bfloat16, use_safetensors=False
+            ).to(device=torch_device)
+
+        inputs = self.get_dummy_tensor_inputs(torch_device)
+        output = loaded_quantized_model(**inputs)[0]
+
+        output_slice = output.flatten()[-9:].detach().float().cpu().numpy()
+
+        self.assertTrue(numpy_cosine_similarity_distance(output_slice, expected_slice) < 1e-3)


I think instead of delegating certain calls to other methods, we can have all of the implementations under this one. This way, everything remains self-contained. Furthermore, since this test class doesn't have other tests, we don't have to modularize too much.

WDYT?

sayakpaul · 2025-06-16T12:22:33Z

tests/quantization/finegrained_fp8/test_finegrained_fp8.py

+        text_encoder = CLIPTextModel.from_pretrained(
+            model_id, subfolder="text_encoder", torch_dtype=torch.bfloat16, cache_dir=cache_dir
+        )
+        text_encoder_2 = T5EncoderModel.from_pretrained(
+            model_id, subfolder="text_encoder_2", torch_dtype=torch.bfloat16, cache_dir=cache_dir
+        )
+        tokenizer = CLIPTokenizer.from_pretrained(
+            model_id, subfolder="tokenizer", cache_dir=cache_dir
+        )
+        tokenizer_2 = AutoTokenizer.from_pretrained(
+            model_id, subfolder="tokenizer_2", cache_dir=cache_dir
+        )
+        vae = AutoencoderKL.from_pretrained(
+            model_id, subfolder="vae", torch_dtype=torch.bfloat16, cache_dir=cache_dir
+        )
+        scheduler = FlowMatchEulerDiscreteScheduler()


I don't think we need to initialize these components like this.

For example, if we do:

transformer = FluxTransformer2DModel.from_pretrained( model_id, subfolder="transformer", quantization_config=quantization_config, torch_dtype=torch.bfloat16, device_map=torch_device, ) pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")

It should work. It's simpler and I would prefer this method.

Sure will change that

tests/quantization/finegrained_fp8/test_finegrained_fp8.py

sayakpaul · 2025-06-16T12:23:27Z

tests/quantization/finegrained_fp8/test_finegrained_fp8.py

+        # A difference of 0.06 in normalized pixel space (-1 to 1), corresponds to a difference of
+        # 0.06 / 2 * 255 = 7.65 in pixel space (0 to 255). On our CI runners, the difference is about 0.04,
+        # on DGX it is 0.06, and on audace it is 0.037. So, we are using a tolerance of 0.06 here.
+        self.assertTrue(np.allclose(output, loaded_output, atol=0.06))


Can we reduce this tolerance?

…izer.py Co-authored-by: Sayak Paul <[email protected]>

sayakpaul

@MekkCyber LMK if there are any unresolved comments on my part. Excited to see this getting merged soon!

sayakpaul · 2025-06-18T01:29:34Z

docs/source/en/quantization/finegrained_fp8.md

+# FinegrainedFP8
+
+## Overview
+
+## Usage


sayakpaul · 2025-06-18T01:36:44Z

src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quantizer.py

+
+        module, tensor_name = get_module_from_name(model, param_name)
+        if isinstance(module, FP8Linear):
+            if self.pre_quantized or tensor_name == "bias":


checks if the param needs to be quantized ?

This.

Example:

diffusers/src/diffusers/models/model_loading_utils.py

Lines 295 to 302 in 1bc6f3d

elif is_quantized and (

hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict, param_device=param_device)

):

hf_quantizer.create_quantized_param(

model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype

)

else:

set_module_tensor_to_device(model, param_name, param_device, value=param, **set_module_kwargs)

I think this is because during this step, we call

diffusers/src/diffusers/models/modeling_utils.py

Line 1228 in 1bc6f3d

hf_quantizer.preprocess_model(

which does:

diffusers/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py

Line 254 in 1bc6f3d

def _process_model_before_weight_loading(

(the job replacing the regular linears with bnb linears)

LMK if you need more clarifications.

sayakpaul · 2025-06-18T01:39:43Z

src/diffusers/quantizers/finegrained_fp8/utils.py

+    return C
+
+
+# Python version of the above triton function, it's much slower than the triton version, for testing


Okay let's then add a comment clarifying the purpose. It will be helpful for the reviewers as well as the future readers.

sayakpaul · 2025-06-18T01:40:06Z

src/diffusers/quantizers/finegrained_fp8/utils.py

+    input_q: torch.Tensor,  # [batch, seq_len, hidden_dim]
+    weight_q: torch.Tensor,  # [out_features, hidden_dim]
+    input_scale: torch.Tensor,  # [batch * seq_len, num_input_groups]
+    weight_scale: torch.Tensor,  # [num_weight_blocks_m, num_weight_blocks_n]
+    block_size: Optional[Tuple[int, int]] = None,  # (M=128, N=128) for weights for example


Prefer adding the comments in the docstrings.

sayakpaul · 2025-06-18T01:43:14Z

tests/quantization/finegrained_fp8/test_finegrained_fp8.py

+            device_map=torch_device,
+        )
+
+        size_quantized_with_not_convert = self.get_model_size_in_bytes(quantized_model_with_not_convert)


Okay let's use it then.

DN6

Looks good 👍🏽 Just one minor comment on importing the quant config. And I think the test file needs clean up. Thanks!

DN6 · 2025-06-18T17:06:35Z

src/diffusers/__init__.py

@@ -96,6 +96,8 @@
 else:
    _import_structure["quantizers.quantization_config"].append("TorchAoConfig")

+_import_structure["quantizers.quantization_config"].append("FinegrainedFP8Config")


We haven't moved to setting torch and accelerate as required deps yet (will happen soon). So this would have to go under a is_torch_available() and not is_accelerate_available() type guard.

MekkCyber and others added 2 commits June 3, 2025 08:01

frst commit

d561f65

Merge branch 'main' into adding_fp8

87c4dc4

MekkCyber added 9 commits June 4, 2025 12:00

buffer to parameter

83163bc

resolve nan issue

e97e2ef

fix style

33b2be7

update

0eb3989

fix shapes

e6e0c4a

small fix

243173f

adding some tests

c3cf41f

rename

c7072fe

update tests

f54dcc0

sayakpaul requested review from sayakpaul and DN6 and removed request for sayakpaul June 16, 2025 08:12

sayakpaul reviewed Jun 16, 2025

View reviewed changes

sayakpaul requested a review from SunMarc June 16, 2025 12:24

MekkCyber and others added 2 commits June 17, 2025 13:51

Update src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quant…

94b2db5

…izer.py Co-authored-by: Sayak Paul <[email protected]>

Update src/diffusers/quantizers/finegrained_fp8/finegrained_fp8_quant…

2b71830

…izer.py Co-authored-by: Sayak Paul <[email protected]>

sayakpaul reviewed Jun 18, 2025

View reviewed changes

DN6 reviewed Jun 18, 2025

View reviewed changes

	elif is_quantized and (
	hf_quantizer.check_if_quantized_param(model, param, param_name, state_dict, param_device=param_device)
	):
	hf_quantizer.create_quantized_param(
	model, param, param_name, param_device, state_dict, unexpected_keys, dtype=dtype
	)
	else:
	set_module_tensor_to_device(model, param_name, param_device, value=param, **set_module_kwargs)

		return C


		# Python version of the above triton function, it's much slower than the triton version, for testing

Add Finegrained FP8 #11647

Are you sure you want to change the base?

Add Finegrained FP8 #11647

Uh oh!

Conversation

MekkCyber commented Jun 3, 2025

What does this PR do?

Uh oh!

HuggingFaceDocBuilderDev commented Jun 3, 2025

Uh oh!

sayakpaul commented Jun 4, 2025

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DN6 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!