fusion-bench 0.2.18__py3-none-any.whl → 0.2.19__py3-none-any.whl

fusion_bench/utils/lazy_state_dict.py

@@ -14,6 +14,7 @@ from torch import nn
 from transformers import AutoConfig
 
 from fusion_bench.utils.dtype import parse_dtype
+from fusion_bench.utils.packages import import_object
 
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
@@ -54,10 +55,13 @@ class LazyStateDict:
     """
 
     _local_path: str
+    """local path to the checkpoint."""
     _state_dict_cache: Optional[Dict]
+    """Cache for the state dict, if enabled."""
     _index_filename: Optional[str]
     _checkpoint_files: Optional[List[str]]
-    _index: Optional[Dict]
+    _index: Optional[Dict[str, str]]
+    """Mapping of parameter names to checkpoint files."""
 
     def __init__(
         self,
@@ -71,7 +75,22 @@ class LazyStateDict:
         hf_cache_dir: Optional[str] = None,
         hf_proxies: Optional[Dict] = None,
     ):
+        """
+        Args:
+            checkpoint (str): Path to the checkpoint file or directory.
+            meta_module_class (Type[nn.Module], optional): Class of the meta module to instantiate.
+            meta_module (nn.Module, optional): Pre-initialized meta module.
+            cache_state_dict (bool): Whether to cache the state dict in memory.
+            torch_dtype (torch.dtype, optional): The dtype to use for the tensors.
+            device (str): The device to load the tensors onto.
+            hf_revision (str, optional): The revision of the model to download from Hugging Face Hub.
+            hf_cache_dir (str, optional): The cache directory for Hugging Face models.
+            hf_proxies (Dict, optional): Proxies to use for downloading from Hugging Face Hub.
+        """
+        self.cache_state_dict = cache_state_dict
         self.meta_module_class = meta_module_class
+        if isinstance(self.meta_module_class, str):
+            self.meta_module_class = import_object(self.meta_module_class)
         self.meta_module = meta_module
         if self.meta_module_class is not None:
             if self.meta_module is not None:
@@ -110,6 +129,18 @@ class LazyStateDict:
                 self._state_dict_cache = {}
             else:
                 self._state_dict_cache = None
+        elif len(self._checkpoint_files) == 1 and self._checkpoint_files[0].endswith(
+            SAFE_WEIGHTS_NAME
+        ):
+            # let the keys of self._index be the keys of the state dict, the values are the checkpoint file
+            with safe_open(
+                self._checkpoint_files[0], framework="pt", device=device
+            ) as f:
+                self._index = {key: self._checkpoint_files[0] for key in f.keys()}
+                if cache_state_dict:
+                    self._state_dict_cache = {}
+                else:
+                    self._state_dict_cache = None
         elif len(self._checkpoint_files) == 1 and self._checkpoint_files[0].endswith(
             WEIGHTS_NAME
         ):
@@ -137,7 +168,11 @@ class LazyStateDict:
     def config(self) -> "PretrainedConfig":
         return AutoConfig.from_pretrained(self._checkpoint)
 
-    def state_dict(self) -> "LazyStateDict":
+    def state_dict(self, keep_vars: bool = False) -> "LazyStateDict":
+        """
+        Args:
+            keep_vars (bool): Ignored, as LazyStateDict does not support keep_vars. Just for compatibility.
+        """
         return self
 
     def _resolve_checkpoint_files(self, checkpoint: str):
@@ -255,6 +290,21 @@ class LazyStateDict:
             )
             return tensor
 
+    def __setitem__(self, key: str, value: torch.Tensor) -> None:
+        """
+        Set a tensor in the LazyStateDict. This will update the state dict cache if it is enabled.
+        """
+        assert key in list(
+            self.keys()
+        ), "KeyError: Cannot set a tensor for a key that does not exist in the LazyStateDict."
+        if self._state_dict_cache is not None:
+            self._state_dict_cache[key] = value
+        else:
+            log.warning(
+                "State dict cache is disabled, setting a tensor will not update the cache."
+            )
+            self._state_dict_cache = {key: value}
+
     def __contains__(self, key: str) -> bool:
         if self._state_dict_cache is not None and key in self._state_dict_cache:
             return True
@@ -314,7 +364,7 @@ class LazyStateDict:
 
     def __repr__(self) -> str:
         if self._index is not None:
-            return f"{self.__class__.__name__}(index={self._index})"
+            return f"{self.__class__.__name__}(keys={list(self.keys())})"
         else:
             return (
                 f"{self.__class__.__name__}(checkpoint_files={self._checkpoint_files})"
@@ -336,3 +386,25 @@ class LazyStateDict:
             raise RuntimeError(
                 "Cannot get submodule because meta_module is not provided."
             )
+
+    def load_state_dict(
+        self, state_dict: Dict[str, torch.Tensor], strict: bool = True
+    ) -> None:
+        """
+        Load a state dict into this LazyStateDict.
+        This method is only for compatibility with nn.Module and it overrides the cache of LazyStateDict.
+
+        Args:
+            state_dict (Dict[str, torch.Tensor]): The state dict to load.
+            strict (bool): Whether to enforce that all keys in the state dict are present in this LazyStateDict.
+        """
+        log.warning(
+            "Loading state dict into LazyStateDict is not recommended, as it may lead to unexpected behavior. "
+            "Use with caution."
+        )
+        if strict:
+            for key in state_dict:
+                if key not in self:
+                    raise KeyError(f"Key {key} not found in LazyStateDict.")
+        for key, value in state_dict.items():
+            self[key] = value

fusion_bench/utils/misc.py

@@ -1,6 +1,6 @@
 from typing import Iterable, List
 
-__all__ = ["first", "has_length", "join_list"]
+__all__ = ["first", "has_length", "join_list", "attr_equal"]
 
 
 def first(iterable: Iterable):
@@ -23,3 +23,21 @@ def join_list(list_of_list: List[List]):
     for item in list_of_list:
         ans.extend(item)
     return ans
+
+
+def attr_equal(obj, attr: str, value):
+    """
+    Check if the attribute of the object is equal to the given value.
+    Returns False if the attribute does not exist or is not equal to the value.
+
+    Args:
+        obj: The object to check.
+        attr (str): The attribute name to check.
+        value: The value to compare against.
+
+    Returns:
+        bool: True if the attribute exists and is equal to the value, False otherwise.
+    """
+    if not hasattr(obj, attr):
+        return False
+    return getattr(obj, attr) == value

{fusion_bench-0.2.18.dist-info → fusion_bench-0.2.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.18
+Version: 0.2.19
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License

{fusion_bench-0.2.18.dist-info → fusion_bench-0.2.19.dist-info}/RECORD

@@ -1,4 +1,4 @@
-fusion_bench/__init__.py,sha256=68dF-zPvb8E2MgYnmgIJsxIHJBy1MApKeOrRZvQEVlg,421
+fusion_bench/__init__.py,sha256=vu3nMzKuiiHkbH13m8SOzj8qYU-n1PreBipWs_xjZig,1937
 fusion_bench/__main__.py,sha256=weUjxpP3ULnDgUxCehdbmoCM9cqfkhDhGB85tAF5qoE,81
 fusion_bench/compat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fusion_bench/compat/method/__init__.py,sha256=A9pbskEYB4_ryD6sVrR0qI4eVYsbI7sExbhPeypP3fQ,5757
@@ -12,6 +12,7 @@ fusion_bench/compat/taskpool/base_pool.py,sha256=1AIZBxqUJgshq0Xo3Yo9es4b-8X8ksN
 fusion_bench/compat/taskpool/clip_image_classification.py,sha256=ZYZsbsE-fPzm6yafA0p-6wcDwVGryLmtXXtuEXeQbTY,7425
 fusion_bench/compat/taskpool/flan_t5_glue_text_generation.py,sha256=JsdAE72V1C1eDcA1WCa0PIcSDTrGPclNKFDQ9G-hYts,5786
 fusion_bench/constants/__init__.py,sha256=Pyc4dLbl6oNduOCdnpeXQ9LDyVoIrkdl9eZ_l2axIv8,41
+fusion_bench/constants/banner.py,sha256=fuIO36ETKlS6a3wbwZn-rA2OswSCfOYyyhZ0Fnal1s4,1656
 fusion_bench/constants/clip_vision.py,sha256=k0NRfiacxRaswdxUj91-e3jcP1u-RmvsaaYdqohcQVU,310
 fusion_bench/constants/paths.py,sha256=DVZyQ9FLhkyUdw6ARpXUCAMf_B8hFyJ6UNI-oYly3pE,591
 fusion_bench/dataset/__init__.py,sha256=OJiYmcqz0Vm5O7mE4PB5QFJeL_KjrsseQTRsQATGTm4,1050
@@ -47,7 +48,7 @@ fusion_bench/method/base_algorithm.py,sha256=UuITuGnSskcKEwUVINuPoWJUwqGm9AIgyQI
 fusion_bench/method/dummy.py,sha256=hb1y6LR_geRZ5eRgGwt5zJUcHYorCeIbs5i76CvurUc,1031
 fusion_bench/method/ensemble.py,sha256=rGxvJTeorfcBuE_e0XO-0-MAc9un7ZCC46ikKGuAcN4,3077
 fusion_bench/method/model_recombination.py,sha256=2tviqmYSPOL0_Ktv8_gt_YzQ4tyCANHxXquUot_3Cgo,5360
-fusion_bench/method/simple_average.py,sha256=vVzlfdf0mPHeY3VeOLrcWI4sWoLBW0gaX0lusjePVyQ,4539
+fusion_bench/method/simple_average.py,sha256=A_VYtHhECcxY0_Mppe5ThOY-ip6XUvvtPHsaQKSmDPc,4971
 fusion_bench/method/ada_svd/__init__.py,sha256=4XzQbbvE9HI3NtEmEFvo8iC3ds_85vJXe7P7qJfL7kk,77
 fusion_bench/method/ada_svd/clip_vision.py,sha256=XvXgIdlShAREMsubRgphyycGrhWqSnuVBo6S9bNYSd0,12581
 fusion_bench/method/adamerging/__init__.py,sha256=nt0saBT_3bqghk-pINQ-XCWm9UWwSZllu4R1sDuAJAA,376
@@ -120,7 +121,7 @@ fusion_bench/method/linear/__init__.py,sha256=ChfkoOEAb-rUKwpowFPel-a1hRfS8gCrbn
 fusion_bench/method/linear/expo.py,sha256=LCHTWlsPm1Mjhrq0mfpWLVC7skkI9ZksGduy3TxULoU,3939
 fusion_bench/method/linear/linear_interpolation.py,sha256=IONw9BPiRJouY8bE9Abfyz7qVI_1B1n8KGZa0f7Pza8,2157
 fusion_bench/method/linear/llama_expo.py,sha256=ccECjhAqcFmzOIDyZ7e_aPzTM2Kj8u2D8TJytyz18YM,8476
-fusion_bench/method/linear/simple_average_for_llama.py,sha256=7JlVrmTMmrePvNGnZNoxSuCSq2Vu7cPQzjGC3WWUXBE,2079
+fusion_bench/method/linear/simple_average_for_llama.py,sha256=OcjvfG5nuUzdo3P4Xi1mO6ApRu51YAUYXG5lAMeD6rg,2711
 fusion_bench/method/linear/task_arithmetic_for_llama.py,sha256=4SZpiTD7OzhWUXtcdK3PYdXbBGyDqiZd7oZOQ0lraN0,1963
 fusion_bench/method/lm_finetune/__init__.py,sha256=IFGAqXujX3Fabzl_tC6zZyOyPFJfVziL0qFtj5MVxj0,149
 fusion_bench/method/lm_finetune/bradley_terry_rm.py,sha256=ys_td1IeL3bzPTE0Cixlj2JooCaB7qseRwSDwroAk5A,18777
@@ -252,7 +253,7 @@ fusion_bench/modelpool/huggingface_gpt2_classification.py,sha256=j8nicVwtoLXY4RP
 fusion_bench/modelpool/lazy_state_dict_pool.py,sha256=HtEA85rqSCHfsIddI5sKDcZf5kSuHNwrb8fF1TUSTr0,652
 fusion_bench/modelpool/nyuv2_modelpool.py,sha256=btuXmYxwfjI6MnGakhoOf53Iyb9fxYH20CavGTrTcnA,1375
 fusion_bench/modelpool/causal_lm/__init__.py,sha256=F432-aDIgAbUITj4GNZS9dgUKKhaDMCbTeHB-9MecaQ,99
-fusion_bench/modelpool/causal_lm/causal_lm.py,sha256=dkumbKspfEJhp3gtlZC71zUutdfJOKpKZnHy5z97qbc,6727
+fusion_bench/modelpool/causal_lm/causal_lm.py,sha256=7-mUWVGVsXyljH_06CmIyReClKx_xVjy5zeXTJcLQIk,8085
 fusion_bench/modelpool/clip_vision/__init__.py,sha256=3b9gN2bWUsoA1EmpitnIMnIlX7nklxbkn4WJ0QJtS2c,43
 fusion_bench/modelpool/clip_vision/modelpool.py,sha256=ADgzslXwYd95x42V26XvgS09WEKGfhH_AYuQmWKdT0w,5887
 fusion_bench/modelpool/openclip_vision/__init__.py,sha256=QDmAitKqUwRygN9QncdS_kGWZdfTKL4uUifC8xh9c10,47
@@ -384,7 +385,7 @@ fusion_bench/tasks/clip_classification/fer2013.py,sha256=_oc2fdV308ywcb16rLZxBAd
 fusion_bench/tasks/clip_classification/flower102.py,sha256=p_JMs6HCCPZBKe7PTXt0WABsd-KcgmpBkxDSlJJaVVY,2096
 fusion_bench/tasks/clip_classification/food101.py,sha256=Oepvws5byGxrHswXt3ILG3UEPiZaFXYqK1yJqm1uYVE,1968
 fusion_bench/tasks/clip_classification/gtsrb.py,sha256=Dsaz-XNz6oA9nNTF2C2iXmmhhVz-gsw-WcGuFTqjzl4,2677
-fusion_bench/tasks/clip_classification/imagenet.py,sha256=Az7gnFuecVCDkP3mMjiLwOgrYAf_cxz177kkdivComU,38815
+fusion_bench/tasks/clip_classification/imagenet.py,sha256=EhZ2iYAc8oApr5BU_vgM3cDY879anTkvb-5hfi-B7m4,48826
 fusion_bench/tasks/clip_classification/kmnist.py,sha256=Ohce6aVaXkPnviDaZYXANMhhBNHZXO3FnXYxYG-ISVg,311
 fusion_bench/tasks/clip_classification/mnist.py,sha256=-gQpHz_kCXmUOtAsM8FBUYFjlwcbAgnqpuVtRfCJ3JM,129
 fusion_bench/tasks/clip_classification/mongo_leaf_disease.py,sha256=L_2IgnzbZdGZrX27VNGu1rC-N3Aj4fetIXB9HM1QZkI,519
@@ -417,8 +418,8 @@ fusion_bench/utils/hydra_utils.py,sha256=TklUDKDEZlg4keI-TEZiqh4gFjr9-61Rt1RMlqk
 fusion_bench/utils/instantiate_utils.py,sha256=57D8YP25OO-ArltOSsHDKtnNcA44m1yAq-1wKZc2YVI,17523
 fusion_bench/utils/json.py,sha256=sVCqbm9mmyHybiui-O57KFt_ULrjLtN2wipSo6VDvqE,2533
 fusion_bench/utils/lazy_imports.py,sha256=v5l9cpHXPMaz1IVBmB5oOqefYr9vA3XvP340xT7Wy18,2796
-fusion_bench/utils/lazy_state_dict.py,sha256=xb_NM4F653_HiPK8OClG3oTPuPk4SaarfCtLLg87Yi8,13347
-fusion_bench/utils/misc.py,sha256=Rgec7eKcGIcp9BaFVdm2pzx0J-L8AyX5qWuiYNTGvTc,530
+fusion_bench/utils/lazy_state_dict.py,sha256=Hu8PkhbJcUikXJxWUJ7vabu2uDbnUUF6UsRS0k8i71U,16841
+fusion_bench/utils/misc.py,sha256=Qc3_H8UMooOp81Ow89zqvM1sNPIybq1cbq7s4-4lsfU,1082
 fusion_bench/utils/packages.py,sha256=L64paDi1SmeT3gRvRV6LaqB8AeGdzIYWIRI31qSQbSk,2110
 fusion_bench/utils/parameters.py,sha256=2vs8vo2o-nRA9NOMOYFye-X8-aHQZoYe54tM6n0r0RE,11757
 fusion_bench/utils/path.py,sha256=hRA1CPHNnTYBUmzbftH77sHvn4aTuybEK5Tth1skP-k,531
@@ -436,7 +437,7 @@ fusion_bench/utils/plot/token_notebook.py,sha256=bsntXf46Zz_RavTxNiB9c3-KvHw7LFw
 fusion_bench/utils/strenum/__init__.py,sha256=id9ORi1uXrDxhbmVxitJ1KDwLS4H3AAwFpaK5h1cQzw,8531
 fusion_bench/utils/strenum/_name_mangler.py,sha256=o11M5-bURW2RBvRTYXFQIPNeqLzburdoWLIqk8X3ydw,3397
 fusion_bench/utils/strenum/_version.py,sha256=6JQRo9LcvODbCOeVFYQb9HNJ_J9XiG_Zbn8ws2A3BV8,18466
-fusion_bench-0.2.18.dist-info/licenses/LICENSE,sha256=nhnOJlw4CPuPVE0qvkGmxfFgHmKi-6nzXvTu8t0NUdg,1066
+fusion_bench-0.2.19.dist-info/licenses/LICENSE,sha256=nhnOJlw4CPuPVE0qvkGmxfFgHmKi-6nzXvTu8t0NUdg,1066
 fusion_bench_config/README.md,sha256=Lc8YSBJ5oxf9KV5kKDivJ9LRyGuraGQPmBbgbdVA-j4,703
 fusion_bench_config/clip-vit-base-patch32_robustness_corrupted.yaml,sha256=7IxLQoLRz-sRWyV8Vqc5kQcmYE_9YQz2_77pmvAkum8,1207
 fusion_bench_config/fabric_model_fusion.yaml,sha256=YwJx_aUXm4ca4_mVItKVUOesMvmBBRGudQIOqgc1EP8,974
@@ -795,6 +796,7 @@ fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-large-patch14_TALL14_
 fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-large-patch14_TALL20.yaml,sha256=yC2U_IoBAhawgSahY_mdi7ea5kJ2SSRPJ2FM-bA-E9M,510
 fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-large-patch14_TALL20_model_only.yaml,sha256=a2nviqKSRNoQScYVbj5buq0PbUzmYJwNWdPBUoLaeV8,386
 fusion_bench_config/modelpool/CLIPVisionModelPool/clip-vit-large-patch14_individual.yaml,sha256=G6yvZuWOKb75RLn6tu2LPnwHUyvoxPfL_wqb_B11aZo,549
+fusion_bench_config/modelpool/CausalLMPool/Qwen2.5-1.5B_math_and_coder.yaml,sha256=HZXjqbZKpSZCHb-G8qjj03PcvXg_8mrAuewDHZp0oEw,263
 fusion_bench_config/modelpool/CausalLMPool/deepseek-v2-lite.yaml,sha256=8gr8ZtgegSHV0GHtJBiEgdYbRe8UHhO4_y8dayxZChk,506
 fusion_bench_config/modelpool/CausalLMPool/llama_alpaca_cleaned.yaml,sha256=oDsZkuAoh1mWUC7jZNzw8794zgX2bV5Z0esXpvbTs-c,643
 fusion_bench_config/modelpool/CausalLMPool/llama_codealpaca.yaml,sha256=FuUsBrvk3_bQiciMRlNsO5vp6AKHQM_-g-8bmU8251w,641
@@ -875,8 +877,8 @@ fusion_bench_config/taskpool/LMEvalHarnessTaskPool/lm_eval.yaml,sha256=3q-KMuFaM
 fusion_bench_config/taskpool/OpenCLIPVisionModelTaskPool/ViT-B-16_TA8.yaml,sha256=GjpiiRownrBCpl-TNwWRW2PYePbF-Cl99jlLNPrK5T4,1017
 fusion_bench_config/taskpool/OpenCLIPVisionModelTaskPool/ViT-B-32_TA8.yaml,sha256=WwiYMQKehtJixDPnu5o3vcWe4yJksXTWRqOzm3uVWXQ,1017
 fusion_bench_config/taskpool/OpenCLIPVisionModelTaskPool/ViT-L-14_TA8.yaml,sha256=xGRt0J9joXTzWUew6DvoYprAWlPXhaVFw5AX4im5VQw,1017
-fusion_bench-0.2.18.dist-info/METADATA,sha256=igyW5oJQzJfuEagCgyNbi0MvXp-Rz56u3FPtIHJFG5Y,21966
-fusion_bench-0.2.18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-fusion_bench-0.2.18.dist-info/entry_points.txt,sha256=iUQ8MCJvda7HP4vYh2n1Teoapb4G9PBVYZkAfcc5SHU,116
-fusion_bench-0.2.18.dist-info/top_level.txt,sha256=BuO4TL6iHL_2yPBUX9-LlIrHRczA_BNMIFwweK0PQEI,13
-fusion_bench-0.2.18.dist-info/RECORD,,
+fusion_bench-0.2.19.dist-info/METADATA,sha256=5pl4dtlAYklMMiMLBeKNaHqCQRd7sLSct7aIh9JIoGY,21966
+fusion_bench-0.2.19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+fusion_bench-0.2.19.dist-info/entry_points.txt,sha256=iUQ8MCJvda7HP4vYh2n1Teoapb4G9PBVYZkAfcc5SHU,116
+fusion_bench-0.2.19.dist-info/top_level.txt,sha256=BuO4TL6iHL_2yPBUX9-LlIrHRczA_BNMIFwweK0PQEI,13
+fusion_bench-0.2.19.dist-info/RECORD,,

fusion_bench_config/modelpool/CausalLMPool/Qwen2.5-1.5B_math_and_coder.yaml

@@ -0,0 +1,11 @@
+_target_: fusion_bench.modelpool.CausalLMPool
+_recursive_: false
+
+load_lazy: false
+models:
+  _pretrained_: Qwen/Qwen2.5-1.5B
+  expert_1: Qwen/Qwen2.5-Math-1.5B
+  expert_2: Qwen/Qwen2.5-Coder-1.5B
+model_kwargs:
+  torch_dtype: bfloat16
+tokenizer: Qwen/Qwen2.5-1.5B