support inflight quant

Author: bukejiyu

fastdeploy/config.py

@@ -665,6 +665,7 @@ class LoadChoices(str, Enum):
     DEFAULT = "default"
     # only support qwen3-bf16 now
     DEFAULT_V1 = "default_v1"
+    INFLIGHT_QUANT = "inflight_quant"
 
 
 class LoadConfig:
@@ -684,6 +685,7 @@ def __init__(
         args,
     ):
         self.load_choices: Union[str, LoadChoices] = LoadChoices.DEFAULT.value
+        self.is_inflight_quant = False
         self.use_fastsafetensor = int(envs.FD_USE_FASTSAFETENSOR) == 1
         self.dynamic_load_weight: bool = False
         self.load_strategy: Optional[Literal["ipc", "ipc_snapshot"]] = None

fastdeploy/model_executor/layers/linear.py

@@ -23,8 +23,10 @@
 from fastdeploy.distributed.communication import tensor_model_parallel_all_reduce
 from fastdeploy.model_executor.layers.quantization.quant_base import QuantMethodBase
 from fastdeploy.model_executor.models.utils import (
-    default_weight_loader,
+    default_load_weights_into_param,
+    default_weights_processor,
     set_weight_attrs,
+    slice_fn,
 )
 from fastdeploy.platforms import current_platform
 
@@ -37,24 +39,29 @@ class UnquantizedLinearMethod(QuantMethodBase):
     def create_weights(self, layer: nn.Layer, **extra_weight_attrs):
         """
         extra_weight_attrs is a dictionary that may include parameters like:
-        - split_axis: specifies which axis to split the weight tensor on (for distributed weight partitioning)
-        - output_dim: determines whether the split is applied along the output dimension (rows) or input dimension (columns)
-        - weight_loader: a callable or method responsible for loading the weight data
+        - weights_processor: a callable or method responsible for loading the weight data
         """
         layer.weight = layer.create_parameter(
             shape=layer.weight_shape,
             dtype=layer.weight_dtype,
             is_bias=False,
             default_initializer=paddle.nn.initializer.Constant(0),
         )
+        split_axis = extra_weight_attrs.get("split_axis")
+        if hasattr(layer, "nranks") and layer.nranks > 0:
+            _set_var_distributed(layer.weight, split_axis=split_axis)
         set_weight_attrs(
             layer.weight,
-            {"weight_loader": extra_weight_attrs.get("weight_loader", default_weight_loader(layer.fd_config))},
+            {
+                **extra_weight_attrs,
+                "weights_processor": extra_weight_attrs.get(
+                    "weights_processor", default_weights_processor(layer.fd_config)
+                ),
+                "load_weights_into_param": extra_weight_attrs.get(
+                    "load_weights_into_param", default_load_weights_into_param()
+                ),
+            },
         )
-        if hasattr(layer, "nranks") and layer.nranks > 0:
-            split_axis = extra_weight_attrs.get("split_axis")
-            _set_var_distributed(layer.weight, split_axis=split_axis)
-            set_weight_attrs(layer.weight, {"output_dim": extra_weight_attrs.get("output_dim")})
 
     def process_loaded_weights(self, layer, weights) -> None:
         # mlp.gate.weight is precision-sensitive, so we cast it to float32 for computation
@@ -157,6 +164,7 @@ def __init__(
                 is_bias=True,
             )
 
+        self.is_quantized = fd_config.model_config.is_quantized
         # smooth quant
         self.linear_shift = None
         self.linear_smooth = None
