[feat]add fast_weights_iterator (#3258)

* add fast_weights_iterator

* update

* update

Author: bukejiyu

fastdeploy/model_executor/layers/linear.py

@@ -162,11 +162,6 @@ def __init__(
         self.linear_shift = None
         self.linear_smooth = None
 
-        if fd_config.model_config.is_quantized:
-            self.weight_key = f"{prefix}.quant_weight"
-            self.weight_scale_key = f"{prefix}.weight_scale"
-            self.act_scale_key = f"{prefix}.activation_scale"
-
     def load_prequant_weight(self, state_dict: dict):
         """
         Load the prequantized weight from the state dictionary.

fastdeploy/model_executor/load_weight_utils.py

@@ -24,6 +24,7 @@
 from paddleformers.transformers import PretrainedModel
 from paddleformers.transformers.model_utils import load_tp_checkpoint
 from paddleformers.utils.log import logger
+from paddleformers.utils.safetensors import fast_safe_open
 from safetensors import safe_open
 from tqdm import tqdm
 
@@ -155,18 +156,28 @@ def get_expert_ranges(fd_config):
     return state_dict
 
 
-def safetensors_weights_iterator(
-    safe_tensor_list: list[str],
-):
+def safetensors_weights_iterator(safe_tensor_list: list[str]):
     """
     safetensors_weights_iterator
     """
     for st_file in tqdm(
         safe_tensor_list,
         desc="Loading safetensors checkpoint shards",
     ):
-        from paddleformers.utils.safetensors import fast_safe_open
+        with safe_open(st_file, framework="np") as f:
+            for name in f.keys():
+                param = f.get_tensor(name)
+                yield name, param
 
+
+def fast_weights_iterator(safe_tensor_list: list[str]):
+    """
+    paddleformers' iterator for safetensors
+    """
+    for st_file in tqdm(
+        safe_tensor_list,
+        desc="Loading safetensors checkpoint shards",
+    ):
         with fast_safe_open(st_file, framework="np") as f:
             for name in f.keys():
                 param = f.get_slice(name)
@@ -215,13 +226,12 @@ def load_pre_sharded_checkpoint(model_path: str, local_rank: int, use_fastsafete
     """
     load_pre_sharded_checkpoint
     """
-    from fastdeploy.model_executor.layers.utils import get_tensor
 
     state_dict = {}
     _, safetensor_files = get_all_safetensors(os.path.join(model_path, f"rank{local_rank}"))
     weights_iterator = safetensors_weights_iterator(safetensor_files)
     for name, weight in weights_iterator:
-        state_dict[name] = get_tensor(weight)
+        state_dict[name] = weight
     return state_dict
 
 

fastdeploy/model_executor/model_loader/default_loader_v1.py

@@ -22,9 +22,9 @@
 
 from fastdeploy.config import FDConfig, LoadConfig, ModelConfig
 from fastdeploy.model_executor.load_weight_utils import (
+    fast_weights_iterator,
     get_all_safetensors,
     measure_time,
-    safetensors_weights_iterator,
 )
 from fastdeploy.model_executor.model_loader.base_loader import BaseModelLoader
 from fastdeploy.model_executor.models.model_base import ModelRegistry
@@ -49,7 +49,7 @@ def clean_memory_fragments(self) -> None:
     @measure_time
     def load_weights(self, model, fd_config: FDConfig) -> None:
         _, safetensor_files = get_all_safetensors(fd_config.model_config.model)
-        weights_iterator = safetensors_weights_iterator(safetensor_files)
+        weights_iterator = fast_weights_iterator(safetensor_files)
         model.load_weights(weights_iterator)
         self.clean_memory_fragments()