loadworker > 1

Author: shihaobai

lightllm/common/basemodel/layer_weights/meta_weights/__init__.py

@@ -1,4 +1,12 @@
 from .base_weight import BaseWeight
-from .mm_weight import MMWeight, ROWMMWeight, COLMMWeight, MultiROWMMWeight, CustomMMWeight, CustomBMMWeight
+from .mm_weight import (
+    MMWeight,
+    MultiMMWeight,
+    ROWMMWeight,
+    COLMMWeight,
+    MultiROWMMWeight,
+    CustomMMWeight,
+    CustomBMMWeight,
+)
 from .norm_weight import NormWeight, GEMMANormWeight, TpNormWeight
 from .fused_moe_weight import FusedMoeWeight

lightllm/common/basemodel/layer_weights/transformer_layer_weight.py

@@ -2,7 +2,7 @@
 
 # from lightllm.common.layers.mm import MM
 from .base_layer_weight import BaseLayerWeight
-from .meta_weights import MMWeight, ROWMMWeight, FusedMoeWeight
+from .meta_weights import BaseWeight, MultiMMWeight, MMWeight, FusedMoeWeight
 from lightllm.utils.log_utils import init_logger
 
 logger = init_logger(__name__)
@@ -34,6 +34,18 @@ def _init_weight_names(self):
     def _init_weight(self):
         pass
 
+    def load_hf_weights(self, weights):
+        """
+        load weights
+        """
+        for attr_name in dir(self):
+            attr = getattr(self, attr_name, None)
+            if isinstance(attr, MultiMMWeight):
+                with self.lock:
+                    attr.load_hf_weights(weights)
+            elif isinstance(attr, BaseWeight):
+                attr.load_hf_weights(weights)
+
     def set_quantization(self):
         if self.quant_cfg.quant_type is None:
             return