@@ -274,9 +282,17 @@ def __init__(
         assert self.quant_method is not None
         self.quant_method.create_weights(
             self,
-            weight_loader=(
-                self.weight_loader if hasattr(self, "weight_loader") else default_weight_loader(self.fd_config)
+            weights_processor=(
+                self.weights_processor
+                if hasattr(self, "weights_processor")
+                else default_weights_processor(self.fd_config)
+            ),
+            load_weights_into_param=(
+                self.load_weights_into_param
+                if hasattr(self, "load_weights_into_param")
+                else default_load_weights_into_param()
             ),
+            inflight_quant=fd_config.quant_config and not skip_quant,
         )
 
 
@@ -335,16 +351,23 @@ def __init__(
             self,
             split_axis=1,
             output_dim=True,
-            weight_loader=(
-                self.weight_loader if hasattr(self, "weight_loader") else default_weight_loader(self.fd_config)
+            weights_processor=(
+                self.weights_processor
+                if hasattr(self, "weights_processor")
+                else default_weights_processor(self.fd_config)
             ),
+            load_weights_into_param=(
+                self.load_weights_into_param
+                if hasattr(self, "load_weights_into_param")
+                else default_load_weights_into_param()
+            ),
+            inflight_quant=fd_config.quant_config and not skip_quant,
         )
 
-        if self.with_bias:
-            if self.nranks > 0:
+        if self.nranks > 0:
+            if self.with_bias:
                 # col parallel
                 _set_var_distributed(self.bias, split_axis=1)
-                set_weight_attrs(self.bias, {"output_dim": True})
 
 
 class MergedColumnParallelLinear(ColumnParallelLinear):
@@ -397,31 +420,33 @@ def __init__(
             skip_quant=skip_quant,
         )
 
-    def weight_loader(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
+    def load_weights_into_param(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
+        assert loaded_shard_id in ["gate", "up"]
+        output_dim = getattr(param, "output_dim", None)
+        if loaded_shard_id == "gate":
+            param = slice_fn(param, output_dim, start=0, end=self.output_size // 2)
+        elif loaded_shard_id == "up":
+            param = slice_fn(param, output_dim, start=self.output_size // 2, end=self.output_size)
+        assert param.shape == loaded_weight.shape, (
+            f" Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({param.shape})"
+        )
+        param.copy_(loaded_weight, False)
+
+    def weights_processor(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
         # 1.fused gate_up in disk
         # 2.split gate up
         assert loaded_shard_id in ["gate", "up"]
         output_dim = getattr(param, "output_dim", None)
         # Tensor parallelism splits the weight along the output_dim
-        if output_dim is not None:
+        if output_dim is not None and self.nranks > 1:
             dim = -1
             size = loaded_weight.get_shape()[dim]
             block_size = size // self.nranks
             shard_offset = self.local_rank * block_size
             shard_size = (self.local_rank + 1) * block_size
-            loaded_weight = loaded_weight[..., shard_offset:shard_size]
-
+            loaded_weight = slice_fn(loaded_weight, output_dim, shard_offset, shard_size)
         loaded_weight = get_tensor(loaded_weight)
-
-        if loaded_shard_id == "gate":
-            param = param[:, : self.output_size // 2]
-        elif loaded_shard_id == "up":
-            param = param[:, self.output_size // 2 :]
-
-        assert param.shape == loaded_weight.shape, (
-            f" Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({param.shape})"
-        )
-        param.copy_(loaded_weight, False)
+        yield loaded_weight
 
     def load_state_dict(self, state_dict: dict):
         """
@@ -491,33 +516,44 @@ def __init__(self, fd_config, prefix, with_bias=False, add_bias=True):
             add_bias=add_bias,
         )
 
-    def weight_loader(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
+    def weights_processor(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
         # 1.fused qkv in disk
         # 2.split q k v
         assert loaded_shard_id in ["q", "k", "v"]
         output_dim = getattr(param, "output_dim", None)
         # Tensor parallelism splits the weight along the output_dim
-        if output_dim is not None:
+        if output_dim is not None and self.nranks > 1:
             dim = -1
             size = loaded_weight.get_shape()[dim]
             block_size = size // self.nranks
             shard_offset = self.local_rank * block_size
             shard_size = (self.local_rank + 1) * block_size
-            loaded_weight = loaded_weight[..., shard_offset:shard_size]
+            loaded_weight = slice_fn(loaded_weight, output_dim, shard_offset, shard_size)
 
         loaded_weight = get_tensor(loaded_weight)
+        yield loaded_weight
 
+    def load_weights_into_param(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
+        assert loaded_shard_id in ["q", "k", "v"]
+        output_dim = getattr(param, "output_dim", None)
         if loaded_shard_id == "q":
-            param = param[:, : self.num_heads_per_rank * self.head_dim]
+            param = slice_fn(param, output_dim, 0, self.num_heads_per_rank * self.head_dim)
+
         elif loaded_shard_id == "k":
-            param = param[
-                :,
-                self.num_heads_per_rank
-                * self.head_dim : (self.num_heads_per_rank + self.kv_num_heads_per_rank)
-                * self.head_dim,
-            ]
+            param = slice_fn(
+                param,
+                output_dim,
+                self.num_heads_per_rank * self.head_dim,
+                (self.num_heads_per_rank + self.kv_num_heads_per_rank) * self.head_dim,
+            )
+
         elif loaded_shard_id == "v":
-            param = param[:, (self.num_heads_per_rank + self.kv_num_heads_per_rank) * self.head_dim :]
+            param = slice_fn(
+                param,
+                output_dim,
+                (self.num_heads_per_rank + self.kv_num_heads_per_rank) * self.head_dim,
+                (self.num_heads_per_rank + 2 * self.kv_num_heads_per_rank) * self.head_dim,
+            )
 
         assert param.shape == loaded_weight.shape, (
             f" Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({param.shape})"
@@ -665,19 +701,30 @@ def __init__(
             self,
             split_axis=0,
             output_dim=False,
-            weight_loader=(
-                self.weight_loader if hasattr(self, "weight_loader") else default_weight_loader(self.fd_config)
+            weights_processor=(
+                self.weights_processor
+                if hasattr(self, "weights_processor")
+                else default_weights_processor(self.fd_config)
             ),
+            load_weights_into_param=(
+                self.load_weights_into_param
+                if hasattr(self, "load_weights_into_param")
+                else default_load_weights_into_param()
+            ),
+            inflight_quant=fd_config.quant_config and not skip_quant,
         )
 
-        if self.with_bias:
-            _set_var_distributed(self.bias, split_axis=0)
-            set_weight_attrs(
-                self.bias,
-                {
-                    "output_dim": False,
-                },
-            )
+        if self.nranks > 0:
+            if self.with_bias:
+                # col parallel
+                _set_var_distributed(self.bias, split_axis=0)
+                set_weight_attrs(
+                    self.bias,
+                    {
+                        "output_dim": False,
+                    },
+                )
+
         self.reduce_results = reduce_results
 
     def forward_cuda(self, x: paddle.Tensor) -> paddle.Tensor:

fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py

@@ -185,9 +185,11 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs):
         if current_platform.is_cuda():
             self.up_gate_proj_weight_shape = [layer.num_experts, layer.hidden_size, layer.moe_intermediate_size * 2]
             self.down_proj_weight_shape = [layer.num_experts, layer.moe_intermediate_size, layer.hidden_size]
+            extra_weight_attrs = {**extra_weight_attrs, "SHARD_ID_TO_SHARDED_DIM": {"gate": 1, "down": 0, "up": 1}}
         else:
             self.up_gate_proj_weight_shape = [layer.num_experts, layer.moe_intermediate_size * 2, layer.hidden_size]
             self.down_proj_weight_shape = [layer.num_experts, layer.hidden_size, layer.moe_intermediate_size]
+            extra_weight_attrs = {**extra_weight_attrs, "SHARD_ID_TO_SHARDED_DIM": {"gate": 0, "down": 1, "up": 0}}
 
         layer.up_gate_proj_weight = layer.create_parameter(
             shape=self.up_gate_proj_weight_shape,