transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc2__py3-none-any.whl

transformers/integrations/eetq.py

@@ -12,104 +12,107 @@
 # 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 ..utils import is_accelerate_available, is_eetq_available, logging
+from ..core_model_loading import ConversionOps
+from ..quantizers.quantizers_utils import should_convert_module
+from ..utils import is_torch_available, logging
 
 
-if is_eetq_available():
-    import eetq
+if is_torch_available():
+    import torch
     import torch.nn as nn
 
-if is_accelerate_available():
-    from accelerate import init_empty_weights
 
 logger = logging.get_logger(__name__)
 
 
-def _replace_with_eetq_linear(
-    model,
-    modules_to_not_convert=None,
-    current_key_name=None,
-    quantization_config=None,
-    has_been_replaced=False,
-    pre_quantized=False,
-):
-    """
-    Private method that wraps the recursion for module replacement.
+class EetqQuantize(ConversionOps):
+    def __init__(self, hf_quantizer):
+        self.hf_quantizer = hf_quantizer
 
-    Returns the converted model and a boolean that indicates if the conversion has been successful or not.
-    """
-    if current_key_name is None:
-        current_key_name = []
-
-    for name, module in model.named_children():
-        current_key_name.append(name)
-
-        if (isinstance(module, nn.Linear)) and name not in modules_to_not_convert:
-            # Check if the current key is not in the `modules_to_not_convert`
-            current_key_name_str = ".".join(current_key_name)
-            if not any(
-                (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert
-            ):
-                with init_empty_weights():
-                    in_features = module.in_features
-                    out_features = module.out_features
-                    model._modules[name] = eetq.EetqLinear(
-                        in_features, out_features, module.bias is not None, module.weight.device
-                    )
-                    if pre_quantized:
-                        model._modules[name].register_scale(module.weight.device)
-                    has_been_replaced = True
-
-                    # Force requires grad to False to avoid unexpected errors
-                    model._modules[name].requires_grad_(False)
-        if len(list(module.children())) > 0:
-            _, has_been_replaced = _replace_with_eetq_linear(
-                module,
-                modules_to_not_convert,
-                current_key_name,
-                quantization_config,
-                has_been_replaced=has_been_replaced,
-                pre_quantized=pre_quantized,
-            )
-        # Remove the last key for recursion
-        current_key_name.pop(-1)
-    return model, has_been_replaced
-
-
-def replace_with_eetq_linear(
-    model, modules_to_not_convert=None, current_key_name=None, quantization_config=None, pre_quantized=False
-):
-    """
-    A helper function to replace all `torch.nn.Linear` modules by `eetq.EetqLinear` modules from the `eetq`
-    library. This will enable running your models using high performance int8 weight-only gemm kerner from
-    FasterTransformer and TensorRT-LLM. Make sure `eetq` compiled with the correct CUDA
-    version of your hardware is installed before running this function. EETQ shall be installed via the source
-    'https://github.com/NetEase-FuXi/EETQ'
+    def convert(
+        self, input_dict: dict[str, list[torch.Tensor]], full_layer_name: str | None = None, **kwargs
+    ) -> dict[str, torch.Tensor]:
+        _, value = tuple(input_dict.items())[0]
+        value = value[0]
+
+        value_device = value.device
+        int8_weight = torch.t(value).contiguous().cpu()
+        int8_weight, scales = eetq_kernels_hub.quant_weights(int8_weight, torch.int8, False)
+
+        int8_weight = int8_weight.to(value_device)
+        scales = scales.to(value_device)
+
+        return {full_layer_name: int8_weight, f"{full_layer_name}_scales": scales}
+
+
+class EetqLinearMMFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, weight, scales, bias=None):
+        # The forward pass can use ctx.
+        ctx.save_for_backward(x, weight, scales, bias)
+        output = eetq_kernels_hub.w8_a16_gemm(x, weight, scales)
+        output = output + bias if bias is not None else output
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weight, scales, bias = ctx.saved_tensors
+        identity = torch.eye(weight.shape[0]).to(weight.device).to(input.dtype)
+
+        # Dequantize the weight
+        weight = eetq_kernels_hub.w8_a16_gemm(identity, weight, scales)
 
-    The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should
-    be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no
-    CPU/GPU memory is required to run this function. Each weight will be quantized along the channel.
+        if ctx.needs_input_grad[0]:
+            # 2D matrix multiplication, unsqueeze to 3D
+            grad_input = grad_output.squeeze(0).matmul(weight.transpose(0, 1)).unsqueeze(0)
+
+        return grad_input, None, None, None
+
+
+class EetqLinear(nn.Module):
+    def __init__(self, in_features, out_features, dtype=torch.int8, bias=False):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty((in_features, out_features), dtype=dtype), requires_grad=False)
+        self.weight_scales = nn.Parameter(torch.empty((out_features), dtype=torch.float16))
+        if bias:
+            self.bias = nn.Parameter(torch.empty((out_features), dtype=torch.float16))
+        else:
+            self.bias = None
+
+    def forward(self, input):
+        output = EetqLinearMMFunction.apply(input, self.weight, self.weight_scales, self.bias)
+        return output
+
+
+def replace_with_eetq_linear(model, modules_to_not_convert: list[str] | None = None, pre_quantized=False):
+    """
+    A helper function to replace all `torch.nn.Linear` modules by `EetqLinear` modules.
 
     Parameters:
         model (`torch.nn.Module`):
             Input model or `torch.nn.Module` as the function is run recursively.
-        modules_to_not_convert (`list[`str`]`, *optional*, defaults to `["lm_head"]`):
+        modules_to_not_convert (`list[`str`]`, *optional*, defaults to `None`):
             Names of the modules to not convert in `EetqLinear`. In practice we keep the `lm_head` in full precision
             for numerical stability reasons.
-        current_key_name (`list[`str`]`, *optional*):
-            An array to track the current key of the recursion. This is used to check whether the current key (part of
-            it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or
-            `disk`).
     """
-
-    modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert
-
-    if quantization_config.modules_to_not_convert is not None:
-        modules_to_not_convert.extend(quantization_config.modules_to_not_convert)
-    modules_to_not_convert = list(set(modules_to_not_convert))
-    model, has_been_replaced = _replace_with_eetq_linear(
-        model, modules_to_not_convert, current_key_name, quantization_config, pre_quantized=pre_quantized
-    )
+    from .hub_kernels import get_kernel
+
+    global eetq_kernels_hub
+    eetq_kernels_hub = get_kernel("kernels-community/quantization-eetq")
+
+    has_been_replaced = False
+    # we need this to correctly materialize the weights during quantization
+    module_kwargs = {} if pre_quantized else {"dtype": None}
+    for module_name, module in model.named_modules():
+        if not should_convert_module(module_name, modules_to_not_convert):
+            continue
+        with torch.device("meta"):
+            if isinstance(module, nn.Linear):
+                new_module = EetqLinear(
+                    module.in_features, module.out_features, bias=module.bias is not None, **module_kwargs
+                )
+                model.set_submodule(module_name, new_module)
+                has_been_replaced = True
 
     if not has_been_replaced:
         logger.warning(

transformers/integrations/fbgemm_fp8.py

@@ -12,8 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import lru_cache
+from typing import Optional
+
 from ..activations import ACT2FN
-from ..utils import is_accelerate_available, is_fbgemm_gpu_available, is_torch_available, logging
+from ..core_model_loading import ConversionOps
+from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module
+from ..utils import (
+    is_accelerate_available,
+    is_fbgemm_gpu_available,
+    is_torch_available,
+    is_torch_xpu_available,
+    logging,
+)
 
 
 if is_torch_available():
@@ -23,24 +34,83 @@ if is_torch_available():
 if is_accelerate_available():
     from accelerate import init_empty_weights
 
-if is_fbgemm_gpu_available():
+_is_torch_xpu_available = is_torch_xpu_available()
+
+if is_fbgemm_gpu_available() and not _is_torch_xpu_available:
     import fbgemm_gpu.experimental.gen_ai  # noqa: F401
 
 logger = logging.get_logger(__name__)
 
 
+class FbgemmFp8Quantize(ConversionOps):
+    def __init__(self, hf_quantizer):
+        self.hf_quantizer = hf_quantizer
+
+    def convert(
+        self,
+        input_dict: dict[str, torch.Tensor | list[torch.Tensor]],
+        model: Optional[torch.nn.Module] = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        target_key, value = tuple(input_dict.items())[0]
+        value = value[0]
+
+        from ..integrations import FbgemmFp8Llama4TextExperts
+
+        module, tensor_name = get_module_from_name(model, target_key)
+
+        if isinstance(module, FbgemmFp8Llama4TextExperts):
+            if tensor_name == "gate_up_proj":
+                # Process each expert separately
+                # Transpose the second and third dimension
+                transposed_param = value.transpose(1, 2)
+
+                # Reshape to 2D for quantization
+                original_shape = transposed_param.shape
+                flattened_param = transposed_param.reshape(-1, original_shape[-1])
+
+                # Quantize using per row instead of per column
+                new_value_flat, weight_scale_flat = quantize_fp8_per_row(flattened_param)
+
+                # Reshape back to original dimensions
+                new_value = new_value_flat.reshape(original_shape)
+                new_value = new_value.transpose(1, 2)
+                weight_scale = weight_scale_flat.reshape(original_shape[0], 1, original_shape[1])
+            elif tensor_name == "down_proj":
+                # Process each expert separately
+                # Transpose the weights for proper quantization
+                transposed_param = value.transpose(1, 2)
+
+                # Reshape to 2D for quantization
+                original_shape = transposed_param.shape
+                flattened_param = transposed_param.reshape(-1, original_shape[-1])
+
+                # Quantize using per column
+                new_value_flat, weight_scale_flat = quantize_fp8_per_row(flattened_param)
+
+                # Reshape back to original dimensions
+                new_value = new_value_flat.reshape(original_shape)
+                new_value = new_value.transpose(1, 2)
+                weight_scale = weight_scale_flat.reshape(original_shape[0], original_shape[1], 1)
+        else:
+            new_value, weight_scale = quantize_fp8_per_row(value)
+            weight_scale = torch.nn.Parameter(weight_scale.view(weight_scale.shape[0], 1))
+
+        return {target_key: torch.nn.Parameter(new_value), f"{target_key}_scale": weight_scale}
+
+
 class FbgemmFp8Linear(torch.nn.Linear):
-    def __init__(self, in_features, out_features, bias, weight_dtype=torch.float32):
+    def __init__(self, in_features, out_features, bias, dtype=torch.float8_e4m3fn):
         super().__init__(in_features, out_features, bias)
         self.in_features = in_features
         self.out_features = out_features
 
-        self.weight = torch.nn.Parameter(torch.zeros((out_features, in_features), dtype=torch.float8_e4m3fn))
-        self.weight_scale = torch.nn.Parameter(torch.zeros((out_features, 1), dtype=weight_dtype))
+        self.weight = torch.nn.Parameter(torch.zeros((out_features, in_features), dtype=dtype))
+        self.weight_scale = torch.nn.Parameter(torch.zeros((out_features, 1), dtype=torch.float32))
         self.register_buffer("input_scale_ub", torch.zeros([1], dtype=torch.float), persistent=False)
 
         if bias:
-            self.bias = torch.nn.Parameter(torch.zeros((self.out_features), dtype=weight_dtype))
+            self.bias = torch.nn.Parameter(torch.zeros((self.out_features), dtype=torch.float32))
         else:
             self.bias = None
 
@@ -49,18 +119,26 @@ class FbgemmFp8Linear(torch.nn.Linear):
         output_shape = (*x.shape[:-1], -1)
         # x_quantized and x_scale are not necessarily on the same device as x, this is an issue.
         # https://github.com/pytorch/FBGEMM/blob/e08af8539c391437f447173863df0f3f6f6f1855/fbgemm_gpu/experimental/gen_ai/src/quantize/quantize.cu#L1237C3-L1237C45
-        x_quantized, x_scale = torch.ops.fbgemm.quantize_fp8_per_row(
-            x.view(-1, x.shape[-1]).contiguous(), scale_ub=self.input_scale_ub
-        )
+        x_quantized, x_scale = quantize_fp8_per_row(x.view(-1, x.shape[-1]).contiguous(), scale_ub=self.input_scale_ub)
         # moving x_quantized, x_scale here creates glibberish output ... However, if we move the output, it works
         # x_quantized, x_scale = x_quantized.to(x.device), x_scale.to(x.device)
 
         # The computation still happens on the device where self.weight is even if x_quantized is not on the same device as self.weight
         weight_scale_float32 = self.weight_scale.to(torch.float32)
-        output = torch.ops.fbgemm.f8f8bf16_rowwise(
-            x_quantized, self.weight, x_scale, weight_scale_float32, use_fast_accum=True
-        )
-        output = output + self.bias if self.bias is not None else output
+        if _is_torch_xpu_available:
+            output = torch._scaled_mm(
+                x_quantized,
+                self.weight.t(),
+                scale_a=x_scale.unsqueeze(-1),
+                scale_b=weight_scale_float32.t(),
+                out_dtype=x.dtype,
+                bias=self.bias,
+            )
+        else:
+            output = torch.ops.fbgemm.f8f8bf16_rowwise(
+                x_quantized, self.weight, x_scale, weight_scale_float32, use_fast_accum=True
+            )
+            output = output + self.bias if self.bias is not None else output
         # Hacky for now, we have the output to the device of x
         output = output.to(x.device)
         output = output.reshape(output_shape)
@@ -112,168 +190,136 @@ class FbgemmFp8Llama4TextExperts(nn.Module):
             expert_hidden = hidden_states[i]
             expert_hidden_reshaped = expert_hidden.reshape(-1, self.hidden_size)
             # Quantize for this expert
-            expert_quantized, expert_scale = torch.ops.fbgemm.quantize_fp8_per_row(
+            expert_quantized, expert_scale = quantize_fp8_per_row(
                 expert_hidden_reshaped, num_tokens, self.input_scale_ub
             )
             sharded_expert_dim = self.gate_up_proj.shape[-1] // 2
             gate_up_proj_scale_float32 = self.gate_up_proj_scale.to(torch.float32)
+            if _is_torch_xpu_available:
+                gate = torch._scaled_mm(
+                    expert_quantized,
+                    self.gate_up_proj[i].transpose(0, 1)[:sharded_expert_dim].contiguous().t(),
+                    scale_a=expert_scale.unsqueeze(-1),
+                    scale_b=gate_up_proj_scale_float32[i][0][:sharded_expert_dim].view(-1, 1).contiguous().t(),
+                    out_dtype=hidden_states.dtype,
+                )
+                up = torch._scaled_mm(
+                    expert_quantized,
+                    self.gate_up_proj[i].transpose(0, 1)[sharded_expert_dim:].contiguous().t(),
+                    scale_a=expert_scale.unsqueeze(-1),
+                    scale_b=gate_up_proj_scale_float32[i][0][sharded_expert_dim:].view(-1, 1).contiguous().t(),
+                    out_dtype=hidden_states.dtype,
+                )
+            else:
+                gate = torch.ops.fbgemm.f8f8bf16_rowwise(
+                    expert_quantized,
+                    self.gate_up_proj[i].transpose(0, 1)[:sharded_expert_dim].contiguous(),
+                    expert_scale,
+                    gate_up_proj_scale_float32[i][0][:sharded_expert_dim].view(-1, 1).contiguous(),
+                    use_fast_accum=True,
+                )
 
-            gate = torch.ops.fbgemm.f8f8bf16_rowwise(
-                expert_quantized,
-                self.gate_up_proj[i].transpose(0, 1)[:sharded_expert_dim].contiguous(),
-                expert_scale,
-                gate_up_proj_scale_float32[i][0][:sharded_expert_dim].view(-1, 1).contiguous(),
-                use_fast_accum=True,
-            )
-
-            up = torch.ops.fbgemm.f8f8bf16_rowwise(
-                expert_quantized,
-                self.gate_up_proj[i].transpose(0, 1)[sharded_expert_dim:].contiguous(),
-                expert_scale,
-                gate_up_proj_scale_float32[i][0][sharded_expert_dim:].view(-1, 1).contiguous(),
-                use_fast_accum=True,
-            )
+                up = torch.ops.fbgemm.f8f8bf16_rowwise(
+                    expert_quantized,
+                    self.gate_up_proj[i].transpose(0, 1)[sharded_expert_dim:].contiguous(),
+                    expert_scale,
+                    gate_up_proj_scale_float32[i][0][sharded_expert_dim:].view(-1, 1).contiguous(),
+                    use_fast_accum=True,
+                )
 
             activated = up * self.act_fn(gate)
 
-            activated_quantized, activated_scale = torch.ops.fbgemm.quantize_fp8_per_row(
-                activated, num_tokens, self.input_scale_ub
-            )
+            activated_quantized, activated_scale = quantize_fp8_per_row(activated, num_tokens, self.input_scale_ub)
 
             down_proj_scale_float32 = self.down_proj_scale.to(torch.float32)
-            expert_output = torch.ops.fbgemm.f8f8bf16_rowwise(
-                activated_quantized,
-                self.down_proj[i].transpose(0, 1).contiguous(),
-                activated_scale,
-                down_proj_scale_float32[i].view(-1, 1).contiguous(),
-                use_fast_accum=True,
-            )
+            if _is_torch_xpu_available:
+                expert_output = torch._scaled_mm(
+                    activated_quantized,
+                    self.down_proj[i].transpose(0, 1).contiguous(),
+                    scale_a=activated_scale.unsqueeze(-1),
+                    scale_b=down_proj_scale_float32[i].view(-1, 1).contiguous().t(),
+                    out_dtype=hidden_states.dtype,
+                )
+            else:
+                expert_output = torch.ops.fbgemm.f8f8bf16_rowwise(
+                    activated_quantized,
+                    self.down_proj[i].transpose(0, 1).contiguous(),
+                    activated_scale,
+                    down_proj_scale_float32[i].view(-1, 1).contiguous(),
+                    use_fast_accum=True,
+                )
 
             next_states[i] = expert_output
         next_states = next_states.to(hidden_states.device)
         return next_states.view(-1, self.hidden_size)
 
 
-def _replace_with_fbgemm_fp8_linear(
-    model,
-    modules_to_not_convert=None,
-    current_key_name=None,
-    quantization_config=None,
-    has_been_replaced=False,
-    pre_quantized=False,
-    config=None,
-    tp_plan=None,
-):
-    """
-    Private method that wraps the recursion for module replacement.
-
-    Returns the converted model and a boolean that indicates if the conversion has been successful or not.
-    """
-
-    import re
-
-    if current_key_name is None:
-        current_key_name = []
-
-    for name, module in model.named_children():
-        current_key_name.append(name)
-
-        if (isinstance(module, nn.Linear)) and name not in modules_to_not_convert:
-            # Check if the current key is not in the `modules_to_not_convert`
-            current_key_name_str = ".".join(current_key_name)
-            if not any(
-                (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert
-            ):
-                with init_empty_weights(include_buffers=True):
-                    in_features = module.in_features
-                    out_features = module.out_features
-                    model._modules[name] = FbgemmFp8Linear(
-                        in_features,
-                        out_features,
-                        module.bias is not None,
-                    )
-                    has_been_replaced = True
-
-                    # Force requires grad to False to avoid unexpected errors
-                    model._modules[name].requires_grad_(False)
-                # set non persistent buffer outside of init_empty_weights
-                model._modules[name].input_scale_ub = torch.tensor(
-                    [quantization_config.activation_scale_ub],
-                    dtype=torch.float,
-                )
-        if module.__class__.__name__ == "Llama4TextExperts" and name not in modules_to_not_convert:
-            current_key_name_str = ".".join(current_key_name)
-            if not any(
-                (key + "." in current_key_name_str) or (key == current_key_name_str) for key in modules_to_not_convert
-            ):
-                with init_empty_weights(include_buffers=True):
-                    tp_plan[re.sub(r"\d+", "*", current_key_name_str + ".down_proj_scale")] = None
-                    model._modules[name] = FbgemmFp8Llama4TextExperts(
-                        config.text_config,
-                    )
-                model._modules[name].input_scale_ub = torch.tensor(
-                    [quantization_config.activation_scale_ub], dtype=torch.float
-                )
+@lru_cache(maxsize=1)
+def get_quantize_fp8_per_row():
+    if _is_torch_xpu_available:
+        from .hub_kernels import get_kernel
 
-        if len(list(module.children())) > 0:
-            _, has_been_replaced = _replace_with_fbgemm_fp8_linear(
-                module,
-                modules_to_not_convert,
-                current_key_name,
-                quantization_config,
-                has_been_replaced=has_been_replaced,
-                pre_quantized=pre_quantized,
-                config=config,
-                tp_plan=tp_plan,
-            )
-        # Remove the last key for recursion
-        current_key_name.pop(-1)
-    return model, has_been_replaced
+        return get_kernel("kernels-community/fp8-fbgemm").quantize_fp8_per_row
+    return torch.ops.fbgemm.quantize_fp8_per_row
 
 
 def replace_with_fbgemm_fp8_linear(
-    model,
-    modules_to_not_convert=None,
-    current_key_name=None,
-    quantization_config=None,
-    pre_quantized=False,
-    config=None,
-    tp_plan=None,
+    model, modules_to_not_convert: list[str] | None = None, quantization_config=None, pre_quantized=False, tp_plan=None
 ):
     """
     A helper function to replace all `torch.nn.Linear` modules by `FbgemmFp8Linear` modules.
     This will enable running your models using high performance fp8 kernel from FBGEMM library.
 
-    The function will be run recursively and replace all `torch.nn.Linear` modules except for the `lm_head` that should
-    be kept as a `torch.nn.Linear` module. The replacement is done under `init_empty_weights` context manager so no
-    CPU/GPU memory is required to run this function. Each weight will be quantized along the channel.
-
     Parameters:
         model (`torch.nn.Module`):
             Input model or `torch.nn.Module` as the function is run recursively.
-        modules_to_not_convert (`list[`str`]`, *optional*, defaults to `["lm_head"]`):
-            Names of the modules to not convert in `FP8Linear`. In practice we keep the `lm_head` in full precision
-            for numerical stability reasons.
-        current_key_name (`list[`str`]`, *optional*):
-            An array to track the current key of the recursion. This is used to check whether the current key (part of
-            it) is not in the list of modules to not convert (for instances modules that are offloaded to `cpu` or
-            `disk`).
+        modules_to_not_convert (`list[`str`]`, *optional*, defaults to `None`):
+            Names of the modules to not convert. In practice we keep the `lm_head` in full precision for numerical stability reasons.
+        quantization_config (`FbgemmFp8Config`):
+            The quantization config object that contains the quantization parameters.
+        pre_quantized (`book`, defaults to `False`):
+            Whether the model is pre-quantized or not
     """
+    global quantize_fp8_per_row
+    quantize_fp8_per_row = get_quantize_fp8_per_row()
+
+    has_been_replaced = False
+    module_kwargs = {} if pre_quantized else {"dtype": None}
+
+    for module_name, module in model.named_modules():
+        if not should_convert_module(module_name, modules_to_not_convert):
+            continue
+
+        new_module = None
+        with init_empty_weights(include_buffers=True):
+            if module.__class__.__name__ == "Llama4TextExperts":
+                # TODO: make sure tp works later
+                # if tp_plan is not None:
+                #     tp_key = re.sub(r"\d+", "*", f"{module_name}.down_proj_scale")
+                #     tp_plan[tp_key] = None
+                text_config = getattr(model.config, "text_config", model.config)
+                new_module = FbgemmFp8Llama4TextExperts(text_config or model.config)
+            elif isinstance(module, nn.Linear):
+                new_module = FbgemmFp8Linear(
+                    module.in_features,
+                    module.out_features,
+                    module.bias is not None,
+                    **module_kwargs,
+                )
+                new_module.requires_grad_(False)
+
+        if new_module is None:
+            continue
+
+        if hasattr(new_module, "input_scale_ub"):
+            new_module.input_scale_ub = torch.tensor(
+                [quantization_config.activation_scale_ub],
+                dtype=torch.float,
+            )
+
+        model.set_submodule(module_name, new_module)
+        has_been_replaced = True
 
-    modules_to_not_convert = ["lm_head"] if modules_to_not_convert is None else modules_to_not_convert
-
-    if quantization_config.modules_to_not_convert is not None:
-        modules_to_not_convert.extend(quantization_config.modules_to_not_convert)
-    modules_to_not_convert = list(set(modules_to_not_convert))
-    model, has_been_replaced = _replace_with_fbgemm_fp8_linear(
-        model,
-        modules_to_not_convert,
-        current_key_name,
-        quantization_config,
-        pre_quantized=pre_quantized,
-        config=config,
-        tp_plan=tp_plan,
-    )
     if not has_been_replaced:
         logger.warning(
             "You are loading your model using FP8 quantization but no linear modules were found in your model."