fusion-bench 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

fusion_bench/method/we_moe/flan_t5_we_moe.py

@@ -16,11 +16,14 @@ from transformers.data import default_data_collator
 
 from fusion_bench.method import BaseAlgorithm
 from fusion_bench.method.task_arithmetic.task_arithmetic import task_arithmetic_merge
-from fusion_bench.mixins.lightning_fabric import LightningFabricMixin
-from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.mixins import (
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+    auto_register_config,
+)
 from fusion_bench.modelpool import Seq2SeqLMPool
 from fusion_bench.models.we_moe import WeightEnsemblingMoE
-from fusion_bench.utils import timeit_context
+from fusion_bench.utils import print_parameters, timeit_context
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
 from fusion_bench.utils.instantiate_utils import instantiate
 from fusion_bench.utils.parameters import print_parameters
@@ -31,10 +34,11 @@ from .utils import get_memory_usage
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class FlanT5WeightEnsemblingMoEAlgorithm(
-    BaseAlgorithm,
     LightningFabricMixin,
     SimpleProfilerMixin,
+    BaseAlgorithm,
 ):
     """
     FlanT5WeightEnsemblingMoEAlgorithm is a class that implements the WeightEnsemblingMoEAlgorithm
@@ -60,7 +64,6 @@ class FlanT5WeightEnsemblingMoEAlgorithm(
         num_workers: int = 0,
         max_steps: int = 1000,
         use_grad_accumulate: bool = True,
-        cache_dir: bool = "outputs",
         fast_dev_run: bool = False,
         **kwargs,
     ):
@@ -70,23 +73,9 @@ class FlanT5WeightEnsemblingMoEAlgorithm(
         Args:
             algorithm_config (DictConfig): The configuration for the algorithm.
         """
-        self.checkpoint = checkpoint
-        self.save_checkpoint = save_checkpoint
-        self.router_hidden_layers = router_hidden_layers
-        self.init_lambda = init_lambda
-        self.batch_reduce = batch_reduce
-        self.lr = lr
-        self.optimizer = optimizer
-        self.devices = devices
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.max_steps = max_steps
-        self.use_grad_accumulate = use_grad_accumulate
-        self.cache_dir = cache_dir
-        self.fast_dev_run = fast_dev_run
         super().__init__(**kwargs)
 
