qwen3/qwen3moe support dyn quant

Author: bukejiyu

fastdeploy/model_executor/layers/embeddings.py

@@ -22,7 +22,7 @@
 from paddle.distributed import fleet
 
 from fastdeploy.config import FDConfig
-from fastdeploy.model_executor.models.utils import set_weight_attrs
+from fastdeploy.model_executor.utils import set_weight_attrs
 
 from .utils import get_tensor
 

fastdeploy/model_executor/layers/linear.py

@@ -23,9 +23,10 @@
 from fastdeploy.config import FDConfig
 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 (
+from fastdeploy.model_executor.utils import (
     default_weight_loader,
     set_weight_attrs,
+    slice_fn,
 )
 from fastdeploy.platforms import current_platform
 
@@ -38,6 +39,7 @@ 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: axis along which to split the tensor in a distributed environment
         - 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
         """
@@ -47,12 +49,16 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs):
             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,
+                "weight_loader": extra_weight_attrs.get("weight_loader", default_weight_loader(layer.fd_config)),
+            },
         )
-        if hasattr(layer, "nranks") and layer.nranks > 1:
-            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
@@ -334,7 +340,6 @@ def __init__(
             ),
         )
         if self.nranks > 0:
-            _set_var_distributed(self.weight, split_axis=1)
             if self.with_bias:
                 # col parallel
                 _set_var_distributed(self.bias, split_axis=1)
@@ -393,44 +398,47 @@ def __init__(
         )
 
     def weight_loader(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
+        output_dim = getattr(param, "output_dim", None)
+        shard_dim = -1 if output_dim else 0
+        output_size = param.shape[shard_dim] // 2
         if loaded_shard_id is None:
             # Loaded weight is already fused on disk.
-            if self.nranks != 1:
-                shard_offsets = [
-                    # (shard_id, shard_offset, shard_size)
-                    ("gate", 0, self.output_size * self.nranks // 2),
-                    ("up", self.output_size * self.nranks // 2, self.output_size * self.nranks // 2),
-                ]
-                for shard_id, shard_offset, shard_size in shard_offsets:
-                    loaded_weight_shard = loaded_weight[..., shard_offset : shard_offset + shard_size]
-                    self.weight_loader(param, loaded_weight_shard, shard_id)
-            else:
-                loaded_weight = get_tensor(loaded_weight)
-                param.copy_(loaded_weight, False)
+            shard_offsets = [
+                # (shard_id, shard_offset, shard_size)
+                ("gate", 0, output_size * self.nranks // 2),
+                ("up", output_size * self.nranks // 2, output_size * self.nranks // 2),
+            ]
+            for shard_id, shard_offset, shard_size in shard_offsets:
+                loaded_weight_shard = slice_fn(
+                    loaded_weight, output_dim, start=shard_offset, end=shard_offset + shard_size
+                )
+                self.weight_loader(param, loaded_weight_shard, shard_id)
         else:
-            # 1.fused gate_up in disk
-            # 2.split gate up
+            # 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:
-                dim = -1
+            if self.nranks != 1:
+                dim = -1 if output_dim else 0
                 if isinstance(loaded_weight, np.ndarray):
                     size = loaded_weight.shape[dim]
                 else:
                     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, start=shard_offset, end=shard_size)
 
             loaded_weight = get_tensor(loaded_weight)
-
+            if not param._is_initialized():
+                param.initialize()
             if loaded_shard_id == "gate":
-                param = param[:, : self.output_size // 2]
-            elif loaded_shard_id == "up":
-                param = param[:, self.output_size // 2 :]
-
+                param_shard_offset = 0
+            else:
+                # loaded_shard_id == "up"
+                param_shard_offset = output_size
+            if hasattr(param, "tensor_track"):
+                param.tensor_track.mark(start=param_shard_offset, end=param_shard_offset + output_size)
+            param = slice_fn(param, output_dim, start=param_shard_offset, end=param_shard_offset + output_size)
             assert param.shape == loaded_weight.shape, (
                 f" Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({param.shape})"
             )
@@ -505,53 +513,54 @@ def __init__(self, fd_config, prefix, with_bias=False, add_bias=True):
         )
 
     def weight_loader(self, param, loaded_weight, loaded_shard_id: Optional[str] = None):
+        output_dim = getattr(param, "output_dim", None)
+        head_dim = param.shape[output_dim] // (self.num_heads_per_rank + 2 * self.kv_num_heads_per_rank)
         if loaded_shard_id is None:
             # Loaded weight is already fused on disk
-            if self.nranks != 1:
-                shard_offsets = [
-                    # (shard_id, shard_offset, shard_size)
-                    ("q", 0, self.num_heads * self.head_dim),
-                    ("k", self.num_heads * self.head_dim, self.kv_num_heads * self.head_dim),
-                    ("v", (self.num_heads + self.kv_num_heads) * self.head_dim, self.kv_num_heads * self.head_dim),
-                ]
-                for shard_id, shard_offset, shard_size in shard_offsets:
-                    loaded_weight_shard = loaded_weight[..., shard_offset : shard_offset + shard_size]
-                    self.weight_loader(param, loaded_weight_shard, shard_id)
-            else:
-                loaded_weight = get_tensor(loaded_weight)
-                split_loaded_weight = loaded_weight
-                param.copy_(split_loaded_weight, False)
+            shard_offsets = [
+                # (shard_id, shard_offset, shard_size)
+                ("q", 0, self.num_heads * head_dim),
+                ("k", self.num_heads * head_dim, self.kv_num_heads * head_dim),
+                ("v", (self.num_heads + self.kv_num_heads) * head_dim, self.kv_num_heads * head_dim),
+            ]
+            for shard_id, shard_offset, shard_size in shard_offsets:
+                loaded_weight_shard = slice_fn(
+                    loaded_weight, output_dim, start=shard_offset, end=shard_offset + shard_size
+                )
+                self.weight_loader(param, loaded_weight_shard, shard_id)
         else:
-            # 1.fused qkv in disk
-            # 2.split q k v
+            # 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:
-                dim = -1
+            if self.nranks != 1:
+                dim = -1 if output_dim else 0
                 if isinstance(loaded_weight, np.ndarray):
                     size = loaded_weight.shape[dim]
                 else:
                     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, start=shard_offset, end=shard_size)
 
             loaded_weight = get_tensor(loaded_weight)
+            if not param._is_initialized():
+                param.initialize()
 
             if loaded_shard_id == "q":
-                param = param[:, : 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,
-                ]
-            elif loaded_shard_id == "v":
-                param = param[:, (self.num_heads_per_rank + self.kv_num_heads_per_rank) * self.head_dim :]
 
+                param_shard_offset = 0
+                param_shard_size = self.num_heads_per_rank * head_dim
+            elif loaded_shard_id == "k":
+                param_shard_offset = self.num_heads_per_rank * head_dim
+                param_shard_size = self.kv_num_heads_per_rank * head_dim
+            else:
+                # loaded_shard_id == "v"
+                param_shard_offset = (self.num_heads_per_rank + self.kv_num_heads_per_rank) * head_dim
+                param_shard_size = self.kv_num_heads_per_rank * head_dim
+            if hasattr(param, "tensor_track"):
+                param.tensor_track.mark(start=param_shard_offset, end=param_shard_offset + param_shard_size)
+            param = slice_fn(param, output_dim, start=param_shard_offset, end=param_shard_offset + param_shard_size)
             assert param.shape == loaded_weight.shape, (
                 f" Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({param.shape})"
             )
@@ -698,7 +707,6 @@ def __init__(
             ),
         )
         if self.nranks > 0:
-            _set_var_distributed(self.weight, split_axis=0)
             if self.with_bias:
                 # col parallel
                 _set_var_distributed(self.bias, split_axis=0)

fastdeploy/model_executor/layers/lm_head.py

@@ -22,7 +22,7 @@
 from paddle.distributed import fleet
 
 from fastdeploy.config import FDConfig
-from fastdeploy.model_executor.models.utils import set_weight_attrs
+from fastdeploy.model_executor.utils import set_weight_attrs
 
 from .utils import get_tensor
 

fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py

@@ -19,7 +19,7 @@
 import paddle
 from paddle import nn
 
-from fastdeploy.model_executor.layers.utils import set_weight_attrs
+from fastdeploy.model_executor.utils import set_weight_attrs
 from fastdeploy.platforms import current_platform
 
 from ..quantization.quant_base import QuantMethodBase
@@ -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,