fusion-bench 0.2.18__tar.gz → 0.2.20__tar.gz

{fusion_bench-0.2.18 → fusion_bench-0.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.18
+Version: 0.2.20
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License
@@ -72,6 +72,14 @@ FusionBench is a benchmark suite designed to evaluate the performance of various
 
 Projects based on FusionBench and news from the community (descending order of date. If you have any work based on FusionBench, please feel free to let us know, we are willing to add it to the list. :partying_face:):
 
+<details>
+  <summary>The-Hai Nguyen, Dang Huu-Tien, Takeshi Suzuki, and Le-Minh Nguyen. RegMean++: Enhancing Effectiveness and Generalization of Regression Mean for Model Merging. Aug, 2025. https://www.arxiv.org/abs/2508.03121</summary>
+
+Regression Mean (RegMean), an approach that formulates model merging as a linear regression problem, aims to find the optimal weights for each linear layer in the merge model by minimizing the discrepancy in predictions between the merge and candidate models. RegMean provides a precise closed-form solution for the merging problem; therefore, it offers explainability and computational efficiency. However, RegMean merges each linear layer independently, overlooking how the features and information in the earlier layers propagate through the layers and influence the final prediction in the merge model. In this paper, we introduce RegMean++, a simple yet effective alternative to RegMean, that explicitly incorporates both intra- and cross-layer dependencies between merge models' layers into RegMean's objective. By accounting for these dependencies, RegMean++ better captures the behaviors of the merge model. Extensive experiments demonstrate that RegMean++ consistently outperforms RegMean across diverse settings, including in-domain (ID) and out-of-domain (OOD) generalization, sequential merging, large-scale tasks, and robustness under several types of distribution shifts. Furthermore, RegMean++ achieves competitive or state-of-the-art performance compared to various recent advanced model merging methods. 
+
+  <img width="1000" alt="image" src="docs/algorithms/images/regmean_vs_regmean_plusplus.png">
+</details>
+
 <details>
   <summary>Hao Mark Chen, et al. FW-Merging: Scaling Model Merging with Frank-Wolfe Optimization. Mar 2025. https://arxiv.org/abs/2503.12649</summary>
 

{fusion_bench-0.2.18 → fusion_bench-0.2.20}/README.md

@@ -21,6 +21,14 @@ FusionBench is a benchmark suite designed to evaluate the performance of various
 
 Projects based on FusionBench and news from the community (descending order of date. If you have any work based on FusionBench, please feel free to let us know, we are willing to add it to the list. :partying_face:):
 
+<details>
+  <summary>The-Hai Nguyen, Dang Huu-Tien, Takeshi Suzuki, and Le-Minh Nguyen. RegMean++: Enhancing Effectiveness and Generalization of Regression Mean for Model Merging. Aug, 2025. https://www.arxiv.org/abs/2508.03121</summary>
+
+Regression Mean (RegMean), an approach that formulates model merging as a linear regression problem, aims to find the optimal weights for each linear layer in the merge model by minimizing the discrepancy in predictions between the merge and candidate models. RegMean provides a precise closed-form solution for the merging problem; therefore, it offers explainability and computational efficiency. However, RegMean merges each linear layer independently, overlooking how the features and information in the earlier layers propagate through the layers and influence the final prediction in the merge model. In this paper, we introduce RegMean++, a simple yet effective alternative to RegMean, that explicitly incorporates both intra- and cross-layer dependencies between merge models' layers into RegMean's objective. By accounting for these dependencies, RegMean++ better captures the behaviors of the merge model. Extensive experiments demonstrate that RegMean++ consistently outperforms RegMean across diverse settings, including in-domain (ID) and out-of-domain (OOD) generalization, sequential merging, large-scale tasks, and robustness under several types of distribution shifts. Furthermore, RegMean++ achieves competitive or state-of-the-art performance compared to various recent advanced model merging methods. 
+
+  <img width="1000" alt="image" src="docs/algorithms/images/regmean_vs_regmean_plusplus.png">
+</details>
+
 <details>
   <summary>Hao Mark Chen, et al. FW-Merging: Scaling Model Merging with Frank-Wolfe Optimization. Mar 2025. https://arxiv.org/abs/2503.12649</summary>
 

fusion_bench-0.2.20/fusion_bench/__init__.py

@@ -0,0 +1,26 @@
+# ███████╗██╗   ██╗███████╗██╗ ██████╗ ███╗   ██╗      ██████╗ ███████╗███╗   ██╗ ██████╗██╗  ██╗
+# ██╔════╝██║   ██║██╔════╝██║██╔═══██╗████╗  ██║      ██╔══██╗██╔════╝████╗  ██║██╔════╝██║  ██║
+# █████╗  ██║   ██║███████╗██║██║   ██║██╔██╗ ██║█████╗██████╔╝█████╗  ██╔██╗ ██║██║     ███████║
+# ██╔══╝  ██║   ██║╚════██║██║██║   ██║██║╚██╗██║╚════╝██╔══██╗██╔══╝  ██║╚██╗██║██║     ██╔══██║
+# ██║     ╚██████╔╝███████║██║╚██████╔╝██║ ╚████║      ██████╔╝███████╗██║ ╚████║╚██████╗██║  ██║
+# ╚═╝      ╚═════╝ ╚══════╝╚═╝ ╚═════╝ ╚═╝  ╚═══╝      ╚═════╝ ╚══════╝╚═╝  ╚═══╝ ╚═════╝╚═╝  ╚═╝
+# flake8: noqa: F401
+from . import (
+    constants,
+    dataset,
+    method,
+    metrics,
+    mixins,
+    modelpool,
+    models,
+    optim,
+    programs,
+    taskpool,
+    tasks,
+    utils,
+)
+from .method import BaseAlgorithm, BaseModelFusionAlgorithm
+from .modelpool import BaseModelPool
+from .models import separate_io
+from .taskpool import BaseTaskPool
+from .utils import parse_dtype, print_parameters, timeit_context

fusion_bench-0.2.20/fusion_bench/constants/banner.py

@@ -0,0 +1,12 @@
+FUSION_BENCH_BANNER = (
+    ""
+    + "███████╗██╗   ██╗███████╗██╗ ██████╗ ███╗   ██╗      ██████╗ ███████╗███╗   ██╗ ██████╗██╗  ██╗\n"
+    + "██╔════╝██║   ██║██╔════╝██║██╔═══██╗████╗  ██║      ██╔══██╗██╔════╝████╗  ██║██╔════╝██║  ██║\n"
+    + "█████╗  ██║   ██║███████╗██║██║   ██║██╔██╗ ██║█████╗██████╔╝█████╗  ██╔██╗ ██║██║     ███████║\n"
+    + "██╔══╝  ██║   ██║╚════██║██║██║   ██║██║╚██╗██║╚════╝██╔══██╗██╔══╝  ██║╚██╗██║██║     ██╔══██║\n"
+    + "██║     ╚██████╔╝███████║██║╚██████╔╝██║ ╚████║      ██████╔╝███████╗██║ ╚████║╚██████╗██║  ██║\n"
+    + "╚═╝      ╚═════╝ ╚══════╝╚═╝ ╚═════╝ ╚═╝  ╚═══╝      ╚═════╝ ╚══════╝╚═╝  ╚═══╝ ╚═════╝╚═╝  ╚═╝\n"
+)
+
+if __name__ == "__main__":
+    print(FUSION_BENCH_BANNER)

{fusion_bench-0.2.18 → fusion_bench-0.2.20}/fusion_bench/method/__init__.py

@@ -41,6 +41,7 @@ _import_structure = {
         "FisherMergingAlgorithmForGPT2",
     ],
     "regmean": ["RegMeanAlgorithmForCLIP", "RegMeanAlgorithmForGPT2"],
+    "regmean_plusplus": ["RegMeanAlgorithmForCLIPPlusPlus"],
     "adamerging": [
         "CLIPTaskWiseAdaMergingAlgorithm",
         "CLIPLayerWiseAdaMergingAlgorithm",
@@ -195,6 +196,7 @@ if TYPE_CHECKING:
     )
     from .rankone_moe import CLIPRankOneMoEAlgorithm, RankOneMoEAlgorithm
     from .regmean import RegMeanAlgorithmForCLIP, RegMeanAlgorithmForGPT2
+    from .regmean_plusplus import RegMeanAlgorithmForCLIPPlusPlus
     from .simple_average import SimpleAverageAlgorithm
     from .slerp import SlerpMergeAlgorithm
     from .smile_upscaling import (

{fusion_bench-0.2.18 → fusion_bench-0.2.20}/fusion_bench/method/linear/simple_average_for_llama.py

@@ -1,4 +1,5 @@
-from typing import Optional
+from copy import deepcopy
+from typing import TYPE_CHECKING, Optional
 
 from typing_extensions import override
 
@@ -6,6 +7,11 @@ from fusion_bench import timeit_context
 from fusion_bench.method.base_algorithm import BaseAlgorithm
 from fusion_bench.method.simple_average import SimpleAverageAlgorithm
 from fusion_bench.modelpool import CausalLMBackbonePool, CausalLMPool
+from fusion_bench.utils.pylogger import getRankZeroLogger
+from omegaconf import flag_override
+from fusion_bench.utils import instantiate
+
+log = getRankZeroLogger(__name__)
 
 
 class SimpleAverageForLlama(BaseAlgorithm):
@@ -26,12 +32,19 @@ class SimpleAverageForLlama(BaseAlgorithm):
 
     _config_mapping = BaseAlgorithm._config_mapping | {
         "merge_backbone": "merge_backbone",
+        "show_pbar": "show_pbar",
     }
 
-    def __init__(self, merge_backbone: bool, model_save_path: Optional[str] = None):
+    def __init__(
+        self,
+        merge_backbone: bool,
+        model_save_path: Optional[str] = None,
+        show_pbar: bool = False,
+    ):
         super().__init__()
         self.merge_backbone = merge_backbone
         self.model_save_path = model_save_path
+        self.show_pbar = show_pbar
 
     @override
     def run(self, modelpool: CausalLMPool):
@@ -40,12 +53,24 @@ class SimpleAverageForLlama(BaseAlgorithm):
 
         if self.merge_backbone:
             assert modelpool.has_pretrained
-            backbone_modelpool = CausalLMBackbonePool(**modelpool.config)
+            log.info(
+                "Merging backbone of the model pool, use CausalLMBackbonePool instead of CausalLMPool."
+            )
+            modelpool_config = deepcopy(modelpool.config)
+            with flag_override(modelpool_config, "allow_objects", True):
+                modelpool_config._target_ = (
+                    "fusion_bench.modelpool.causal_lm.CausalLMBackbonePool"
+                )
+            backbone_modelpool = instantiate(modelpool_config)
             model = modelpool.load_model("_pretrained_")
-            backbone_model = SimpleAverageAlgorithm().run(backbone_modelpool)
+            backbone_model = SimpleAverageAlgorithm(show_pbar=self.show_pbar).run(
+                backbone_modelpool
+            )
             model.model.layers = backbone_model
         else:
-            model = SimpleAverageAlgorithm().run()
+            model = SimpleAverageAlgorithm(show_pbar=self.show_pbar).run(
+                modelpool=modelpool
+            )
 
         if self.model_save_path is not None:
             with timeit_context(f"Saving the model to {self.model_save_path}"):

fusion_bench-0.2.20/fusion_bench/method/regmean_plusplus/__init__.py

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

fusion_bench-0.2.20/fusion_bench/method/regmean_plusplus/clip_regmean_plusplus.py

@@ -0,0 +1,192 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, cast  # noqa: F401
+
+import torch
+import torch.utils.data
+from omegaconf import DictConfig
+from torch import Tensor, nn
+from torch.nn.modules import Module
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+from fusion_bench.dataset.clip_dataset import CLIPDataset
+from fusion_bench.mixins import CLIPClassificationMixin
+
+from .regmean_plusplus import RegMeanAlgorithmPlusPlus
+
+log = logging.getLogger(__name__)
+
+
+class RegMeanAlgorithmForCLIPPlusPlus(
+    RegMeanAlgorithmPlusPlus,
+    CLIPClassificationMixin,
+):
+    _config_mapping = {
+        "_dataloader_kwargs": "dataloader_kwargs",
+    }
+
+    def __init__(self, *, dataloader_kwargs: DictConfig, **kwargs):
+        super().__init__(**kwargs)
+        self._dataloader_kwargs = dataloader_kwargs
+
+    def on_regmean_start(self):
+        self.setup_zero_shot_classification_head()
+
+    def compute_logits(self, module, batch, task: str) -> Tensor:
+        images, _ = batch
+        text_embeds = self.zeroshot_weights[task]
+
+        image_embeds = module(images)[1]
+        image_embeds = self.visual_projection(image_embeds)
+
+        # normalize embeddings
+        image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+
+        # cosine similarity
+        logits_per_text = (
+            torch.matmul(text_embeds, image_embeds.t()) * self.logit_scale_exp
+        )
+        logits_per_image = logits_per_text.t()
+
+        return logits_per_image
+
+    def get_regmean_weights(
+        self,
+        model_name: str,
+        layer: Module,
+        batches_input: List[Tensor],
+        linear_modules_to_merge: Dict[str, Module],
+    ):
+        layer = self.fabric.setup(layer)
+
+        def compute_regmean_weights(module_name: str):
+            """
+            compute the regmean weights, a hook function to deal with each module's input
+            :param module_name: str, module name
+            :return:
+            """
+
+            def hook(module: nn.Module, input: tuple, output: torch.Tensor):
+                # Tensor, shape (batch_size, sequence_length, hidden_dim)
+                x = cast(Tensor, input[0]).detach()
+                batch_num_actual_examples = x.shape[0]
+                # Tensor, shape (batch_size * sequence_length, hidden_dim)
+                x = x.reshape(-1, x.shape[-1])
+                # Tensor, shape (hidden_dim, hidden_dim)
+                xtx = torch.matmul(x.transpose(0, 1), x)
+                # store the averaged weights in regmean_weights
+                if module_name not in regmean_weights.keys():
+                    regmean_weights[module_name] = xtx / x.shape[0]
+                    num_computed_examples[module_name] = x.shape[0]
+                    num_actual_examples[module_name] = batch_num_actual_examples
+                else:
+                    regmean_weights[module_name] = (
+                        regmean_weights[module_name]
+                        * num_computed_examples[module_name]
+                        + xtx
+                    ) / (num_computed_examples[module_name] + x.shape[0])
+                    num_computed_examples[module_name] += x.shape[0]
+                    num_actual_examples[module_name] += batch_num_actual_examples
+
+            return hook
+
+        handles = []
+        # dictionary, regmean matrices for each linear module inputs
+        regmean_weights = {}
+        # dictionary, number of examples (multiplied the sequence length) used for computing regmean matrices
+        num_computed_examples = {}
+        # dictionary, number of actual examples used for computing regmean matrices
+        num_actual_examples = {}
+
+        for module_name, linear_module_to_merge in linear_modules_to_merge.items():
+            # register a hook in the forward process
+            handle = linear_module_to_merge.register_forward_hook(
+                compute_regmean_weights(module_name=module_name)
+            )
+            handles.append(handle)
+        _ = self.layer_batches_forward(layer, batches_input)
+
+        # remove the added hook
+        for handle in handles:
+            handle.remove()
+
+        for module_name in regmean_weights.keys():
+            regmean_weights[module_name] = regmean_weights[module_name].detach().cpu()
+
+        return regmean_weights
+
+    def merge_embedding_layer(self, models_to_merge_dict: Dict[str, nn.Module]):
+        models_to_merge_param_dict = defaultdict(list)
+
+        # get the parameters of the embedding layer from each model
+        for model_to_merge in models_to_merge_dict.values():
+            model_to_merge_state_dict = model_to_merge.state_dict()
+
+            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"):
+                    param_dict[name] = param
+
+            for param_name in param_dict.keys():
+                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_dataloader = self.fabric.setup_dataloaders(train_dataloader)
+        model = self.fabric.setup(model)
+
+        def compute_input(model, batch):
+            images, _ = batch
+
+            images = images.to(model.device)
+            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
+
+        batches_input = []
+        for batch in train_dataloader:
+            if num_computed_examples >= num_regmean_examples:
+                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):
+        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:
+        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()
+            batches_output.append(logits)
+        return batches_output