-    def construct_moe_model(self) -> WeightEnsemblingMoE:
+    def construct_moe_model(self):
         """
         Construct the Mixture of Experts (MoE) model using the models in the model pool.
 

fusion_bench/method/wudi/__init__.py

@@ -0,0 +1 @@
+from .wudi import WUDIMerging, wudi_merging

fusion_bench/method/wudi/wudi.py

@@ -0,0 +1,105 @@
+"""
+Whoever Started the Interference Should End It:  Guiding Data-Free Model Merging via Task Vectors
+Arxiv: http://arxiv.org/abs/2503.08099
+"""
+
+from typing import List
+
+import torch
+from tqdm import tqdm
+
+from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.utils import timeit_context
+from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_sub
+
+
+def wudi_merging(
+    task_vectors: List[torch.Tensor],
+    accelerator="cuda",
+    iter_num: int = 300,
+    exclude_keys: List[str] = None,
+):
+    exclude_keys = [] if exclude_keys is None else exclude_keys
+
+    with timeit_context("WUDI Merging"):
+        new_vector = {}
+        for key in tqdm(task_vectors[0], desc="WUDI Merging", leave=False):
+            tqdm.write(f"key: {key}")
+            original_device = task_vectors[0][key].device
+            tvs = torch.stack(
+                [
+                    task_vector[key].to(device=accelerator, non_blocking=True)
+                    for task_vector in task_vectors
+                ]
+            )
+            num_tvs = len(tvs)
+            new_vector[key] = torch.nn.Parameter(torch.sum(tvs, dim=0))
+
+            if len(task_vectors[0][key].shape) == 2 and key not in exclude_keys:
+                optimizer = torch.optim.Adam([new_vector[key]], lr=1e-5, weight_decay=0)
+                l2_norms = torch.square(
+                    torch.norm(tvs.reshape(tvs.shape[0], -1), p=2, dim=-1)
+                )
+                for i in tqdm(
+                    range(iter_num),
+                ):
+                    disturbing_vectors = new_vector[key].unsqueeze(0) - tvs
+                    product = torch.matmul(disturbing_vectors, tvs.transpose(1, 2))
+                    loss = torch.sum(
+                        torch.square(product) / l2_norms.unsqueeze(-1).unsqueeze(-1)
+                    )
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+            else:
+                new_vector[key] = new_vector[key] / num_tvs
+            new_vector[key] = new_vector[key].to(
+                device=original_device, non_blocking=True
+            )
+    return new_vector
+
+
+@auto_register_config
+class WUDIMerging(
+    LightningFabricMixin,
+    BaseAlgorithm,
+):
+    """
+    Whoever Started the Interference Should End It:  Guiding Data-Free Model Merging via Task Vectors
+    """
+
+    def __init__(
+        self,
+        iter_num: int,
+        exclude_keys: List[str] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+    def run(self, modelpool: BaseModelPool):
+        # load the pretrained model and the task vectors of all the finetuned models
+        with torch.no_grad():
+            pretrained_model = modelpool.load_pretrained_model()
+            task_vectors = []
+            for model_name in modelpool.model_names:
+                finetuned_model = modelpool.load_model(model_name)
+                task_vectors.append(
+                    state_dict_sub(
+                        finetuned_model.state_dict(), pretrained_model.state_dict()
+                    )
+                )
+                del finetuned_model  # free memory
+
+        merged_tv = wudi_merging(
+            task_vectors,
+            accelerator=self.fabric.device,
+            iter_num=self.iter_num,
+            exclude_keys=self.exclude_keys,
+        )
+
+        pretrained_model.load_state_dict(
+            state_dict_add(pretrained_model.state_dict(), merged_tv)
+        )
+
+        return pretrained_model

fusion_bench/mixins/clip_classification.py

@@ -113,11 +113,27 @@ class CLIPClassificationMixin(LightningFabricMixin):
         clip_model: Optional[CLIPModel] = None,
         task_names: Optional[List[str]] = None,
     ):
+        """
+        Initializes a zero-shot classification head.
+
+        This method constructs a zero-shot classification head by generating text embeddings for each class name using a set of templates.
+        These embeddings function as the weights of the classification layer. The method also extracts the `visual_projection` and `logit_scale`
+        from the provided CLIP model, which are necessary for calculating the final logits.
+
+        Args:
+            clip_processor (Optional[CLIPProcessor]): The processor for the CLIP model. If not provided, it is loaded from the model pool.
+            clip_model (Optional[CLIPModel]): The CLIP model to use. If not provided, a pretrained model is loaded from the model pool.
+            task_names (Optional[List[str]]): A list of task names to set up the classification head for. If not provided, all models in the model pool will be used.
+        """
         self.whether_setup_zero_shot_classification_head = True
+        # load clip model if not provided
         if clip_model is None:
             if self.modelpool.has_pretrained:
                 clip_model = self.modelpool.load_clip_model("_pretrained_")
             else:
+                log.warning(
+                    f"No pretrained CLIP model found, using the model from the model pool: {self.modelpool.model_names[0]}."
+                )
                 clip_model = self.modelpool.load_clip_model(
                     self.modelpool.model_names[0]
                 )
@@ -166,16 +182,20 @@ class CLIPClassificationMixin(LightningFabricMixin):
         image_embeds: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """
-        Compute the logits of the images for a given task.
+        Computes the classification logits for a batch of images for a specific task.
+
+        This method performs zero-shot classification by calculating the cosine similarity between image and text embeddings.
+        The image embeddings are obtained from the provided vision model, and the text embeddings (zero-shot weights) are pre-computed for the task.
+        The similarity scores are then scaled by the CLIP model's `logit_scale` to produce the final logits.
 
         Args:
