qwen loader (#3057)

fastdeploy/model_executor/models/qwen3.py

@@ -228,7 +228,7 @@ class Qwen3ForCausalLM(ModelForCasualLM):
             fd_config (FDConfig): Configurations for the LLM model.
         """
         super(Qwen3ForCausalLM, self).__init__(fd_config)
-
+        self.fd_config = fd_config
         self.model = Qwen3Model(fd_config=fd_config)
 
         self.ori_vocab_size = fd_config.model_config.ori_vocab_size
@@ -245,6 +245,47 @@ class Qwen3ForCausalLM(ModelForCasualLM):
         """ """
         return "Qwen3ForCausalLM"
 
+    @paddle.no_grad()
+    def load_weights(self, weights_iterator) -> None:
+        """
+        Load model parameters from a given weights_iterator object.
+
+        Args:
+            weights_iterator (Iterator): An iterator yielding (name, weight) pairs.
+        """
+
+        from fastdeploy.model_executor.models.utils import default_weight_loader
+
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("up_gate_proj", "gate_proj", "gate"),
+            ("up_gate_proj", "up_proj", "up"),
+            ("embed_tokens.embeddings", "embed_tokens", None),
+            ("lm_head.linear", "lm_head", None),
+        ]
+
+        params_dict = dict(self.named_parameters())
+        for loaded_weight_name, loaded_weight in weights_iterator:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in loaded_weight_name:
+                    continue
+                model_param_name = loaded_weight_name.replace(weight_name, param_name)
+                if model_param_name not in params_dict:
+                    continue
+                param = params_dict[model_param_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader(self.fd_config))
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if loaded_weight_name not in params_dict:
+                    continue
+                param = params_dict[loaded_weight_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader(self.fd_config))
+                weight_loader(param, loaded_weight)
+
     @paddle.no_grad()
     def set_state_dict(self, state_dict):
         """

fastdeploy/model_executor/models/utils.py

@@ -24,7 +24,7 @@ import random
 import re
 import struct
 from functools import partial
-from typing import NamedTuple, Optional
+from typing import Any, NamedTuple, Optional, Union
 
 import numpy as np
 import paddle
@@ -40,10 +40,51 @@ from paddleformers.utils.env import (
 from paddleformers.utils.log import logger
 from tqdm import tqdm
 
+from fastdeploy.config import FDConfig
+from fastdeploy.model_executor.layers.utils import get_tensor
+
 MAX_BSZ = 512
 MAX_DRAFT_TOKENS = 6
 
 
+def set_weight_attrs(param, param_attr_map: Optional[dict[str, Any]]):
+    if param_attr_map is None:
+        return
+    for key, value in param_attr_map.items():
+        setattr(param, key, value)
+
+
+def default_weight_loader(fd_config: FDConfig) -> None:
+    """Default weight loader"""
+
+    def fn(param, loaded_weight, shard_id: Optional[Union[int, str]] = None):
+        """fn"""
+        try:
+            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 output_dim else 0
+                size = loaded_weight.get_shape()[dim]
+                block_size = size // fd_config.parallel_config.tensor_parallel_size
+                shard_offset = fd_config.parallel_config.tensor_parallel_rank * block_size
+                shard_size = (fd_config.parallel_config.tensor_parallel_rank + 1) * block_size
+                if output_dim:
+                    loaded_weight = loaded_weight[..., shard_offset:shard_size]
+                else:
+                    loaded_weight = loaded_weight[shard_offset:shard_size, ...]
+            loaded_weight = get_tensor(loaded_weight)
+
+            assert param.shape == loaded_weight.shape, (
+                f" Attempted to load weight ({loaded_weight.shape}) " f"into parameter ({param.shape})"
+            )
+
+            param.copy_(loaded_weight, False)
+        except Exception:
+            raise
+
+    return fn
+
+
 class LayerIdPlaceholder(str, enum.Enum):
     """LayerIdPlaceholder"""