From 3d7f4ad77c69c62b4b7cc7e9428300ee7ce2f484 Mon Sep 17 00:00:00 2001 From: LiangSu8899 <1342163629@qq.com> Date: Wed, 1 Jul 2026 06:40:24 -0400 Subject: [PATCH] Add after-load fusion for static quantized MLPs --- docs/source/en/fusion_mapping.md | 31 +- src/transformers/fusion_mapping.py | 56 +- src/transformers/integrations/__init__.py | 12 + .../integrations/static_quantized_mlp.py | 502 ++++++++++++++++++ .../quantizers/quantizer_finegrained_fp8.py | 6 + .../quantization/finegrained_fp8/test_fp8.py | 252 +++++++++ 6 files changed, 856 insertions(+), 3 deletions(-) create mode 100644 src/transformers/integrations/static_quantized_mlp.py diff --git a/docs/source/en/fusion_mapping.md b/docs/source/en/fusion_mapping.md index 2e4fd57305be..0f1e4b72f16f 100644 --- a/docs/source/en/fusion_mapping.md +++ b/docs/source/en/fusion_mapping.md @@ -58,9 +58,11 @@ This lets a fusion use a different runtime module structure while still loading Note: With the current monkey-patching mechanism, fusion registration is class-level: one compatible module class maps to one fused replacement class. +Some fusions require loaded quantized weights and scales. Those fusions still use `fusion_config` for validation, but are applied after checkpoint loading by the relevant quantizer instead of using pre-initialization monkey patching. + ## Current fusion families -Currently, `fusion_config` supports one fusion family: +Currently, `fusion_config` supports these fusion families: - `patch_embeddings` Enable with: @@ -71,6 +73,33 @@ Currently, `fusion_config` supports one fusion family: Effect: Replaces compatible `nn.Conv3d` patch embedding projections with equivalent flattened `nn.Linear` projections at runtime. +- `static_quantized_mlp` + Configure the Hub kernels per MLP pattern: + + ```python + fusion_config = { + "static_quantized_mlp": { + "input_quant": { + "repo_id": "owner/static-fp8-input-quant", + "version": 1, + }, + "gated_mlp": { + "repo_id": "owner/static-fp8-gated-mlp", + "version": 1, + }, + "dense_gelu_mlp": { + "repo_id": "owner/static-fp8-dense-gelu-mlp", + "version": 1, + }, + } + } + ``` + + Effect: + After static FP8 weights and activation scales are loaded, replaces eligible static per-tensor FineGrainedFP8 + `gate_proj/up_proj/down_proj` MLPs and `fc1/GELU/fc2` MLPs with fused Hub kernels. Unsupported quantization + layouts or activations keep the existing module path. The configured repositories must expose the functions required + by the selected pattern. Set `trust_remote_code=True` only for kernel repositories you trust. ## Extending fusion mapping diff --git a/src/transformers/fusion_mapping.py b/src/transformers/fusion_mapping.py index c3e0157024a3..053cde363bd9 100644 --- a/src/transformers/fusion_mapping.py +++ b/src/transformers/fusion_mapping.py @@ -67,6 +67,36 @@ def make_transforms(self, config: "PretrainedConfig") -> list[WeightTransform]: """Build the weight transforms needed to load and save the fused runtime layout.""" raise NotImplementedError + def validate_options(self, options: bool | Mapping[str, Any]) -> None: + """Validate user options for this fusion family.""" + return + + +class AfterLoadFusionSpec: + """Base recipe for fusions that require loaded weights and buffers.""" + + def validate_options(self, options: bool | Mapping[str, Any]) -> None: + """Validate user options for this fusion family.""" + raise NotImplementedError + + def apply(self, model: "PreTrainedModel", options: bool | Mapping[str, Any]) -> "PreTrainedModel": + """Apply the fusion after the model weights have been loaded.""" + raise NotImplementedError + + +class StaticQuantizedMLPAfterLoadFusionSpec(AfterLoadFusionSpec): + """Fuse eligible static quantized MLP modules after quantized weights and scales are loaded.""" + + def validate_options(self, options: bool | Mapping[str, Any]) -> None: + from .integrations.static_quantized_mlp import get_static_quantized_mlp_fusion_spec + + get_static_quantized_mlp_fusion_spec(options) + + def apply(self, model: "PreTrainedModel", options: bool | Mapping[str, Any]) -> "PreTrainedModel": + from .integrations.static_quantized_mlp import replace_with_static_quantized_mlp + + return replace_with_static_quantized_mlp(model, options) + class _FusedPatchEmbeddingMixin: def __init__(self, *args, **kwargs): @@ -232,7 +262,10 @@ def _register_module_fusion( register_checkpoint_conversion_mapping(model_type, converters, overwrite=True) -_FUSION_REGISTRY: dict[str, ModuleFusionSpec] = {"patch_embeddings": PatchEmbeddingsFusionSpec()} +_FUSION_REGISTRY: dict[str, ModuleFusionSpec | AfterLoadFusionSpec] = { + "patch_embeddings": PatchEmbeddingsFusionSpec(), + "static_quantized_mlp": StaticQuantizedMLPAfterLoadFusionSpec(), +} def _iter_enabled_fusions(fusion_config: Mapping[str, bool | Mapping[str, Any]]) -> list[str]: @@ -248,6 +281,7 @@ def _iter_enabled_fusions(fusion_config: Mapping[str, bool | Mapping[str, Any]]) raise ValueError( f"Invalid fusion config for {fusion_name}: expected `True`, `False`, or a mapping of options." ) + _FUSION_REGISTRY[fusion_name].validate_options(fusion_options) enabled_fusions.append(fusion_name) return enabled_fusions @@ -267,4 +301,22 @@ def register_fusion_patches( return for fusion_name in _iter_enabled_fusions(fusion_config): - _register_module_fusion(cls, config, fusion_name, _FUSION_REGISTRY[fusion_name]) + spec = _FUSION_REGISTRY[fusion_name] + if isinstance(spec, ModuleFusionSpec): + _register_module_fusion(cls, config, fusion_name, spec) + + +def apply_after_load_fusions( + model: "PreTrainedModel", fusion_config: Mapping[str, bool | Mapping[str, Any]] | None = None +) -> "PreTrainedModel": + """Apply fusions that require loaded weights and buffers.""" + + if not fusion_config: + return model + + for fusion_name in _iter_enabled_fusions(fusion_config): + spec = _FUSION_REGISTRY[fusion_name] + if isinstance(spec, AfterLoadFusionSpec): + model = spec.apply(model, fusion_config[fusion_name]) + + return model diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 2ecc31ae54cf..737544f9fcc5 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -153,6 +153,12 @@ "quanto": ["replace_with_quanto_layers"], "sinq": ["SinqDeserialize", "SinqQuantize"], "spqr": ["replace_with_spqr_linear"], + "static_quantized_mlp": [ + "StaticFP8DenseGELUMLP", + "StaticFP8GatedMLP", + "get_static_quantized_mlp_fusion_spec", + "replace_with_static_quantized_mlp", + ], "vptq": ["replace_with_vptq_linear"], } @@ -309,6 +315,12 @@ from .quanto import replace_with_quanto_layers from .sinq import SinqDeserialize, SinqQuantize from .spqr import replace_with_spqr_linear + from .static_quantized_mlp import ( + StaticFP8DenseGELUMLP, + StaticFP8GatedMLP, + get_static_quantized_mlp_fusion_spec, + replace_with_static_quantized_mlp, + ) from .vptq import replace_with_vptq_linear try: diff --git a/src/transformers/integrations/static_quantized_mlp.py b/src/transformers/integrations/static_quantized_mlp.py new file mode 100644 index 000000000000..177257bc1fe3 --- /dev/null +++ b/src/transformers/integrations/static_quantized_mlp.py @@ -0,0 +1,502 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import functools +from collections.abc import Mapping +from dataclasses import dataclass +from types import ModuleType +from typing import Any + +import torch +import torch.nn as nn + +from ..activations import ACT2FN +from ..utils import logging +from .finegrained_fp8 import _FP8_DTYPE, FP8Linear +from .hub_kernels import get_kernel +from .tensor_parallel import to_local + + +logger = logging.get_logger(__name__) + +_GATED_MLP_PATTERN = "gated_mlp" +_DENSE_GELU_MLP_PATTERN = "dense_gelu_mlp" +_INPUT_QUANT_KEY = "input_quant" +_SUPPORTED_PATTERN_KEYS = {_GATED_MLP_PATTERN, _DENSE_GELU_MLP_PATTERN, _INPUT_QUANT_KEY} + + +@dataclass(frozen=True) +class HubKernelEndpoint: + repo_id: str + version: int | str | None = 1 + revision: str | None = None + trust_remote_code: bool = False + + +@dataclass(frozen=True) +class StaticQuantizedMLPInputQuantSpec: + kernel: HubKernelEndpoint + func: str = "channel_scale_quantize_fp8_static_bf16" + + +@dataclass(frozen=True) +class StaticQuantizedGatedMLPSpec: + kernel: HubKernelEndpoint + swiglu_func: str = "fp8_swiglu_mlp_bf16" + geglu_func: str = "fp8_geglu_mlp_bf16" + + @property + def required_functions(self) -> tuple[str, str]: + return (self.swiglu_func, self.geglu_func) + + +@dataclass(frozen=True) +class StaticQuantizedDenseGELUMLPSpec: + kernel: HubKernelEndpoint + gelu_func: str = "fp8_gelu_mlp_bf16" + + @property + def required_functions(self) -> tuple[str]: + return (self.gelu_func,) + + +@dataclass(frozen=True) +class StaticQuantizedMLPFusionSpec: + input_quant: StaticQuantizedMLPInputQuantSpec + gated_mlp: StaticQuantizedGatedMLPSpec | None = None + dense_gelu_mlp: StaticQuantizedDenseGELUMLPSpec | None = None + + +def _endpoint_from_options( + options: str | Mapping[str, Any], + *, + pattern_name: str, +) -> HubKernelEndpoint: + if isinstance(options, str): + return HubKernelEndpoint(repo_id=options) + if not isinstance(options, Mapping): + raise ValueError(f"Static quantized MLP kernel options for {pattern_name!r} must be a repo string or mapping.") + + repo_id = options.get("repo_id") + if not isinstance(repo_id, str) or "/" not in repo_id: + raise ValueError( + f"Expected a valid Hub repo id for static quantized MLP fusion option {pattern_name!r}, got {repo_id!r}" + ) + + trust_remote_code = bool(options.get("trust_remote_code", False)) + return HubKernelEndpoint( + repo_id=repo_id, + version=options.get("version", 1), + revision=options.get("revision", None), + trust_remote_code=trust_remote_code, + ) + + +def _input_quant_spec_from_options(options: Mapping[str, Any]) -> StaticQuantizedMLPInputQuantSpec: + if _INPUT_QUANT_KEY not in options: + raise ValueError("Static quantized MLP fusion requires an explicit 'input_quant' kernel configuration.") + input_quant_options = options[_INPUT_QUANT_KEY] + endpoint = _endpoint_from_options(input_quant_options, pattern_name=_INPUT_QUANT_KEY) + func = "channel_scale_quantize_fp8_static_bf16" + if isinstance(input_quant_options, Mapping): + func = input_quant_options.get("func", func) + return StaticQuantizedMLPInputQuantSpec(kernel=endpoint, func=func) + + +def _gated_spec_from_options(options: str | Mapping[str, Any]) -> StaticQuantizedGatedMLPSpec: + endpoint = _endpoint_from_options(options, pattern_name=_GATED_MLP_PATTERN) + swiglu_func = "fp8_swiglu_mlp_bf16" + geglu_func = "fp8_geglu_mlp_bf16" + if isinstance(options, Mapping): + swiglu_func = options.get("swiglu_func", swiglu_func) + geglu_func = options.get("geglu_func", geglu_func) + return StaticQuantizedGatedMLPSpec(kernel=endpoint, swiglu_func=swiglu_func, geglu_func=geglu_func) + + +def _dense_gelu_spec_from_options(options: str | Mapping[str, Any]) -> StaticQuantizedDenseGELUMLPSpec: + endpoint = _endpoint_from_options(options, pattern_name=_DENSE_GELU_MLP_PATTERN) + gelu_func = "fp8_gelu_mlp_bf16" + if isinstance(options, Mapping): + gelu_func = options.get("gelu_func", gelu_func) + return StaticQuantizedDenseGELUMLPSpec(kernel=endpoint, gelu_func=gelu_func) + + +def get_static_quantized_mlp_fusion_spec( + options: bool | Mapping[str, Any], +) -> StaticQuantizedMLPFusionSpec | None: + if options is False: + return None + + if options is True: + raise ValueError("Static quantized MLP fusion requires explicit Hub kernel configurations.") + if not isinstance(options, Mapping): + raise ValueError(f"Static quantized MLP fusion config must be True, False, or a mapping, got {type(options)}") + + unknown_keys = set(options) - _SUPPORTED_PATTERN_KEYS + if unknown_keys: + raise ValueError( + "Unknown static quantized MLP fusion option(s): " + f"{sorted(unknown_keys)}. Supported options are {sorted(_SUPPORTED_PATTERN_KEYS)}." + ) + + gated_mlp = None + dense_gelu_mlp = None + if options.get(_GATED_MLP_PATTERN, False): + gated_mlp = _gated_spec_from_options(options[_GATED_MLP_PATTERN]) + if options.get(_DENSE_GELU_MLP_PATTERN, False): + dense_gelu_mlp = _dense_gelu_spec_from_options(options[_DENSE_GELU_MLP_PATTERN]) + + if gated_mlp is None and dense_gelu_mlp is None: + return None + + return StaticQuantizedMLPFusionSpec( + input_quant=_input_quant_spec_from_options(options), + gated_mlp=gated_mlp, + dense_gelu_mlp=dense_gelu_mlp, + ) + + +@functools.cache +def _load_hub_kernel( + repo_id: str, + version: int | str | None, + revision: str | None, + trust_remote_code: bool, +) -> ModuleType: + return get_kernel(repo_id, version=version, revision=revision, allow_all_kernels=trust_remote_code) + + +def _load_endpoint(endpoint: HubKernelEndpoint) -> ModuleType: + return _load_hub_kernel(endpoint.repo_id, endpoint.version, endpoint.revision, endpoint.trust_remote_code) + + +def _is_scalar_float_scale(tensor: torch.Tensor | None) -> bool: + return tensor is not None and tensor.numel() == 1 and tensor.dtype == torch.float32 + + +def _same_scalar_scale(lhs: torch.Tensor, rhs: torch.Tensor) -> bool: + if lhs.device.type == "meta" or rhs.device.type == "meta": + return True + return bool(torch.equal(lhs.detach().float().cpu(), rhs.detach().float().cpu())) + + +def _is_supported_static_fp8_linear(module: nn.Module, *, allow_bias: bool) -> bool: + if not isinstance(module, FP8Linear): + return False + if module.activation_scheme != "static" or module.block_size is not None: + return False + if module.weight.dtype != _FP8_DTYPE: + return False + if not allow_bias and module.bias is not None: + return False + return _is_scalar_float_scale(module.weight_scale_inv) and _is_scalar_float_scale(module.activation_scale) + + +def _activation_matches(module: nn.Module, *activation_names: str) -> bool: + activation = getattr(module, "act_fn", None) + if activation is None: + activation = getattr(module, "activation_fn", None) + if activation is None: + activation = getattr(module, "activation", None) + if activation is None: + return False + + for activation_name in activation_names: + expected = ACT2FN[activation_name] + if activation is expected or activation.__class__ is expected.__class__: + return True + return False + + +def _gated_activation_name(module: nn.Module) -> str | None: + if _activation_matches(module, "silu"): + return "silu" + if _activation_matches(module, "gelu_pytorch_tanh"): + return "gelu_pytorch_tanh" + return None + + +class StaticFP8GatedMLP(nn.Module): + """Runtime fused static-FP8 gated MLP backed by Hub kernels.""" + + def __init__( + self, + original_mlp: nn.Module, + input_quant_kernel: ModuleType, + input_quant_spec: StaticQuantizedMLPInputQuantSpec, + mlp_kernel: ModuleType, + mlp_spec: StaticQuantizedGatedMLPSpec, + activation_name: str, + ): + super().__init__() + self.input_quant_kernel = input_quant_kernel + self.input_quant_spec = input_quant_spec + self.mlp_kernel = mlp_kernel + self.mlp_spec = mlp_spec + self.activation_name = activation_name + + self.gate_proj = original_mlp.gate_proj + self.up_proj = original_mlp.up_proj + self.down_proj = original_mlp.down_proj + + self.register_buffer( + "gate_up_weight", + torch.cat([self.gate_proj.weight.detach(), self.up_proj.weight.detach()], dim=0).contiguous(), + persistent=False, + ) + self.register_buffer( + "input_channel_scale", + torch.ones(self.gate_proj.in_features, device=self.gate_proj.weight.device, dtype=torch.bfloat16), + persistent=False, + ) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + input_shape = input.shape + x = input.reshape(-1, input_shape[-1]).to(torch.bfloat16) + quant_fn = getattr(self.input_quant_kernel, self.input_quant_spec.func) + x_fp8 = quant_fn( + x, + to_local(self.input_channel_scale), + to_local(self.gate_proj.activation_scale), + ) + + mlp_fn_name = self.mlp_spec.swiglu_func if self.activation_name == "silu" else self.mlp_spec.geglu_func + mlp_fn = getattr(self.mlp_kernel, mlp_fn_name) + output = mlp_fn( + x_fp8, + to_local(self.gate_up_weight), + to_local(self.down_proj.weight), + to_local(self.gate_proj.activation_scale), + to_local(self.gate_proj.weight_scale_inv), + to_local(self.down_proj.activation_scale), + to_local(self.down_proj.weight_scale_inv), + ) + return output.reshape(*input_shape[:-1], output.shape[-1]).to(input.dtype) + + +class StaticFP8DenseGELUMLP(nn.Module): + """Runtime fused static-FP8 dense GELU MLP backed by Hub kernels.""" + + def __init__( + self, + original_mlp: nn.Module, + input_quant_kernel: ModuleType, + input_quant_spec: StaticQuantizedMLPInputQuantSpec, + mlp_kernel: ModuleType, + mlp_spec: StaticQuantizedDenseGELUMLPSpec, + ): + super().__init__() + self.input_quant_kernel = input_quant_kernel + self.input_quant_spec = input_quant_spec + self.mlp_kernel = mlp_kernel + self.mlp_spec = mlp_spec + + self.fc1 = original_mlp.fc1 + self.fc2 = original_mlp.fc2 + + fc1_bias = ( + self.fc1.bias.detach() + if self.fc1.bias is not None + else torch.zeros(self.fc1.out_features, device=self.fc1.weight.device) + ) + fc2_bias = ( + self.fc2.bias.detach() + if self.fc2.bias is not None + else torch.zeros(self.fc2.out_features, device=self.fc2.weight.device) + ) + self.register_buffer("fc1_bias_bf16", fc1_bias.to(torch.bfloat16), persistent=False) + self.register_buffer("fc2_bias_bf16", fc2_bias.to(torch.bfloat16), persistent=False) + self.register_buffer( + "input_channel_scale", + torch.ones(self.fc1.in_features, device=self.fc1.weight.device, dtype=torch.bfloat16), + persistent=False, + ) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + input_shape = input.shape + x = input.reshape(-1, input_shape[-1]).to(torch.bfloat16) + quant_fn = getattr(self.input_quant_kernel, self.input_quant_spec.func) + x_fp8 = quant_fn( + x, + to_local(self.input_channel_scale), + to_local(self.fc1.activation_scale), + ) + + mlp_fn = getattr(self.mlp_kernel, self.mlp_spec.gelu_func) + output = mlp_fn( + x_fp8, + to_local(self.fc1.weight), + to_local(self.fc1_bias_bf16), + to_local(self.fc2.weight), + to_local(self.fc2_bias_bf16), + to_local(self.fc1.activation_scale), + to_local(self.fc1.weight_scale_inv), + to_local(self.fc2.activation_scale), + to_local(self.fc2.weight_scale_inv), + ) + return output.reshape(*input_shape[:-1], output.shape[-1]).to(input.dtype) + + +def _can_fuse_static_fp8_gated_mlp(module: nn.Module) -> tuple[bool, str | None]: + if not all(hasattr(module, name) for name in ("gate_proj", "up_proj", "down_proj")): + return False, None + + gate_proj = module.gate_proj + up_proj = module.up_proj + down_proj = module.down_proj + if not all(_is_supported_static_fp8_linear(proj, allow_bias=False) for proj in (gate_proj, up_proj, down_proj)): + return False, None + + if gate_proj.in_features != up_proj.in_features or gate_proj.out_features != up_proj.out_features: + return False, None + if down_proj.in_features != gate_proj.out_features: + return False, None + if not _same_scalar_scale(gate_proj.activation_scale, up_proj.activation_scale): + return False, None + if not _same_scalar_scale(gate_proj.weight_scale_inv, up_proj.weight_scale_inv): + return False, None + + activation_name = _gated_activation_name(module) + return activation_name is not None, activation_name + + +def _can_fuse_static_fp8_dense_gelu_mlp(module: nn.Module) -> bool: + if not all(hasattr(module, name) for name in ("fc1", "fc2")): + return False + + fc1 = module.fc1 + fc2 = module.fc2 + if not all(_is_supported_static_fp8_linear(proj, allow_bias=True) for proj in (fc1, fc2)): + return False + if fc2.in_features != fc1.out_features: + return False + return _activation_matches(module, "gelu_pytorch_tanh") + + +def _validate_kernel_functions(kernel: ModuleType, repo_id: str, function_names: tuple[str, ...]) -> bool: + missing = [name for name in function_names if not hasattr(kernel, name)] + if missing: + logger.warning_once( + f"Static quantized MLP Hub kernel {repo_id!r} is missing required function(s) {missing}; " + "skipping this fusion pattern." + ) + return False + return True + + +def replace_with_static_quantized_mlp( + model: nn.Module, + options: bool | Mapping[str, Any], +) -> nn.Module: + fusion_spec = get_static_quantized_mlp_fusion_spec(options) + if fusion_spec is None: + return model + + try: + input_quant_kernel = _load_endpoint(fusion_spec.input_quant.kernel) + except Exception as e: + logger.warning_once( + "Unable to load static quantized MLP input-quant Hub kernel " + f"{fusion_spec.input_quant.kernel.repo_id!r}; keeping existing MLP modules. {e}" + ) + return model + + if not _validate_kernel_functions( + input_quant_kernel, + fusion_spec.input_quant.kernel.repo_id, + (fusion_spec.input_quant.func,), + ): + return model + + gated_kernel = None + if fusion_spec.gated_mlp is not None: + try: + gated_kernel = _load_endpoint(fusion_spec.gated_mlp.kernel) + except Exception as e: + logger.warning_once( + f"Unable to load static quantized gated MLP Hub kernel " + f"{fusion_spec.gated_mlp.kernel.repo_id!r}; skipping gated MLP fusion. {e}" + ) + gated_kernel = None + if gated_kernel is not None and not _validate_kernel_functions( + gated_kernel, + fusion_spec.gated_mlp.kernel.repo_id, + fusion_spec.gated_mlp.required_functions, + ): + gated_kernel = None + + dense_kernel = None + if fusion_spec.dense_gelu_mlp is not None: + try: + dense_kernel = _load_endpoint(fusion_spec.dense_gelu_mlp.kernel) + except Exception as e: + logger.warning_once( + f"Unable to load static quantized dense GELU MLP Hub kernel " + f"{fusion_spec.dense_gelu_mlp.kernel.repo_id!r}; skipping dense GELU MLP fusion. {e}" + ) + dense_kernel = None + if dense_kernel is not None and not _validate_kernel_functions( + dense_kernel, + fusion_spec.dense_gelu_mlp.kernel.repo_id, + fusion_spec.dense_gelu_mlp.required_functions, + ): + dense_kernel = None + + replaced_gated = 0 + replaced_dense_gelu = 0 + for module_name, module in list(model.named_modules()): + if gated_kernel is not None and fusion_spec.gated_mlp is not None: + can_fuse, activation_name = _can_fuse_static_fp8_gated_mlp(module) + if can_fuse: + model.set_submodule( + module_name, + StaticFP8GatedMLP( + module, + input_quant_kernel, + fusion_spec.input_quant, + gated_kernel, + fusion_spec.gated_mlp, + activation_name, + ), + ) + replaced_gated += 1 + continue + + if dense_kernel is not None and fusion_spec.dense_gelu_mlp is not None: + if _can_fuse_static_fp8_dense_gelu_mlp(module): + model.set_submodule( + module_name, + StaticFP8DenseGELUMLP( + module, + input_quant_kernel, + fusion_spec.input_quant, + dense_kernel, + fusion_spec.dense_gelu_mlp, + ), + ) + replaced_dense_gelu += 1 + + if fusion_spec.gated_mlp is not None and gated_kernel is not None and replaced_gated == 0: + logger.warning_once("Static quantized gated MLP fusion was enabled, but no eligible modules were found.") + if fusion_spec.dense_gelu_mlp is not None and dense_kernel is not None and replaced_dense_gelu == 0: + logger.warning_once("Static quantized dense GELU MLP fusion was enabled, but no eligible modules were found.") + + if replaced_gated or replaced_dense_gelu: + logger.info( + "Replaced static quantized MLP module(s) with fused Hub kernels: " + f"{replaced_gated} gated, {replaced_dense_gelu} dense GELU." + ) + + return model diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 60fcd4fef68c..ad6ec275e234 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -153,6 +153,12 @@ def _process_model_after_weight_loading(self, model, **kwargs): module = model.get_submodule(module_name) scale = getattr(module, attr) setattr(module, attr, torch.nn.Parameter(scale.data.to(ue8m0), requires_grad=False)) + + fusion_config = getattr(model.config, "fusion_config", None) + if fusion_config is not None and not self.quantization_config.dequantize: + from ..fusion_mapping import apply_after_load_fusions + + model = apply_after_load_fusions(model, fusion_config) return model def update_tp_plan(self, config): diff --git a/tests/quantization/finegrained_fp8/test_fp8.py b/tests/quantization/finegrained_fp8/test_fp8.py index 3548d9e9f3a9..a8a986963b42 100644 --- a/tests/quantization/finegrained_fp8/test_fp8.py +++ b/tests/quantization/finegrained_fp8/test_fp8.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import gc import tempfile import unittest from contextlib import ExitStack, contextmanager +from types import SimpleNamespace from unittest.mock import patch from parameterized import parameterized @@ -36,6 +38,8 @@ if is_torch_available(): import torch + import torch.nn as nn + import torch.nn.functional as F @contextmanager @@ -73,6 +77,254 @@ def test_from_dict(self): self.assertEqual(dict["quant_method"], quantization_config.quant_method) +def _quantize_static_fp8(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + scale = torch.clamp(tensor.float().abs().max() / 448.0, min=1e-6).reshape(1).to(torch.float32) + quantized = torch.clamp(tensor.float() / scale, -448.0, 448.0).to(torch.float8_e4m3fn) + return quantized.contiguous(), scale + + +def _dequantize_static_fp8(tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return tensor.float() * scale.float() + + +class _ReferenceInputQuantKernel: + def channel_scale_quantize_fp8_static_bf16(self, input, channel_scale, scale): + scaled = input.float() * channel_scale.float().reshape(1, -1) + return torch.clamp(scaled / scale.float(), -448.0, 448.0).to(torch.float8_e4m3fn) + + +class _ReferenceMLPKernel: + def fp8_swiglu_mlp_bf16( + self, + input, + gate_up_weight, + down_weight, + input_scale, + gate_up_weight_scale, + hidden_scale, + down_weight_scale, + ): + x = _dequantize_static_fp8(input, input_scale) + gate_up = F.linear(x, _dequantize_static_fp8(gate_up_weight, gate_up_weight_scale)) + gate, up = gate_up.chunk(2, dim=-1) + hidden = F.silu(gate) * up + hidden_fp8 = torch.clamp(hidden / hidden_scale.float(), -448.0, 448.0).to(torch.float8_e4m3fn) + out = F.linear( + _dequantize_static_fp8(hidden_fp8, hidden_scale), _dequantize_static_fp8(down_weight, down_weight_scale) + ) + return out.to(torch.bfloat16) + + def fp8_geglu_mlp_bf16( + self, + input, + gate_up_weight, + down_weight, + input_scale, + gate_up_weight_scale, + hidden_scale, + down_weight_scale, + ): + x = _dequantize_static_fp8(input, input_scale) + gate_up = F.linear(x, _dequantize_static_fp8(gate_up_weight, gate_up_weight_scale)) + gate, up = gate_up.chunk(2, dim=-1) + hidden = F.gelu(gate, approximate="tanh") * up + hidden_fp8 = torch.clamp(hidden / hidden_scale.float(), -448.0, 448.0).to(torch.float8_e4m3fn) + out = F.linear( + _dequantize_static_fp8(hidden_fp8, hidden_scale), _dequantize_static_fp8(down_weight, down_weight_scale) + ) + return out.to(torch.bfloat16) + + def fp8_gelu_mlp_bf16( + self, + input, + up_weight, + up_bias, + down_weight, + down_bias, + input_scale, + up_weight_scale, + hidden_scale, + down_weight_scale, + ): + x = _dequantize_static_fp8(input, input_scale) + up = F.linear(x, _dequantize_static_fp8(up_weight, up_weight_scale), up_bias.float()) + hidden = F.gelu(up, approximate="tanh") + hidden_fp8 = torch.clamp(hidden / hidden_scale.float(), -448.0, 448.0).to(torch.float8_e4m3fn) + out = F.linear( + _dequantize_static_fp8(hidden_fp8, hidden_scale), + _dequantize_static_fp8(down_weight, down_weight_scale), + down_bias.float(), + ) + return out.to(torch.bfloat16) + + +def _static_fp8_linear(in_features: int, out_features: int, *, has_bias: bool = False) -> nn.Module: + from transformers.integrations.finegrained_fp8 import FP8Linear + + linear = FP8Linear( + in_features, + out_features, + block_size=None, + activation_scheme="static", + has_bias=has_bias, + ) + weight = torch.randn(out_features, in_features, dtype=torch.float32) * 0.1 + weight_fp8, weight_scale = _quantize_static_fp8(weight) + linear.weight = nn.Parameter(weight_fp8, requires_grad=False) + linear.weight_scale_inv = nn.Parameter(weight_scale, requires_grad=False) + linear.activation_scale = nn.Parameter(torch.tensor([0.05], dtype=torch.float32), requires_grad=False) + if has_bias: + linear.bias = nn.Parameter(torch.randn(out_features, dtype=torch.bfloat16) * 0.01, requires_grad=False) + return linear + + +class _TinyStaticFP8GatedMLP(nn.Module): + def __init__(self, *, activation="gelu_pytorch_tanh"): + super().__init__() + from transformers.activations import ACT2FN + + self.gate_proj = _static_fp8_linear(8, 16) + self.up_proj = _static_fp8_linear(8, 16) + self.down_proj = _static_fp8_linear(16, 8) + self.up_proj.activation_scale = nn.Parameter( + self.gate_proj.activation_scale.detach().clone(), requires_grad=False + ) + self.up_proj.weight_scale_inv = nn.Parameter( + self.gate_proj.weight_scale_inv.detach().clone(), requires_grad=False + ) + self.act_fn = ACT2FN[activation] + + def forward(self, x): + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class _TinyStaticFP8DenseGELUMLP(nn.Module): + def __init__(self): + super().__init__() + from transformers.activations import ACT2FN + + self.fc1 = _static_fp8_linear(8, 16, has_bias=True) + self.fc2 = _static_fp8_linear(16, 8, has_bias=True) + self.activation_fn = ACT2FN["gelu_pytorch_tanh"] + + def forward(self, x): + return self.fc2(self.activation_fn(self.fc1(x))) + + +class _TinyStaticFP8Model(nn.Module): + def __init__(self): + super().__init__() + self.gated = _TinyStaticFP8GatedMLP() + self.dense = _TinyStaticFP8DenseGELUMLP() + self.config = SimpleNamespace() + + def forward(self, x): + return self.gated(x), self.dense(x) + + +class StaticQuantizedMLPFusionTest(unittest.TestCase): + fusion_config = { + "input_quant": {"repo_id": "reference/input-quant"}, + "gated_mlp": {"repo_id": "reference/gated"}, + "dense_gelu_mlp": {"repo_id": "reference/dense"}, + } + + def _kernel_loader(self, repo_id, version=None, revision=None, trust_remote_code=False): + del version, revision, trust_remote_code + if repo_id == "reference/input-quant": + return _ReferenceInputQuantKernel() + if repo_id in {"reference/gated", "reference/dense"}: + return _ReferenceMLPKernel() + raise AssertionError(f"Unexpected repo id {repo_id}") + + def test_fusion_config_parsing(self): + from transformers.integrations.static_quantized_mlp import get_static_quantized_mlp_fusion_spec + + spec = get_static_quantized_mlp_fusion_spec( + { + "input_quant": {"repo_id": "reference/input-quant", "trust_remote_code": True}, + "gated_mlp": {"repo_id": "reference/gated", "version": 2}, + } + ) + self.assertEqual(spec.input_quant.kernel.repo_id, "reference/input-quant") + self.assertTrue(spec.input_quant.kernel.trust_remote_code) + self.assertEqual(spec.gated_mlp.kernel.version, 2) + self.assertIsNone(spec.dense_gelu_mlp) + + def test_fusion_config_rejects_unknown_pattern(self): + from transformers.fusion_mapping import register_fusion_patches + + with self.assertRaisesRegex(ValueError, "Unknown static quantized MLP fusion option"): + register_fusion_patches(object, SimpleNamespace(), {"static_quantized_mlp": {"unknown_mlp": True}}) + + def test_static_quantized_mlp_requires_explicit_kernel_config(self): + from transformers.fusion_mapping import register_fusion_patches + + with self.assertRaisesRegex(ValueError, "requires explicit Hub kernel configurations"): + register_fusion_patches(object, SimpleNamespace(), {"static_quantized_mlp": True}) + + def test_static_fp8_gated_and_dense_gelu_mlp_replacement(self): + from transformers.integrations.static_quantized_mlp import ( + StaticFP8DenseGELUMLP, + StaticFP8GatedMLP, + replace_with_static_quantized_mlp, + ) + + torch.manual_seed(0) + model = _TinyStaticFP8Model() + reference = copy.deepcopy(model) + x = torch.randn(2, 3, 8, dtype=torch.bfloat16) + + with patch("transformers.integrations.static_quantized_mlp._load_hub_kernel", side_effect=self._kernel_loader): + model = replace_with_static_quantized_mlp(model, self.fusion_config) + reference = replace_with_static_quantized_mlp(reference, self.fusion_config) + + self.assertIsInstance(model.gated, StaticFP8GatedMLP) + self.assertIsInstance(model.dense, StaticFP8DenseGELUMLP) + torch.testing.assert_close(model.gated(x), reference.gated(x)) + torch.testing.assert_close(model.dense(x), reference.dense(x)) + + state_dict = model.state_dict() + self.assertIn("gated.gate_proj.weight", state_dict) + self.assertIn("dense.fc1.weight", state_dict) + self.assertNotIn("gated.gate_up_weight", state_dict) + self.assertNotIn("dense.input_channel_scale", state_dict) + + def test_static_fp8_dense_gelu_unsupported_activation_falls_back(self): + from transformers.activations import ACT2FN + from transformers.integrations.static_quantized_mlp import ( + StaticFP8DenseGELUMLP, + replace_with_static_quantized_mlp, + ) + + model = _TinyStaticFP8Model() + model.dense.activation_fn = ACT2FN["relu"] + + with patch("transformers.integrations.static_quantized_mlp._load_hub_kernel", side_effect=self._kernel_loader): + model = replace_with_static_quantized_mlp( + model, + { + "input_quant": self.fusion_config["input_quant"], + "dense_gelu_mlp": self.fusion_config["dense_gelu_mlp"], + }, + ) + + self.assertNotIsInstance(model.dense, StaticFP8DenseGELUMLP) + + def test_quantizer_after_load_applies_static_quantized_mlp_fusion(self): + from transformers.integrations.static_quantized_mlp import StaticFP8DenseGELUMLP, StaticFP8GatedMLP + + model = _TinyStaticFP8Model() + model.config.fusion_config = {"static_quantized_mlp": self.fusion_config} + quantizer = FineGrainedFP8HfQuantizer(FineGrainedFP8Config(activation_scheme="static", weight_block_size=None)) + + with patch("transformers.integrations.static_quantized_mlp._load_hub_kernel", side_effect=self._kernel_loader): + model = quantizer._process_model_after_weight_loading(model) + + self.assertIsInstance(model.gated, StaticFP8GatedMLP) + self.assertIsInstance(model.dense, StaticFP8DenseGELUMLP) + + @slow @require_accelerate @require_torch_accelerator