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

fusion_bench/method/pwe_moe/clip_pwe_moe.py

@@ -16,14 +16,18 @@ from transformers import CLIPVisionModel
 from transformers.models.clip.modeling_clip import CLIPEncoderLayer
 from typing_extensions import override
 
-from fusion_bench.method.base_algorithm import BaseAlgorithm
+from fusion_bench import (
+    BaseAlgorithm,
+    auto_register_config,
+    print_parameters,
+    timeit_context,
+)
+from fusion_bench.dataset import CLIPDataset
 from fusion_bench.method.task_arithmetic import task_arithmetic_merge
 from fusion_bench.mixins.clip_classification import CLIPClassificationMixin
 from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
 from fusion_bench.modelpool import CLIPVisionModelPool
-from fusion_bench.utils import timeit_context
 from fusion_bench.utils.data import InfiniteDataLoader
-from fusion_bench.utils.parameters import print_parameters
 
 from .module import ParetoWeightEnsemblingModule
 from .utils import generate_simplex_grid
@@ -31,27 +35,13 @@ from .utils import generate_simplex_grid
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class PWEMoEAlgorithmForCLIP(
     BaseAlgorithm,
     SimpleProfilerMixin,
     CLIPClassificationMixin,
 ):
     modelpool: CLIPVisionModelPool = None
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "upscale_mlp": "upscale_mlp",
-        "upscale_attn": "upscale_attn",
-        "init_lambda": "init_lambda",
-        "router_hidden_layers": "router_hidden_layers",
-        "lr": "lr",
-        "num_steps": "num_steps",
-        "save_interval": "save_interval",
-        "alpha": "alpha",
-        "checkpoint_path": "checkpoint_path",
-        "eval_grid": "eval_grid",
-        "eval_grid_n": "eval_grid_n",
-        "eval_grid_m": "eval_grid_m",
-        "_dataloader_kwargs": "dataloader_kwargs",
-    }
 
     def __init__(
         self,
@@ -72,19 +62,6 @@ class PWEMoEAlgorithmForCLIP(
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.upscale_mlp = upscale_mlp
-        self.upscale_attn = upscale_attn
-        self.init_lambda = init_lambda
-        self.router_hidden_layers = router_hidden_layers
-        self.lr = lr
-        self.num_steps = num_steps
-        self.save_interval = save_interval
-        self.alpha = alpha
-        self.checkpoint_path = checkpoint_path
-        self.eval_grid = eval_grid
-        self.eval_grid_n = eval_grid_n
-        self.eval_grid_m = eval_grid_m
-        self._dataloader_kwargs = dataloader_kwargs
 
     @override
     def run(self, modelpool: CLIPVisionModelPool):
@@ -193,13 +170,14 @@ class PWEMoEAlgorithmForCLIP(
         Loads the datasets specified in the configuration.
         """
         train_datasets = {
-            dataset_name: self.modelpool.load_train_dataset(
-                dataset_name, self.clip_processor
+            dataset_name: CLIPDataset(
+                self.modelpool.load_train_dataset(dataset_name),
+                processor=self.clip_processor,
             )
             for dataset_name in self.modelpool.model_names
         }
         train_loaders = {
-            dataset_name: DataLoader(dataset, shuffle=True, **self._dataloader_kwargs)
+            dataset_name: DataLoader(dataset, shuffle=True, **self.dataloader_kwargs)
             for dataset_name, dataset in train_datasets.items()
         }
         train_loaders = {

fusion_bench/method/randes/modelsoup.py

@@ -5,9 +5,7 @@ import torch
 
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.parameters import count_parameters
-from fusion_bench.utils.state_dict_arithmetic import (
-    state_dict_mul,
-)
+from fusion_bench.utils.state_dict_arithmetic import state_dict_mul
 
 from .base_algorithm import SuperposedAlgorithmBase, compare_models
 

fusion_bench/method/regmean/clip_regmean.py

@@ -27,7 +27,7 @@ class RegMeanAlgorithmForCLIP(
 
     def __init__(self, *, dataloader_kwargs: DictConfig, **kwargs):
         super().__init__(**kwargs)
-        self._dataloader_kwargs = dataloader_kwargs
+        self.dataloader_kwargs = dataloader_kwargs
 
     def on_regmean_start(self):
         self.setup_zero_shot_classification_head()
@@ -60,7 +60,7 @@ class RegMeanAlgorithmForCLIP(
         # setup dataloader
         train_dataset = CLIPDataset(train_dataset, self.clip_processor)
         train_dataloader = DataLoader(
-            train_dataset, shuffle=True, **self._dataloader_kwargs
+            train_dataset, shuffle=True, **self.dataloader_kwargs
         )
         train_dataloader = self.fabric.setup_dataloaders(train_dataloader)
         model = self.fabric.setup(model)

fusion_bench/method/regmean/gpt2_regmean.py

@@ -15,7 +15,7 @@ from transformers import GPT2ForSequenceClassification, GPT2Model
 from transformers.data import default_data_collator
 from transformers.models.gpt2.modeling_gpt2 import Conv1D
 
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, auto_register_config
 from fusion_bench.utils import timeit_context
 
 from .regmean import RegMeanAlgorithm
@@ -23,22 +23,15 @@ from .regmean import RegMeanAlgorithm
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class RegMeanAlgorithmForGPT2(
-    RegMeanAlgorithm,
     LightningFabricMixin,
+    RegMeanAlgorithm,
 ):
     _include_module_type = [Conv1D]
     classifiers = {}
-    _config_mapping = RegMeanAlgorithm._config_mapping | {
-        "cache_dir": "cache_dir",
-        "batch_size": "batch_size",
-        "num_workers": "num_workers",
-    }
 
     def __init__(self, cache_dir: str, batch_size: int, num_workers: int, **kwargs):
-        self.cache_dir = cache_dir
-        self.batch_size = batch_size
-        self.num_workers = num_workers
         super().__init__(**kwargs)
 
     def on_regmean_start(self):

fusion_bench/method/regmean/regmean.py

@@ -13,7 +13,7 @@ from torch import Tensor, nn
 from tqdm.autonotebook import tqdm
 
 from fusion_bench.method import BaseAlgorithm
-from fusion_bench.mixins import SimpleProfilerMixin
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 
 log = logging.getLogger(__name__)
@@ -280,14 +280,9 @@ def regmean_merging(
     return merged_params
 
 
+@auto_register_config
 class RegMeanAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
     _include_module_type = [nn.Linear]
-    _config_mapping = {
-        "num_regmean_examples": "num_regmean_examples",
-        "exclude_param_names_regex": "exclude_param_names_regex",
-        "reduce_non_diagonal_ratio": "reduce_non_diagonal_ratio",
-        "weight_transpose": "weight_transpose",
-    }
 
     def __init__(
         self,
@@ -298,10 +293,6 @@ class RegMeanAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         weight_transpose: bool,
         **kwargs,
     ):
-        self.num_regmean_examples = num_regmean_examples
-        self.exclude_param_names_regex = exclude_param_names_regex
-        self.reduce_non_diagonal_ratio = reduce_non_diagonal_ratio
-        self.weight_transpose = weight_transpose
         super().__init__(**kwargs)
 
     def run(self, modelpool: BaseModelPool, **kwargs):

fusion_bench/method/regmean_plusplus/__init__.py

@@ -1,3 +1,3 @@
 # flake8: noqa F401
 from .clip_regmean_plusplus import RegMeanAlgorithmForCLIPPlusPlus
-from .regmean_plusplus import RegMeanAlgorithmPlusPlus
+from .regmean_plusplus import RegMeanAlgorithmPlusPlus

fusion_bench/method/regmean_plusplus/clip_regmean_plusplus.py

@@ -28,7 +28,7 @@ class RegMeanAlgorithmForCLIPPlusPlus(
 
     def __init__(self, *, dataloader_kwargs: DictConfig, **kwargs):
         super().__init__(**kwargs)
-        self._dataloader_kwargs = dataloader_kwargs
+        self.dataloader_kwargs = dataloader_kwargs
 
     def on_regmean_start(self):
         self.setup_zero_shot_classification_head()
@@ -125,27 +125,26 @@ class RegMeanAlgorithmForCLIPPlusPlus(
 
             param_dict = {}
             for name, param in model_to_merge_state_dict.items():
-                if name.startswith("vision_model.embeddings") or name.startswith("vision_model.pre_layrnorm"):
+                if name.startswith("vision_model.embeddings") or name.startswith(
+                    "vision_model.pre_layrnorm"
+                ):
                     param_dict[name] = param
 
             for param_name in param_dict.keys():
-                models_to_merge_param_dict[param_name].append(
-                    param_dict[param_name]
-                )
+                models_to_merge_param_dict[param_name].append(param_dict[param_name])
 
         # merge the parameters of the embedding layer
         merged_params_dict = {}
         for param_name, param_list in models_to_merge_param_dict.items():
             merged_params_dict[param_name] = torch.stack(param_list).mean(dim=0)
-                
+
         return merged_params_dict
-        
-    
+
     def get_input_for_first_layer(self, model: nn.Module, train_dataset):
         # setup dataloader
         train_dataset = CLIPDataset(train_dataset, self.clip_processor)
         train_dataloader = DataLoader(
-            train_dataset, shuffle=True, **self._dataloader_kwargs
+            train_dataset, shuffle=True, **self.dataloader_kwargs
         )
         train_dataloader = self.fabric.setup_dataloaders(train_dataloader)
         model = self.fabric.setup(model)
@@ -157,9 +156,9 @@ class RegMeanAlgorithmForCLIPPlusPlus(
             image_embeds = model.vision_model.embeddings(images)
             image_embeds = model.vision_model.pre_layrnorm(image_embeds)
             image_embeds = image_embeds.detach().cpu()
-            
+
             return image_embeds
-        
+
         num_computed_examples = 0
         num_regmean_examples = self.num_regmean_examples
 
@@ -169,24 +168,32 @@ class RegMeanAlgorithmForCLIPPlusPlus(
                 break
             batches_input.append(compute_input(model, batch))
             num_computed_examples += batch[0].size(0)
-        
+
         return batches_input
 
     def get_layers(self, model: nn.Module):
         return model.vision_model.encoder.layers
-    
-    def update_merged_params_dict(self, merged_params_dict, new_merged_params, layer_idx):
+
+    def update_merged_params_dict(
+        self, merged_params_dict, new_merged_params, layer_idx
+    ):
         for key, value in new_merged_params.items():
             key = f"vision_model.encoder.layers.{layer_idx}.{key}"
             merged_params_dict[key] = value
 
         return merged_params_dict
-    
-    def layer_batches_forward(self, layer: nn.Module, batches_input: List[Tensor]) -> Tensor:
+
+    def layer_batches_forward(
+        self, layer: nn.Module, batches_input: List[Tensor]
+    ) -> Tensor:
         batches_output = []
         for batch in batches_input:
             device = next(layer.parameters()).device
             batch = batch.to(device)
-            logits = layer(batch, attention_mask=None, causal_attention_mask=None)[0].detach().cpu()
+            logits = (
+                layer(batch, attention_mask=None, causal_attention_mask=None)[0]
+                .detach()
+                .cpu()
+            )
             batches_output.append(logits)
         return batches_output

fusion_bench/method/regmean_plusplus/regmean_plusplus.py

@@ -81,13 +81,11 @@ def regmean_params_merge(
     reduce_non_diagonal_ratio: float = 1.0,
     weight_transpose: bool = True,
     module_name: str = "",
-    device = "cpu"
+    device="cpu",
 ):
     # two lists with length num_models_to_merge
     param_multiplied_results, module_regmean_weights_list = [], []
-    for model_idx, module_regmean_weights in enumerate(
-        param_regmean_list
-    ):
+    for model_idx, module_regmean_weights in enumerate(param_regmean_list):
         # reduce non-diagonal elements
         module_regmean_weights = reduce_non_diagonal_elements(
             regmean_weights=module_regmean_weights,
@@ -113,9 +111,7 @@ def regmean_params_merge(
     sum_param_multiplied_results = sum(param_multiplied_results)
 
     # get the inverse matrix
-    inv_sum_module_regmean_weights = torch.inverse(
-        sum_module_regmean_weights
-    )
+    inv_sum_module_regmean_weights = torch.inverse(sum_module_regmean_weights)
     # merge parameters with regmean
     merged_param = torch.matmul(
         inv_sum_module_regmean_weights, sum_param_multiplied_results
@@ -158,15 +154,19 @@ def merging_with_regmean_weights(
                     device = param_value_list[model_idx].device
 
                     # Tensor, shape (hidden_dim, hidden_dim)
-                    module_regmean_weights = model_to_merge_regmean_weights[module_name].to(device)
+                    module_regmean_weights = model_to_merge_regmean_weights[
+                        module_name
+                    ].to(device)
                     module_regmean_weights_list.append(module_regmean_weights)
 
-                merged_params[param_name] = regmean_params_merge(param_weight_list=param_value_list,
-                                                                 param_regmean_list=module_regmean_weights_list,
-                                                                 reduce_non_diagonal_ratio=reduce_non_diagonal_ratio,
-                                                                 weight_transpose=weight_transpose,
-                                                                 module_name=module_name,
-                                                                 device=device)
+                merged_params[param_name] = regmean_params_merge(
+                    param_weight_list=param_value_list,
+                    param_regmean_list=module_regmean_weights_list,
+                    reduce_non_diagonal_ratio=reduce_non_diagonal_ratio,
+                    weight_transpose=weight_transpose,
+                    module_name=module_name,
+                    device=device,
+                )
 
                 merged_by_regmean = True
         # use average merging for parameters whose names are not end with ".weight" or not in Linear module
@@ -205,7 +205,9 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
             modelpool = BaseModelPool(modelpool)
         self.modelpool = modelpool
         device = "cuda:0" if torch.cuda.is_available() else "cpu"
-        models_to_merge_dict = {name: model.to(device) for name, model in modelpool.named_models()}
+        models_to_merge_dict = {
+            name: model.to(device) for name, model in modelpool.named_models()
+        }
         self.on_regmean_start()
 
         # initialize the merged models as the pretrained model
@@ -213,7 +215,9 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
         merged_params_dict = {}
 
         # 1. merge embedding layer
-        merged_embedding_dict = self.merge_embedding_layer(models_to_merge_dict=models_to_merge_dict)
+        merged_embedding_dict = self.merge_embedding_layer(
+            models_to_merge_dict=models_to_merge_dict
+        )
         merged_model.load_state_dict(merged_embedding_dict, strict=False)
 
         with torch.no_grad():
@@ -223,12 +227,13 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
                 self.profile("computing first layer input"),
             ):
                 batches_input_dict = defaultdict(list)
-                for name in tqdm(models_to_merge_dict.keys(), desc="computing input for first layer"):
+                for name in tqdm(
+                    models_to_merge_dict.keys(), desc="computing input for first layer"
+                ):
                     dataset = modelpool.load_train_dataset(name)
-                    
+
                     batches_input_dict[name] = self.get_input_for_first_layer(
-                        merged_model,
-                        dataset
+                        merged_model, dataset
                     )
 
             # 2. iteratively merge layer by layer with regmean algorithm
@@ -240,9 +245,9 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
                 models_to_merge_layers_dict[name] = self.get_layers(model)
 
             param_names_to_merge = None
-            for layer_idx, backbone_layer in tqdm(enumerate(backbone_layers), 
-                                                  desc="merging layers", 
-                                                  total=num_layers):
+            for layer_idx, backbone_layer in tqdm(
+                enumerate(backbone_layers), desc="merging layers", total=num_layers
+            ):
                 # dictionary of list, where key is the parameter name,
                 # value is a list of the corresponding parameters of all the models that need to be merged
                 models_to_merge_param_dict = defaultdict(list)
@@ -263,16 +268,19 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
                                 "exclude_param_names_regex", []
                             ),
                         )
-                
+
                     for param_name in param_names_to_merge:
                         models_to_merge_param_dict[param_name].append(
                             param_dict[param_name]
                         )
 
                     linear_modules_to_merge = get_modules_to_merge(
-                        model=layer_to_merge, include_module_types=self._include_module_type
+                        model=layer_to_merge,
+                        include_module_types=self._include_module_type,
                     )
-                    assert len(linear_modules_to_merge) > 0, "No linear modules to merge"
+                    assert (
+                        len(linear_modules_to_merge) > 0
+                    ), "No linear modules to merge"
 
                     # 2.1. compute regmean weights for each model
                     with (
@@ -288,12 +296,19 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
 
                         module_subset = get_param_names_to_merge(
                             input_param_names=list(param_dict.keys()),
-                            exclude_param_names_regex=self.exclude_param_names_regex
+                            exclude_param_names_regex=self.exclude_param_names_regex,
                         )
-                        module_subset = [name.replace(".weight", "").replace(".bias", "") for name in module_subset]
+                        module_subset = [
+                            name.replace(".weight", "").replace(".bias", "")
+                            for name in module_subset
+                        ]
                         module_subset = list(set(module_subset))
-                        regmean_weights = {module_name: regmean_weights[module_name] for module_name in module_subset if module_name in regmean_weights}
-                        
+                        regmean_weights = {
+                            module_name: regmean_weights[module_name]
+                            for module_name in module_subset
+                            if module_name in regmean_weights
+                        }
+
                         models_to_merge_regmean_weights_list.append(regmean_weights)
 
                 # 2.2. merge parameters with regmean weights
@@ -318,21 +333,22 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
                     self.profile("forwarding next layer"),
                 ):
                     if layer_idx < num_layers - 1:
-                        backbone_layer.load_state_dict(merged_layer_params, strict=False)
+                        backbone_layer.load_state_dict(
+                            merged_layer_params, strict=False
+                        )
                         batches_output_dict = defaultdict(list)
                         for name in models_to_merge_dict.keys():
                             batches_output_dict[name] = self.layer_batches_forward(
-                                backbone_layer, 
-                                batches_input_dict[name]
+                                backbone_layer, batches_input_dict[name]
                             )
                         batches_input_dict = batches_output_dict
-                
+
             # 3. load state dict to the merged model
             merged_model.load_state_dict(merged_params_dict, strict=False)
 
         self.print_profile_summary()
         return merged_model
-    
+
     def merge_embedding_layer(self, models_to_merge_dict: Dict[str, nn.Module]):
         """
         Merge the embedding layer of the model with the merged model.
@@ -345,10 +361,12 @@ class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
 
     def get_layers(self, model: nn.Module):
         raise NotImplementedError
-    
-    def update_merged_params_dict(self, merged_params_dict, new_merged_params, layer_idx):
+
+    def update_merged_params_dict(
+        self, merged_params_dict, new_merged_params, layer_idx
+    ):
         raise NotImplementedError
-    
+
     def layer_batches_forward(self, layer: nn.Module, batches_input: List[Tensor]):
         raise NotImplementedError
 

fusion_bench/method/simple_average.py

@@ -6,7 +6,7 @@ import torch
 from torch import nn
 
 from fusion_bench.method.base_algorithm import BaseAlgorithm
-from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils import LazyStateDict
 from fusion_bench.utils.state_dict_arithmetic import (
@@ -59,24 +59,20 @@ def simple_average(
         return state_dict_avg(modules)
 
 
+@auto_register_config
 class SimpleAverageAlgorithm(
-    BaseAlgorithm,
     SimpleProfilerMixin,
+    BaseAlgorithm,
 ):
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "show_pbar": "show_pbar",
-    }
-
-    def __init__(self, show_pbar: bool = False):
+    def __init__(self, show_pbar: bool = False, **kwargs):
         """
         Args:
             show_pbar (bool): If True, shows a progress bar during model loading and merging. Default is False.
         """
-        super().__init__()
-        self.show_pbar = show_pbar
+        super().__init__(**kwargs)
 
     @torch.no_grad()
-    def run(self, modelpool: Union[BaseModelPool, Dict[str, nn.Module]]):
+    def run(self, modelpool: Union[BaseModelPool, Dict[str, nn.Module]]) -> nn.Module:
         """
         Fuse the models in the given model pool using simple averaging.
 
@@ -124,13 +120,13 @@ class SimpleAverageAlgorithm(
         if isinstance(forward_model, LazyStateDict):
             # if the model is a LazyStateDict, convert it to an empty module
             forward_model = forward_model.meta_module.to_empty(
-                device=(
-                    "cpu"
-                    if forward_model._torch_dtype is None
-                    else forward_model._torch_dtype
-                )
+                device=forward_model._device
             )
-        forward_model.load_state_dict(sd)
+        result = forward_model.load_state_dict(sd, strict=False)
+        if result.unexpected_keys:
+            raise ValueError(f"Unexpected keys in state dict: {result.unexpected_keys}")
+        if result.missing_keys:
+            log.warning(f"Missing keys in state dict: {result.missing_keys}")
         # print profile report and log the merged models
         self.print_profile_summary()
         log.info(f"merged {len(merged_model_names)} models:")

fusion_bench/method/slerp/slerp.py

@@ -1,10 +1,13 @@
 import logging
+from typing import Any, Dict
 
 import torch
+from torch import nn
 from typing_extensions import override
 
 from fusion_bench.method import BaseAlgorithm
 from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.utils.type import StateDictType
 
 from .slerp_utils import slerp
 
@@ -18,7 +21,7 @@ def slerp_on_state_dicts(
     *,
     DOT_THRESHOLD: float = 0.9995,
     epsilon: float = 1e-8,
-):
+) -> StateDictType:
     """
     Perform spherical linear interpolation (slerp) on the state dictionaries of two models.
 
@@ -72,7 +75,7 @@ class SlerpMergeAlgorithm(BaseAlgorithm):
         super().__init__()
 
     @override
-    def run(self, modelpool: BaseModelPool):
+    def run(self, modelpool: BaseModelPool) -> nn.Module:
         """
         Run the SlerpMergeAlgorithm on the given model pool.