-            module (Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"]): The module to compute the logits.
-            images (torch.Tensor): The images to compute the logits.
-            task (str): The task to compute the logits.
-            image_embeds (Optional[torch.Tensor]): The precomputed image embeddings. If None, the image embeddings will be computed.
+            module (Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"]): The vision encoder part of the CLIP model.
+            images (torch.Tensor): A batch of images to classify.
+            task (str): The name of the classification task.
+            image_embeds (Optional[torch.Tensor]): Pre-computed image embeddings. If provided, the method skips the image encoding step.
 
         Returns:
-            torch.Tensor: The logits of the images.
+            torch.Tensor: A tensor of logits for each image, with shape (batch_size, num_classes).
         """
         text_embeds = self.zeroshot_weights[task]
 

fusion_bench/mixins/lightning_fabric.py

@@ -100,6 +100,10 @@ class LightningFabricMixin:
             self.setup_lightning_fabric(getattr(self, "config", DictConfig({})))
         return self._fabric_instance
 
+    @fabric.setter
+    def fabric(self, instance: L.Fabric):
+        self._fabric_instance = instance
+
     @property
     def log_dir(self):
         """

fusion_bench/mixins/serialization.py

@@ -4,8 +4,9 @@ from copy import deepcopy
 from functools import wraps
 from inspect import Parameter, _ParameterKind
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import Dict, Mapping, Optional, Union
 
+from bidict import MutableBidict, bidict
 from omegaconf import DictConfig, OmegaConf
 
 from fusion_bench.constants import FUSION_BENCH_VERSION
@@ -15,12 +16,33 @@ from fusion_bench.utils.instantiate_utils import set_print_function_call
 log = logging.getLogger(__name__)
 
 __all__ = [
-    "YAMLSerializationMixin",
     "auto_register_config",
+    "YAMLSerializationMixin",
     "BaseYAMLSerializable",
 ]
 
 
+def _set_attr(self, param_name: str, value):
+    """
+    Set an attribute on the object using the parameter name from config mapping.
+
+    This function looks up the corresponding attribute name for the given parameter
+    name using the object's _config_mapping, then sets that attribute to the
+    specified value. It also logs the operation for debugging purposes.
+
+    Args:
+        self: The object instance to set the attribute on.
+        param_name (str): The parameter name (config key) to map to an attribute.
+        value: The value to assign to the attribute.
+
+    Raises:
+        ValueError: If the parameter name is not found in the config mapping.
+    """
+    attr_name = self._config_mapping.inverse[param_name]
+    log.debug(f"set {attr_name} to {value}. Parameter name: {param_name}")
+    setattr(self, attr_name, value)
+
+
 def auto_register_config(cls):
     """
     Decorator to automatically register __init__ parameters in _config_mapping.
@@ -45,37 +67,16 @@ def auto_register_config(cls):
             functionality and modified __init__ behavior.
 
     Behavior:
-        - **Parameter Registration**: All non-variadic parameters (excluding *args, **kwargs)
+        - **Parameter Registration**: All non-variadic parameters (excluding ``*args``, ``**kwargs``)
           from the __init__ method are automatically added to _config_mapping
         - **Positional Arguments**: Handled in order and mapped to corresponding parameter names
         - **Keyword Arguments**: Processed after positional arguments, overriding any conflicts
         - **Default Values**: Applied when parameters are not provided via arguments
         - **Attribute Setting**: All parameters become instance attributes accessible via dot notation
 
-    Example:
-        ```python
-        @auto_register_config
-        class MyAlgorithm(BaseYAMLSerializable):
-            def __init__(self, learning_rate: float = 0.001, batch_size: int = 32, model_name: str = "default"):
-                super().__init__()
-
-        # All instantiation methods work automatically:
-        algo1 = MyAlgorithm(0.01, 64)  # positional args
-        algo2 = MyAlgorithm(learning_rate=0.01, model_name="bert")  # keyword args
-        algo3 = MyAlgorithm(0.01, batch_size=128, model_name="gpt")  # mixed args
-
-        # Attributes are automatically set and can be serialized:
-        print(algo1.learning_rate)  # 0.01
-        print(algo1.batch_size)     # 64
-        print(algo1.model_name)     # "default" (from default value)
-
-        config = algo1.config
-        # DictConfig({'_target_': 'MyAlgorithm', 'learning_rate': 0.01, 'batch_size': 64, 'model_name': 'default'})
-        ```
-
     Note:
         - The decorator wraps the original __init__ method while preserving its signature for IDE support
-        - Parameters with *args or **kwargs signatures are ignored during registration
+        - Parameters with ``*args`` or ``**kwargs`` signatures are ignored during registration
         - The attributes are auto-registered, then the original __init__ method is called,
         - Type hints, method name, and other metadata are preserved using functools.wraps
         - This decorator is designed to work seamlessly with the YAML serialization system
@@ -89,7 +90,10 @@ def auto_register_config(cls):
 
     # Auto-register parameters in _config_mapping
     if not "_config_mapping" in cls.__dict__:
-        cls._config_mapping = deepcopy(getattr(cls, "_config_mapping", {}))
+        cls._config_mapping = deepcopy(getattr(cls, "_config_mapping", bidict()))
+    if not isinstance(cls._config_mapping, bidict):
+        cls._config_mapping = bidict(cls._config_mapping)
+
     registered_parameters = tuple(cls._config_mapping.values())
 
     for param_name in list(sig.parameters.keys())[1:]:  # Skip 'self'
@@ -102,9 +106,9 @@ def auto_register_config(cls):
         ) and (param_name not in registered_parameters):
             cls._config_mapping[param_name] = param_name
 
+    @wraps(original_init)
     def __init__(self, *args, **kwargs):
-        nonlocal original_init, registered_parameters
-
+        log.debug(f"set attributes for {self.__class__.__name__} in {cls.__name__}")
         # auto-register the attributes based on the signature
         sig = inspect.signature(original_init)
         param_names = list(sig.parameters.keys())[1:]  # Skip 'self'
@@ -117,29 +121,26 @@ def auto_register_config(cls):
                     _ParameterKind.VAR_POSITIONAL,
                     _ParameterKind.VAR_KEYWORD,
                 ]:
-                    setattr(self, param_name, arg_value)
+                    _set_attr(self, param_name, arg_value)
 
         # Handle keyword arguments and defaults
         for param_name in param_names:
-            if (
-                sig.parameters[param_name].kind
-                not in [
-                    _ParameterKind.VAR_POSITIONAL,
-                    _ParameterKind.VAR_KEYWORD,
-                ]
-            ) and (param_name not in registered_parameters):
+            if sig.parameters[param_name].kind not in [
+                _ParameterKind.VAR_POSITIONAL,
+                _ParameterKind.VAR_KEYWORD,
+            ]:
                 # Skip if already set by positional argument
                 param_index = param_names.index(param_name)
                 if param_index >= 0 and param_index < len(args):
                     continue
 
                 if param_name in kwargs:
-                    setattr(self, param_name, kwargs[param_name])
+                    _set_attr(self, param_name, kwargs[param_name])
                 else:
                     # Set default value if available and attribute doesn't exist
                     default_value = sig.parameters[param_name].default
                     if default_value is not Parameter.empty:
-                        setattr(self, param_name, default_value)
+                        _set_attr(self, param_name, default_value)
 
         # Call the original __init__
         result = original_init(self, *args, **kwargs)
@@ -152,33 +153,10 @@ def auto_register_config(cls):
 
 class YAMLSerializationMixin:
     _config_key: Optional[str] = None
-    _config_mapping: Dict[str, str] = {}
+    _config_mapping: MutableBidict[str, str] = bidict()
     R"""
     `_config_mapping` is a dictionary mapping the attribute names of the class to the config option names. This is used to convert the class to a DictConfig.
 
-    For example, if an algorithm class is defined as follows:
-    
-    ```python
-    class SomeModelFusionAlgorithm(BaseModelFusionAlgorithm):
-        hyper_parameter_1 = None
-        hyper_parameter_2 = None
-
-        _config_mapping = BaseModelFusionAlgorithm._config_mapping | {
-            "hyper_parameter_1" : "hyper_param_1",
-            "hyper_parameter_2" : "hyper_param_2",
-        }
-        def __init__(self, hyper_param_1: int, hyper_param_2: int):
-            self.hyper_parameter_1 = hyper_param_1
-            self.hyper_parameter_2 = hyper_param_2
-            super().__init__()
-    ```
-
-    The model pool will be converted to a DictConfig as follows:
-        
-    ```python
-    algorithm = SomeModelFusionAlgorithm(hyper_param_1=1, hyper_param_2=2)
-    ```
-
     >>> algorithm.config
         DictCOnfig({'_target_': 'SomeModelFusionAlgorithm', 'hyper_param_1': 1, 'hyper_param_2': 2})
 
@@ -197,17 +175,6 @@ class YAMLSerializationMixin:
         This property converts the model pool instance into a dictionary
         configuration, which can be used for serialization or other purposes.
 
-        Example:
-
-        ```python
-        model = SomeModelFusionAlgorithm(hyper_param_1=1, hyper_param_2=2)
-        config = model.config
-        print(config)
-        # DictConfig({'_target_': 'SomeModelFusionAlgorithm', 'hyper_param_1': 1, 'hyper_param_2': 2})
-        ```
-
-        This is useful for serializing the object to a YAML file or for debugging.
-
         Returns:
             DictConfig: The configuration of the model pool.
         """
@@ -272,16 +239,6 @@ class YAMLSerializationMixin:
                 serialization. This is how the attribute will appear in YAML output.
             value: The value to assign to the attribute.
 
-        Example:
-            ```python
-            model = BaseYAMLSerializable()
-            model.set_option("learning_rate", "lr", 0.001)
-
-            # This sets model.learning_rate = 0.001
-            # and maps it to "lr" in the config output
-            config = model.config
-            # config will contain: {"lr": 0.001, ...}
-            ```
         """
         setattr(self, attr_name, value)
         self._config_mapping[attr_name] = param_name

fusion_bench/modelpool/base_pool.py

@@ -277,7 +277,7 @@ class BaseModelPool(
         for dataset_name in self.test_dataset_names:
             yield self.load_test_dataset(dataset_name)
 
-    def save_model(self, model: nn.Module, path: str):
+    def save_model(self, model: nn.Module, path: str, *args, **kwargs):
         """
         Save the state dictionary of the model to the specified path.