refactor pt loading (#4532)

fastdeploy/model_executor/layers/quantization/block_wise_fp8.py

@@ -22,10 +22,15 @@ import fastdeploy
 from fastdeploy import envs
 from fastdeploy.model_executor.layers.linear import (
     MergedColumnParallelLinear,
+    MergedReplicatedLinear,
     QKVParallelLinear,
 )
 from fastdeploy.model_executor.layers.moe import FusedMoE
-from fastdeploy.model_executor.utils import TensorTracker, set_weight_attrs
+from fastdeploy.model_executor.utils import (
+    TensorTracker,
+    process_weight_transpose,
+    set_weight_attrs,
+)
 
 from ..utils import get_tensor, per_block_cast_to_fp8
 from .quant_base import QuantConfigBase, QuantMethodBase
@@ -90,51 +95,66 @@ class BlockWiseFP8LinearMethod(QuantMethodBase):
 
     def create_weights(self, layer, **extra_weight_attrs):
         # TODO(bukejiyu): remove v1 loader check when v0 loader is removed
+        self.model_format = extra_weight_attrs.get("model_format")
         if self.quant_config.is_checkpoint_bf16 and layer.fd_config.load_config.load_choices == "default_v1":
+            weight_shape = layer.weight_shape[::-1] if self.model_format == "torch" else layer.weight_shape
             layer.weight = layer.create_parameter(
-                shape=layer.weight_shape,
+                shape=weight_shape,
                 dtype=layer.weight_dtype,
                 is_bias=False,
                 default_initializer=paddle.nn.initializer.Constant(0),
             )
-            extra_weight_attrs["weight_need_transpose"] = extra_weight_attrs.get("model_format") == "torch"
             quant_attrs = extra_weight_attrs
-            if isinstance(layer, MergedColumnParallelLinear) or isinstance(layer, QKVParallelLinear):
+            if (
+                isinstance(layer, MergedColumnParallelLinear)
+                or isinstance(layer, QKVParallelLinear)
+                or isinstance(layer, MergedReplicatedLinear)
+            ):
+                tensor_output_dim = (self.model_format == "torch") ^ quant_attrs.get("output_dim", True)
                 quant_attrs = {
                     **extra_weight_attrs,
-                    "tensor_track": TensorTracker(
-                        shape=layer.weight_shape, output_dim=extra_weight_attrs.get("output_dim")
-                    ),
+                    "tensor_track": TensorTracker(shape=weight_shape, output_dim=tensor_output_dim),
                 }
+            if self.model_format == "torch" and "output_dim" in quant_attrs:
+                quant_attrs["output_dim"] = not quant_attrs["output_dim"]
             set_weight_attrs(
                 layer.weight,
                 quant_attrs,
             )
         else:
             layer.weight_shape.reverse()
+            weight_scale_inv_shape = [
+                (layer.weight_shape[0] + self.quant_config.weight_block_size[0] - 1)
+                // self.quant_config.weight_block_size[0],
+                (layer.weight_shape[1] + self.quant_config.weight_block_size[1] - 1)
+                // self.quant_config.weight_block_size[1],
+            ]
+
+            if self.model_format != "torch" and layer.fd_config.load_config.load_choices == "default_v1":
+                weight_shape = layer.weight_shape[::-1]
+                weight_scale_inv_shape = weight_scale_inv_shape[::-1]
+            else:
+                # v0 loader or torch model format
+                weight_shape = layer.weight_shape
+                weight_scale_inv_shape = weight_scale_inv_shape
+                extra_weight_attrs["output_dim"] = (
+                    not extra_weight_attrs["output_dim"] if extra_weight_attrs["output_dim"] is not None else None
+                )
+
             layer.weight_dtype = "float8_e4m3fn"
             layer.weight = layer.create_parameter(
-                shape=layer.weight_shape,
+                shape=weight_shape,
                 dtype=layer.weight_dtype,
                 is_bias=False,
                 default_initializer=paddle.nn.initializer.Constant(0),
             )
 
             layer.weight_scale_inv = layer.create_parameter(
-                shape=[
-                    (layer.weight_shape[0] + self.quant_config.weight_block_size[0] - 1)
-                    // self.quant_config.weight_block_size[0],
-                    (layer.weight_shape[1] + self.quant_config.weight_block_size[1] - 1)
-                    // self.quant_config.weight_block_size[1],
-                ],
+                shape=weight_scale_inv_shape,
                 dtype="float32",
                 is_bias=False,
             )
-            extra_weight_attrs["output_dim"] = (
-                not extra_weight_attrs["output_dim"] if extra_weight_attrs["output_dim"] is not None else None
-            )
 
-            extra_weight_attrs["weight_need_transpose"] = not extra_weight_attrs.get("model_format") == "torch"
             set_weight_attrs(
                 layer.weight,
                 extra_weight_attrs,
@@ -148,31 +168,41 @@ class BlockWiseFP8LinearMethod(QuantMethodBase):
             )
 
     def process_weights_after_loading(self, layer) -> None:
-        if not self.quant_config.is_checkpoint_bf16:
-            return
-        weight_tensor = layer.weight.transpose([1, 0])
-        quanted_weight_tensor, weight_block_scale_tensor = per_block_cast_to_fp8(weight_tensor)
+        def _process_quantize():
+            weight_tensor = layer.weight.transpose([1, 0])
+            quanted_weight_tensor, weight_block_scale_tensor = per_block_cast_to_fp8(weight_tensor)
 
-        if hasattr(layer.weight, "tensor_track"):
-            layer.weight.tensor_track = None
-        layer.weight.value().get_tensor()._clear()
-        del layer.weight
+            if hasattr(layer.weight, "tensor_track"):
+                layer.weight.tensor_track = None
+            layer.weight.value().get_tensor()._clear()
+            del layer.weight
 
-        layer.weight = layer.create_parameter(
-            shape=quanted_weight_tensor.shape,
-            dtype="float8_e4m3fn",
-            is_bias=False,
-            default_initializer=paddle.nn.initializer.Constant(0),
-        )
-        layer.weight_scale_inv = layer.create_parameter(
-            shape=weight_block_scale_tensor.shape,
-            dtype="float32",
-            is_bias=False,
-            default_initializer=paddle.nn.initializer.Constant(0),
-        )
+            layer.weight = layer.create_parameter(
+                shape=quanted_weight_tensor.shape,
+                dtype="float8_e4m3fn",
+                is_bias=False,
+                default_initializer=paddle.nn.initializer.Constant(0),
+            )
+            layer.weight_scale_inv = layer.create_parameter(
+                shape=weight_block_scale_tensor.shape,
+                dtype="float32",
+                is_bias=False,
+                default_initializer=paddle.nn.initializer.Constant(0),
+            )
 
-        layer.weight.copy_(quanted_weight_tensor, False)
-        layer.weight_scale_inv.copy_(weight_block_scale_tensor, False)
+            layer.weight.copy_(quanted_weight_tensor, False)
+            layer.weight_scale_inv.copy_(weight_block_scale_tensor, False)
+
+        if self.quant_config.is_checkpoint_bf16:
+            if self.model_format == "torch":
+                process_weight_transpose(layer, "weight")
+            _process_quantize()
+        else:
+            if self.model_format != "torch":
+                process_weight_transpose(layer, "weight")
+                process_weight_transpose(layer, "weight_scale_inv")
+            else:
+                return
 
     def process_loaded_weights(self, layer, weights) -> None:
         weight_tensor = weights.transpose([1, 0])

fastdeploy/model_executor/layers/quantization/mix_quant.py

@@ -55,7 +55,6 @@ class MixQuantConfig(QuantConfigBase):
         self.quant_round_type = 0
         self.is_permuted = is_permuted
         self.is_checkpoint_bf16 = not is_quantized
-        self.is_quantized = is_quantized
         self.hadamard_block_size = hadamard_block_size
 
     def name(self) -> str:
@@ -83,7 +82,7 @@ class MixQuantConfig(QuantConfigBase):
                     .from_config(
                         {
                             "is_permuted": self.is_permuted,
-                            "is_quantized": self.is_quantized,
+                            "is_quantized": not self.is_checkpoint_bf16,
                             "hadamard_block_size": self.hadamard_block_size,
                         }
                     )
@@ -95,7 +94,7 @@ class MixQuantConfig(QuantConfigBase):
                     .from_config(
                         {
                             "is_permuted": self.is_permuted,
-                            "is_quantized": self.is_quantized,
+                            "is_quantized": not self.is_checkpoint_bf16,
                             "hadamard_block_size": self.hadamard_block_size,
                         }
                     )
@@ -113,6 +112,6 @@ class MixQuantConfig(QuantConfigBase):
         else:
             return (
                 get_quantization_config(self.dense_quant_type)
-                .from_config({"is_quantized": self.is_quantized})
+                .from_config({"is_quantized": not self.is_checkpoint_bf16})
                 .get_quant_method(layer)
             )

fastdeploy/model_executor/layers/quantization/weight_only.py

@@ -28,7 +28,12 @@ from fastdeploy.model_executor.layers.linear import (
     MergedReplicatedLinear,
     QKVParallelLinear,
 )
-from fastdeploy.model_executor.utils import TensorTracker, free_tensor, set_weight_attrs
+from fastdeploy.model_executor.utils import (
+    TensorTracker,
+    free_tensor,
+    process_weight_transpose,
+    set_weight_attrs,
+)
 from fastdeploy.platforms import current_platform
 
 if current_platform.is_xpu():
@@ -231,26 +236,33 @@ class WeightOnlyLinearMethod(QuantMethodBase):
 
     def create_weights(self, layer, **extra_weight_attrs):
         # TODO(bukejiyu): remove v1 loader check when v0 loader is removed
+        self.model_format = extra_weight_attrs.get("model_format")
         if self.quant_config.is_checkpoint_bf16 and layer.fd_config.load_config.load_choices == "default_v1":
+            weight_shape = layer.weight_shape[::-1] if self.model_format == "torch" else layer.weight_shape
             layer.weight = layer.create_parameter(
-                shape=layer.weight_shape,
+                shape=weight_shape,
                 dtype=layer.weight_dtype,
                 is_bias=False,
                 default_initializer=paddle.nn.initializer.Constant(0),
             )
-            extra_weight_attrs["weight_need_transpose"] = extra_weight_attrs.get("model_format") == "torch"
+
             quant_attrs = extra_weight_attrs
+
             if (
                 isinstance(layer, MergedColumnParallelLinear)
                 or isinstance(layer, QKVParallelLinear)
                 or isinstance(layer, MergedReplicatedLinear)
             ):
+                # Only MergedReplicatedLinear uses the default outdim.
+                tensor_output_dim = (self.model_format == "torch") ^ quant_attrs.get("output_dim", True)
                 quant_attrs = {
-                    **extra_weight_attrs,
-                    "tensor_track": TensorTracker(
-                        shape=layer.weight_shape, output_dim=extra_weight_attrs.get("output_dim", True)
-                    ),
+                    **quant_attrs,
+                    "tensor_track": TensorTracker(shape=weight_shape, output_dim=tensor_output_dim),
                 }
+
+            if self.model_format == "torch" and "output_dim" in quant_attrs:
+                quant_attrs["output_dim"] = not quant_attrs["output_dim"]
+
             set_weight_attrs(
                 layer.weight,
                 quant_attrs,
@@ -279,16 +291,11 @@ class WeightOnlyLinearMethod(QuantMethodBase):
                 default_initializer=paddle.nn.initializer.Constant(0),
             )
 
-            output_dim = extra_weight_attrs.get("output_dim")
-            output_dim = not output_dim
-            weight_loader = extra_weight_attrs.get("weight_loader")
+            if "output_dim" in extra_weight_attrs:
+                extra_weight_attrs["output_dim"] = not extra_weight_attrs["output_dim"]
             set_weight_attrs(
                 layer.weight,
-                {
-                    "weight_loader": weight_loader,
-                    "output_dim": output_dim,
-                    "weight_need_transpose": not extra_weight_attrs.get("model_format") == "torch",
-                },
+                extra_weight_attrs,
             )
 
             layer.weight_scale = layer.create_parameter(
@@ -299,47 +306,49 @@ class WeightOnlyLinearMethod(QuantMethodBase):
 
             set_weight_attrs(
                 layer.weight_scale,
-                {
-                    "weight_loader": weight_loader,
-                    "output_dim": output_dim,
-                },
+                extra_weight_attrs,
             )
 
     def process_weights_after_loading(self, layer) -> None:
-        if not self.quant_config.is_checkpoint_bf16:
-            return
-        if isinstance(self, MacheteWeightOnlyLinearMethod):
+        def _process_quantize():
+            if isinstance(self, MacheteWeightOnlyLinearMethod):
+                # Using group scale for machete
+                quanted_weight_tensor, weight_scale_tensor = machete_quantize_and_pack(
+                    w=layer.weight,
+                    atype=layer._dtype,
+                    quant_type="uint4b8" if self.quant_config.name() == "wint4" else "uint8b128",
+                    group_size=self.quant_config.group_size,
+                )
+            else:
+                quanted_weight_tensor, weight_scale_tensor = weight_quantize(
+                    layer.weight,
+                    algo=self.quant_config.algo,
+                    arch=self.quant_config.weight_only_linear_arch,
+                )
 
-            # Using group scale for machete
-            quanted_weight_tensor, weight_scale_tensor = machete_quantize_and_pack(
-                w=layer.weight,
-                atype=layer._dtype,
-                quant_type="uint4b8" if self.quant_config.name() == "wint4" else "uint8b128",
-                group_size=self.quant_config.group_size,
+            free_tensor(layer.weight)
+
+            layer.weight = layer.create_parameter(
+                shape=quanted_weight_tensor.shape,
+                dtype="int8" if not isinstance(self, MacheteWeightOnlyLinearMethod) else "int32",
+                is_bias=False,
+                default_initializer=paddle.nn.initializer.Constant(0),
             )
+            layer.weight_scale = layer.create_parameter(
+                shape=weight_scale_tensor.shape,
+                dtype=layer._dtype,
+                is_bias=False,
+                default_initializer=paddle.nn.initializer.Constant(0),
+            )
+            layer.weight.copy_(quanted_weight_tensor, False)
+            layer.weight_scale.copy_(weight_scale_tensor, False)
+
+        if self.quant_config.is_checkpoint_bf16:
+            if self.model_format == "torch":
+                process_weight_transpose(layer, "weight")
+            _process_quantize()
         else:
-            quanted_weight_tensor, weight_scale_tensor = weight_quantize(
-                layer.weight,
-                algo=self.quant_config.algo,
-                arch=self.quant_config.weight_only_linear_arch,
-            )
-
-        free_tensor(layer.weight)
-
-        layer.weight = layer.create_parameter(
-            shape=quanted_weight_tensor.shape,
-            dtype="int8" if not isinstance(self, MacheteWeightOnlyLinearMethod) else "int32",
-            is_bias=False,
-            default_initializer=paddle.nn.initializer.Constant(0),
-        )
-        layer.weight_scale = layer.create_parameter(
-            shape=weight_scale_tensor.shape,
-            dtype=layer._dtype,
-            is_bias=False,
-            default_initializer=paddle.nn.initializer.Constant(0),
-        )
-        layer.weight.copy_(quanted_weight_tensor, False)
-        layer.weight_scale.copy_(weight_scale_tensor, False)
+            return
 
     @abstractmethod
     def process_loaded_weights(self, layer, weights) -> None:

fastdeploy/model_executor/layers/quantization/wfp8afp8.py

@@ -21,6 +21,7 @@ import paddle
 
 from fastdeploy.model_executor.layers.linear import (
     MergedColumnParallelLinear,
+    MergedReplicatedLinear,
     QKVParallelLinear,
 )
 from fastdeploy.model_executor.layers.moe import FusedMoE
@@ -33,7 +34,11 @@ from fastdeploy.model_executor.layers.quantization.quant_base import (
     QuantMethodBase,
 )
 from fastdeploy.model_executor.layers.utils import per_token_cast_to_fp8
-from fastdeploy.model_executor.utils import TensorTracker, set_weight_attrs
+from fastdeploy.model_executor.utils import (
+    TensorTracker,
+    process_weight_transpose,
+    set_weight_attrs,
+)
 
 
 class WFP8AFP8Config(QuantConfigBase):
@@ -101,22 +106,28 @@ class WFP8AFP8LinearMethod(QuantMethodBase):
                 (weight_shape[i] + weight_block_size[i] - 1) // weight_block_size[i] if weight_block_size[i] > 0 else 1
             )
         scale_shape = scale_shape[::-1]
+        self.model_format = extra_weight_attrs.get("model_format")
         if self.quant_config.is_checkpoint_bf16 and layer.fd_config.load_config.load_choices == "default_v1":
+            weight_shape = weight_shape[::-1] if self.model_format == "torch" else weight_shape
             layer.weight = layer.create_parameter(
                 shape=weight_shape,
                 dtype=layer.weight_dtype,
                 is_bias=False,
                 default_initializer=paddle.nn.initializer.Constant(0),
             )
-            extra_weight_attrs["weight_need_transpose"] = extra_weight_attrs.get("model_format") == "torch"
             quant_attrs = extra_weight_attrs
-            if isinstance(layer, MergedColumnParallelLinear) or isinstance(layer, QKVParallelLinear):
+            if (
+                isinstance(layer, MergedColumnParallelLinear)
+                or isinstance(layer, QKVParallelLinear)
+                or isinstance(layer, MergedReplicatedLinear)
+            ):
+                tensor_output_dim = (self.model_format == "torch") ^ quant_attrs.get("output_dim", True)
                 quant_attrs = {
                     **extra_weight_attrs,
-                    "tensor_track": TensorTracker(
-                        shape=layer.weight_shape, output_dim=extra_weight_attrs.get("output_dim")
-                    ),
+                    "tensor_track": TensorTracker(shape=weight_shape, output_dim=tensor_output_dim),
                 }
+            if self.model_format == "torch" and "output_dim" in quant_attrs:
+                quant_attrs["output_dim"] = not quant_attrs["output_dim"]
             set_weight_attrs(
                 layer.weight,
                 quant_attrs,
@@ -142,30 +153,39 @@ class WFP8AFP8LinearMethod(QuantMethodBase):
     def process_weights_after_loading(self, layer) -> None:
         if not self.quant_config.is_checkpoint_bf16:
             return
-        weight_tensor = layer.weight.transpose([1, 0]).contiguous()
-        assert self.quant_config.weight_block_size == [-1, 1]
-        qweight, weight_scale = per_token_cast_to_fp8(weight_tensor)
 
-        if hasattr(layer.weight, "tensor_track"):
-            layer.weight.tensor_track = None
-        layer.weight.value().get_tensor()._clear()
-        del layer.weight
+        def _process_quantize():
+            weight_tensor = layer.weight.transpose([1, 0]).contiguous()
+            assert self.quant_config.weight_block_size == [-1, 1]
+            qweight, weight_scale = per_token_cast_to_fp8(weight_tensor)
 
-        layer.weight = layer.create_parameter(
-            shape=qweight.shape,
-            dtype="float8_e4m3fn",
-            is_bias=False,
-            default_initializer=paddle.nn.initializer.Constant(0),
-        )
-        layer.weight_scale = layer.create_parameter(
-            shape=weight_scale.shape,
-            dtype="float32",
-            is_bias=False,
-            default_initializer=paddle.nn.initializer.Constant(0),
-        )
+            if hasattr(layer.weight, "tensor_track"):
+                layer.weight.tensor_track = None
+            layer.weight.value().get_tensor()._clear()
+            del layer.weight
 
-        layer.weight.copy_(qweight, False)
-        layer.weight_scale.copy_(weight_scale, False)
+            layer.weight = layer.create_parameter(
+                shape=qweight.shape,
+                dtype="float8_e4m3fn",
+                is_bias=False,
+                default_initializer=paddle.nn.initializer.Constant(0),
+            )
+            layer.weight_scale = layer.create_parameter(
+                shape=weight_scale.shape,
+                dtype="float32",
+                is_bias=False,
+                default_initializer=paddle.nn.initializer.Constant(0),
+            )
+
+            layer.weight.copy_(qweight, False)
+            layer.weight_scale.copy_(weight_scale, False)
+
+        if self.quant_config.is_checkpoint_bf16:
+            if self.model_format == "torch":
+                process_weight_transpose(layer, "weight")
+            _process_quantize()
+        else:
+            return
 
     def process_loaded_weights(self, layer, weights) -> None:
         """ """