refactor pt loading (#4532)

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: