fusion-bench 0.2.11__tar.gz → 0.2.12__tar.gz

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.11
+Version: 0.2.12
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License
@@ -45,6 +45,7 @@ Requires-Dist: rich
 Requires-Dist: scipy
 Requires-Dist: h5py
 Requires-Dist: pytest
+Dynamic: license-file
 
 <div align='center'>
 
@@ -69,6 +70,18 @@ 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):
 
+<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>
+
+Model merging has emerged as a promising approach for multi-task learning (MTL), offering a data-efficient alternative to conventional fine-tuning. However, with the rapid development of the open-source AI ecosystem and the increasing availability of fine-tuned foundation models, existing model merging methods face two key limitations: (i) They are primarily designed for in-house fine-tuned models, making them less adaptable to diverse model sources with partially unknown model and task information, (ii) They struggle to scale effectively when merging numerous model checkpoints. To address these challenges, we formulate model merging as a constrained optimization problem and introduce a novel approach: Frank-Wolfe Merging (FW-Merging). Inspired by Frank-Wolfe optimization, our approach iteratively selects the most relevant model in the pool to minimize a linear approximation of the objective function and then executes a local merging similar to the Frank-Wolfe update. The objective function is designed to capture the desired behavior of the target-merged model, while the fine-tuned candidate models define the constraint set. More importantly, FW-Merging serves as an orthogonal technique for existing merging methods, seamlessly integrating with them to further enhance accuracy performance. Our experiments show that FW-Merging scales across diverse model sources, remaining stable with 16 irrelevant models and improving by 15.3% with 16 relevant models on 20 CV tasks, while maintaining constant memory overhead, unlike the linear overhead of data-informed merging methods. Compared with the state-of-the-art approaches, FW-Merging surpasses the data-free merging method by 32.8% and outperforms the data-informed Adamerging by 8.39% when merging 20 ViT models.
+</details>
+
+<details>
+  <summary>Daniel Marczak, et al. No Task Left Behind: Isotropic Model Merging with Common and Task-Specific Subspaces. Feb 2025. https://arxiv.org/abs/2502.04959</summary>
+
+  Model merging integrates the weights of multiple task-specific models into a single multi-task model. Despite recent interest in the problem, a significant performance gap between the combined and single-task models remains. In this paper, we investigate the key characteristics of task matrices -- weight update matrices applied to a pre-trained model -- that enable effective merging. We show that alignment between singular components of task-specific and merged matrices strongly correlates with performance improvement over the pre-trained model. Based on this, we propose an isotropic merging framework that flattens the singular value spectrum of task matrices, enhances alignment, and reduces the performance gap. Additionally, we incorporate both common and task-specific subspaces to further improve alignment and performance. Our proposed approach achieves state-of-the-art performance across multiple scenarios, including various sets of tasks and model scales. This work advances the understanding of model merging dynamics, offering an effective methodology to merge models without requiring additional training. 
+</details>
+
 <details>
   <summary>Anke Tang, et al. Merging Models on the Fly Without Retraining: A Sequential Approach to Scalable Continual Model Merging. Jan 2025. https://arxiv.org/pdf/2501.09522</summary>
 

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/README.md

@@ -21,6 +21,18 @@ 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):
 
+<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>
+
+Model merging has emerged as a promising approach for multi-task learning (MTL), offering a data-efficient alternative to conventional fine-tuning. However, with the rapid development of the open-source AI ecosystem and the increasing availability of fine-tuned foundation models, existing model merging methods face two key limitations: (i) They are primarily designed for in-house fine-tuned models, making them less adaptable to diverse model sources with partially unknown model and task information, (ii) They struggle to scale effectively when merging numerous model checkpoints. To address these challenges, we formulate model merging as a constrained optimization problem and introduce a novel approach: Frank-Wolfe Merging (FW-Merging). Inspired by Frank-Wolfe optimization, our approach iteratively selects the most relevant model in the pool to minimize a linear approximation of the objective function and then executes a local merging similar to the Frank-Wolfe update. The objective function is designed to capture the desired behavior of the target-merged model, while the fine-tuned candidate models define the constraint set. More importantly, FW-Merging serves as an orthogonal technique for existing merging methods, seamlessly integrating with them to further enhance accuracy performance. Our experiments show that FW-Merging scales across diverse model sources, remaining stable with 16 irrelevant models and improving by 15.3% with 16 relevant models on 20 CV tasks, while maintaining constant memory overhead, unlike the linear overhead of data-informed merging methods. Compared with the state-of-the-art approaches, FW-Merging surpasses the data-free merging method by 32.8% and outperforms the data-informed Adamerging by 8.39% when merging 20 ViT models.
+</details>
+
+<details>
+  <summary>Daniel Marczak, et al. No Task Left Behind: Isotropic Model Merging with Common and Task-Specific Subspaces. Feb 2025. https://arxiv.org/abs/2502.04959</summary>
+
+  Model merging integrates the weights of multiple task-specific models into a single multi-task model. Despite recent interest in the problem, a significant performance gap between the combined and single-task models remains. In this paper, we investigate the key characteristics of task matrices -- weight update matrices applied to a pre-trained model -- that enable effective merging. We show that alignment between singular components of task-specific and merged matrices strongly correlates with performance improvement over the pre-trained model. Based on this, we propose an isotropic merging framework that flattens the singular value spectrum of task matrices, enhances alignment, and reduces the performance gap. Additionally, we incorporate both common and task-specific subspaces to further improve alignment and performance. Our proposed approach achieves state-of-the-art performance across multiple scenarios, including various sets of tasks and model scales. This work advances the understanding of model merging dynamics, offering an effective methodology to merge models without requiring additional training. 
+</details>
+
 <details>
   <summary>Anke Tang, et al. Merging Models on the Fly Without Retraining: A Sequential Approach to Scalable Continual Model Merging. Jan 2025. https://arxiv.org/pdf/2501.09522</summary>
 

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/fusion_bench/compat/method/__init__.py

@@ -20,7 +20,7 @@ class AlgorithmFactory:
         # model merging methods
         "clip_task_wise_adamerging": ".adamerging.clip_task_wise_adamerging.CLIPTaskWiseAdaMergingAlgorithm",
         "clip_layer_wise_adamerging": ".adamerging.clip_layer_wise_adamerging.CLIPLayerWiseAdaMergingAlgorithm",
-        "clip_layer_wise_adamerging_doge_ta": ".DOGE_TA.clip_layer_wise_adamerging.CLIPLayerWiseAdaMergingAlgorithm",
+        "clip_layer_wise_adamerging_doge_ta": ".doge_ta.clip_layer_wise_adamerging.CLIPLayerWiseAdaMergingAlgorithm",
         "singular_projection_merging": "fusion_bench.method.smile_upscaling.singular_projection_merging.SingularProjectionMergingAlgorithm",
         "clip_layer_wise_adamerging_surgery": ".surgery.clip_layer_wise_adamerging_surgery.CLIPLayerWiseAdaMergingSurgeryAlgorithm",
         # plug-and-play model merging methods

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/fusion_bench/dataset/fer2013.py

@@ -7,7 +7,6 @@ def load_fer2013(path: str = "clip-benchmark/wds_fer2013", split: str = "train")
     dataset = dataset.rename_columns({"jpg": "image", "cls": "label"})
     return dataset
 
-
 if __name__ == "__main__":
     dataset = load_fer2013(split="test")
     print(dataset)

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/fusion_bench/method/__init__.py

@@ -53,7 +53,7 @@ _import_structure = {
         "PWEMoExactParetoOptimalForCLIP",
     ],
     "ada_svd": ["AdaSVDMergingForCLIPVisionModel"],
-    "DOGE_TA": ["DOGE_TA_Algorithm"],
+    "doge_ta": ["DOGE_TA_Algorithm"],
     "task_singular_vector": ["TaskSingularVectorMerging"],
     "isotropic_merging": [
         "ISO_C_Merge",  # alias
@@ -128,7 +128,7 @@ if TYPE_CHECKING:
     from .dare import DareSimpleAverage, DareTaskArithmetic, DareTiesMerging
     from .dawe import DataAdaptiveWeightEnsemblingForCLIP
     from .depth_upscaling import DepthUpscalingAlgorithm, DepthUpscalingForLlama
-    from .DOGE_TA import DOGE_TA_Algorithm
+    from .doge_ta import DOGE_TA_Algorithm
     from .dummy import DummyAlgorithm
     from .ensemble import (
         MaxModelPredictorAlgorithm,

{fusion_bench-0.2.11/fusion_bench/method/DOGE_TA → fusion_bench-0.2.12/fusion_bench/method/adamerging}/clip_layer_wise_adamerging.py

@@ -9,7 +9,7 @@ fusion_bench \
     modelpool=clip-vit-base-patch32_TA8 \
     taskpool=clip-vit-classification_TA8 \
     fabric.loggers.root_dir=outputs/logs/ViT-B-32 \
-    fabric.loggers.name=clip_layer_wise_adamerging_adam
+    fabric.loggers.name=clip_layer_wise_adamerging_adamerging
 ```
 """
 

fusion_bench-0.2.12/fusion_bench/method/doge_ta/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa F401
+from .doge_ta import DOGE_TA_Algorithm

{fusion_bench-0.2.11/fusion_bench/method/adamerging → fusion_bench-0.2.12/fusion_bench/method/doge_ta}/clip_layer_wise_adamerging.py

@@ -9,7 +9,7 @@ fusion_bench \
     modelpool=clip-vit-base-patch32_TA8 \
     taskpool=clip-vit-classification_TA8 \
     fabric.loggers.root_dir=outputs/logs/ViT-B-32 \
-    fabric.loggers.name=clip_layer_wise_adamerging_adam
+    fabric.loggers.name=clip_layer_wise_adamerging_adamerging
 ```
 """
 

fusion_bench-0.2.11/fusion_bench/method/DOGE_TA/DOGE_TA.py → fusion_bench-0.2.12/fusion_bench/method/doge_ta/doge_ta.py

@@ -7,7 +7,7 @@ Example Usage:
 
 ```bash
 fusion_bench \
-    method=DOGE_TA/DOGE_TA \
+    method=doge_ta/doge_ta \
     modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8_model_only \
     taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
 

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/fusion_bench/method/opcm/opcm.py

@@ -15,7 +15,7 @@ from tqdm.auto import tqdm
 from transformers import CLIPVisionModel
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
 from fusion_bench.taskpool import CLIPVisionModelTaskPool
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.json import load_from_json, save_to_json
@@ -31,6 +31,7 @@ if TYPE_CHECKING:
 class OPCMForCLIP(
     BaseAlgorithm,
     LightningFabricMixin,
+    SimpleProfilerMixin,
 ):
     def __init__(
         self,
@@ -64,7 +65,8 @@ class OPCMForCLIP(
             L.seed_everything(self.seed)
         accelerator = self.fabric.device
 
-        pretrained_model = modelpool.load_pretrained_model()
+        with self.profile("loading model"):
+            pretrained_model = modelpool.load_pretrained_model()
 
         model_names = modelpool.model_names
         if self.shuffle_order:
@@ -83,15 +85,17 @@ class OPCMForCLIP(
             )
 
         # get the average model
-        merged_model = modelpool.load_model(model_names[0])
+        with self.profile("loading model"):
+            merged_model = modelpool.load_model(model_names[0])
 
         if self.evaluate_on_every_step:
-            self.taskpool._is_setup = False
-            self.taskpool._test_datasets = DictConfig(
-                {model_names[0]: self._test_datasets[model_names[0]]}
-            )
-            report = self.taskpool.evaluate(deepcopy(merged_model))
-            save_to_json(report, Path(self.log_dir) / "report_0.json")
+            with self.profile("evaluating model"):
+                self.taskpool._is_setup = False
+                self.taskpool._test_datasets = DictConfig(
+                    {model_names[0]: self._test_datasets[model_names[0]]}
+                )
+                report = self.taskpool.evaluate(deepcopy(merged_model))
+                save_to_json(report, Path(self.log_dir) / "report_0.json")
 
         self.avg_task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
         self.all_task_vector_norm = [self.avg_task_vector_norm]
@@ -113,90 +117,95 @@ class OPCMForCLIP(
             enumerate(model_names[1:]), desc="Processing models"
         ):
             model_idx += 1
-            task_model = modelpool.load_model(model_name)
+            with self.profile("loading model"):
+                task_model = modelpool.load_model(model_name)
 
-            self.all_task_vector_norm.append(
-                get_task_vector_norm(task_model, pretrained_model)
-            )
-            self.avg_task_vector_norm = np.mean(self.all_task_vector_norm)
-            self.fabric.log(
-                "model/task_vector_norm", self.all_task_vector_norm[-1], step=model_idx
-            )
-            self.fabric.log(
-                "model/avg_task_vector_norm", self.avg_task_vector_norm, step=model_idx
-            )
+            with self.profile("merging model"):
+                self.all_task_vector_norm.append(
+                    get_task_vector_norm(task_model, pretrained_model)
+                )
+                self.avg_task_vector_norm = np.mean(self.all_task_vector_norm)
+                self.fabric.log(
+                    "model/task_vector_norm", self.all_task_vector_norm[-1], step=model_idx
+                )
+                self.fabric.log(
+                    "model/avg_task_vector_norm", self.avg_task_vector_norm, step=model_idx
+                )
 
-            self.lambda_t = 1  # temporary value
-
-            for module_name, module in tqdm(
-                list(merged_model.named_modules()),
-                desc=f"Processing {model_name}",
-                leave=False,
-            ):
-                if not is_leaf_module(module):
-                    continue
-
-                if isinstance(module, nn.Linear):
-                    module.weight.data = self.merge_linear_weights(
-                        module.weight,
-                        pretrained_model.get_submodule(module_name).weight,
-                        task_model.get_submodule(module_name).weight,
-                        param_name=".".join([module_name, "weight"]),
-                        alpha=self.alpha,
-                        accelerator=accelerator,
-                    )
-                    if module.bias is not None:
-                        module.bias.data = self.merge_other_parameters(
-                            module.bias,
-                            pretrained_model.get_submodule(module_name).bias,
-                            task_model.get_submodule(module_name).bias,
-                            param_name=".".join([module_name, "bias"]),
+                self.lambda_t = 1  # temporary value
+
+                for module_name, module in tqdm(
+                    list(merged_model.named_modules()),
+                    desc=f"Processing {model_name}",
+                    leave=False,
+                ):
+                    if not is_leaf_module(module):
+                        continue
+
+                    if isinstance(module, nn.Linear):
+                        module.weight.data = self.merge_linear_weights(
+                            module.weight,
+                            pretrained_model.get_submodule(module_name).weight,
+                            task_model.get_submodule(module_name).weight,
+                            param_name=".".join([module_name, "weight"]),
+                            alpha=self.alpha,
                             accelerator=accelerator,
                         )
-                else:
-                    for param_name, param in module.named_parameters():
-                        param.data = self.merge_other_parameters(
-                            merged_W=param,
-                            pretrained_W=pretrained_model.get_submodule(
-                                module_name
-                            ).get_parameter(param_name),
-                            task_W=task_model.get_submodule(module_name).get_parameter(
-                                param_name
-                            ),
-                            param_name=".".join([module_name, param_name]),
-                            accelerator=accelerator,
-                        )
-
-            task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
-            self.lambda_t *= task_vector_norm / self.avg_task_vector_norm
-            for param_name, param in merged_model.named_parameters():
-                param.data = pretrained_model.get_parameter(param_name) + (
-                    param - pretrained_model.get_parameter(param_name)
-                ) * (self.avg_task_vector_norm / task_vector_norm)
-            self.fabric.log("model/lambda_t", self.lambda_t, step=model_idx)
-            self.fabric.log(
-                "empirical/lambda_t", np.sqrt(model_idx + 1), step=model_idx
-            )
-            self.previous_lambda_t = self.lambda_t
-            self.lambda_t = None
+                        if module.bias is not None:
+                            module.bias.data = self.merge_other_parameters(
+                                module.bias,
+                                pretrained_model.get_submodule(module_name).bias,
+                                task_model.get_submodule(module_name).bias,
+                                param_name=".".join([module_name, "bias"]),
+                                accelerator=accelerator,
+                            )
+                    else:
+                        for param_name, param in module.named_parameters():
+                            param.data = self.merge_other_parameters(
+                                merged_W=param,
+                                pretrained_W=pretrained_model.get_submodule(
+                                    module_name
+                                ).get_parameter(param_name),
+                                task_W=task_model.get_submodule(module_name).get_parameter(
+                                    param_name
+                                ),
+                                param_name=".".join([module_name, param_name]),
+                                accelerator=accelerator,
+                            )
+
+                task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
+                self.lambda_t *= task_vector_norm / self.avg_task_vector_norm
+                for param_name, param in merged_model.named_parameters():
+                    param.data = pretrained_model.get_parameter(param_name) + (
+                        param - pretrained_model.get_parameter(param_name)
+                    ) * (self.avg_task_vector_norm / task_vector_norm)
+                self.fabric.log("model/lambda_t", self.lambda_t, step=model_idx)
+                self.fabric.log(
+                    "empirical/lambda_t", np.sqrt(model_idx + 1), step=model_idx
+                )
+                self.previous_lambda_t = self.lambda_t
+                self.lambda_t = None
 
-            self.fabric.log(
-                "model/merged_task_vector_norm",
-                get_task_vector_norm(merged_model, pretrained_model),
-                step=model_idx,
-            )
+                self.fabric.log(
+                    "model/merged_task_vector_norm",
+                    get_task_vector_norm(merged_model, pretrained_model),
+                    step=model_idx,
+                )
 
             if self.save_on_every_step:
-                self.save_merged_model(merged_model, model_idx)
+                with self.profile("saving model"):
+                    self.save_merged_model(merged_model, model_idx)
 
             if self.evaluate_on_every_step:
-                self.taskpool._is_setup = False
-                self.taskpool._test_datasets = DictConfig(
-                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
-                )
-                report = self.taskpool.evaluate(deepcopy(merged_model))
-                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
+                with self.profile("evaluating model"):
+                    self.taskpool._is_setup = False
+                    self.taskpool._test_datasets = DictConfig(
+                        {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
+                    )
+                    report = self.taskpool.evaluate(deepcopy(merged_model))
+                    save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
 
+        self.print_profile_summary()
         return merged_model
 
     def save_merged_model(self, merged_model: CLIPVisionModel, step: int):
@@ -227,7 +236,7 @@ class OPCMForCLIP(
         split_rank = (s.cumsum(dim=0) / s.sum() > alpha).float().argmax().item()
 
         projected_task_tv = u.T @ task_tv @ v
-        projected_task_tv.diag().fill_(0)
+        projected_task_tv.diagonal().fill_(0)
 
         projected_task_tv[:split_rank, :split_rank] = 0
 

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/fusion_bench/method/opcm/task_arithmetic.py

@@ -15,7 +15,7 @@ from tqdm.auto import tqdm
 from transformers import CLIPVisionModel
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
 from fusion_bench.taskpool import CLIPVisionModelTaskPool
 from fusion_bench.utils.json import load_from_json, save_to_json
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_sub
@@ -24,7 +24,11 @@ if TYPE_CHECKING:
     from torch.utils.tensorboard import SummaryWriter
 
 
-class ContinualTaskArithmeticForCLIP(BaseAlgorithm, LightningFabricMixin):
+class ContinualTaskArithmeticForCLIP(
+    BaseAlgorithm,
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+):
     def __init__(
         self,
         scaling_factor: float,
@@ -79,32 +83,42 @@ class ContinualTaskArithmeticForCLIP(BaseAlgorithm, LightningFabricMixin):
         for model_idx, model_name in tqdm(
             enumerate(model_names), desc="Processing models"
         ):
-            task_model = modelpool.load_model(model_name)
+            with self.profile("loading model"):
+                task_model = modelpool.load_model(model_name)
 
-            for param_name, param in task_model.named_parameters():
-                if not param.requires_grad:
-                    continue
+            with self.profile("merging model"):
+                for param_name, param in task_model.named_parameters():
+                    if not param.requires_grad:
+                        continue
 
-                task_param = param
-                merged_param = merged_model.get_parameter(param_name)
-                pretrained_param = pretrained_model.get_parameter(param_name)
+                    task_param = param
+                    merged_param = merged_model.get_parameter(param_name)
+                    pretrained_param = pretrained_model.get_parameter(param_name)
 
-                new_param = merged_param + self.scaling_factor * (
-                    task_param - pretrained_param
-                )
-                merged_model.get_parameter(param_name).data = new_param
+                    new_param = merged_param + self.scaling_factor * (
+                        task_param - pretrained_param
+                    )
+                    merged_model.get_parameter(param_name).data = new_param
 
             if self.save_on_every_step:
-                self.save_merged_model(merged_model, model_idx)
+                with self.profile("saving model"):
+                    self.save_merged_model(merged_model, model_idx)
 
             if self.evaluate_on_every_step:
-                self.taskpool._is_setup = False
-                self.taskpool._test_datasets = DictConfig(
-                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
-                )
-                report = self.taskpool.evaluate(deepcopy(merged_model))
-                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
-
+                with self.profile("evaluating model"):
+                    self.taskpool._is_setup = False
+                    self.taskpool._test_datasets = DictConfig(
+                        {
+                            n: self._test_datasets[n]
+                            for n in model_names[: model_idx + 1]
+                        }
+                    )
+                    report = self.taskpool.evaluate(deepcopy(merged_model))
+                    save_to_json(
+                        report, Path(self.log_dir) / f"report_{model_idx}.json"
+                    )
+
+        self.print_profile_summary()
         return merged_model
 
     def save_merged_model(self, merged_model: CLIPVisionModel, step: int):

{fusion_bench-0.2.11 → fusion_bench-0.2.12}/fusion_bench/method/opcm/ties_merging.py

@@ -20,7 +20,7 @@ from fusion_bench.method.ties_merging.ties_merging_utils import (
     ties_merging,
     vector_to_state_dict,
 )
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
 from fusion_bench.taskpool import CLIPVisionModelTaskPool
 from fusion_bench.utils.json import load_from_json, save_to_json
 from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_sub
@@ -29,7 +29,11 @@ if TYPE_CHECKING:
     from torch.utils.tensorboard import SummaryWriter
 
 
-class ContinualTiesMergingForCLIP(BaseAlgorithm, LightningFabricMixin):
+class ContinualTiesMergingForCLIP(
+    BaseAlgorithm,
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+):
     def __init__(
         self,
         scaling_factor: float,
@@ -84,68 +88,83 @@ class ContinualTiesMergingForCLIP(BaseAlgorithm, LightningFabricMixin):
             )
 
         # get the average model
-        pretrained_model = modelpool.load_pretrained_model()
+        with self.profile("loading model"):
+            pretrained_model = modelpool.load_pretrained_model()
         merged_model = deepcopy(pretrained_model)
 
         for model_idx, model_name in tqdm(
             enumerate(model_names), desc="Processing models"
         ):
-            task_model = modelpool.load_model(model_name)
+            with self.profile("loading model"):
+                task_model = modelpool.load_model(model_name)
 
-            task_vector = state_dict_sub(
-                task_model.state_dict(),
-                pretrained_model.state_dict(),
-            )
-            if model_idx == 0:
-                # if is the first model, the merged task vector is equal to the task vector
-                ties_merging_state_dict = task_vector
-            else:
-                # if is not the first model, we need to merge the task vector with the previous merged task vector
-                merged_tv = state_dict_sub(
-                    merged_model.state_dict(),
+            with self.profile("merging model"):
+                task_vector = state_dict_sub(
+                    task_model.state_dict(),
                     pretrained_model.state_dict(),
                 )
-                tv_flat_checks = torch.vstack(
-                    [
-                        state_dict_to_vector(merged_tv, remove_keys=self.remove_keys),
-                        state_dict_to_vector(task_vector, remove_keys=self.remove_keys),
-                    ]
-                )
-                # perform the TIES merging
-                ties_merging_tv = ties_merging(
-                    tv_flat_checks,
-                    reset_thresh=self.threshold,
-                    merge_func=self.merge_func,
-                )
-                # convert the merged task vector back to a state dict
-                ties_merging_state_dict = vector_to_state_dict(
-                    ties_merging_tv,
-                    merged_model.state_dict(),
-                    remove_keys=self.remove_keys,
-                )
-
-            for param_name, param in task_model.named_parameters():
-                if not param.requires_grad:
-                    continue
-
-                merged_param = merged_model.get_parameter(param_name)
-                new_param = (
-                    merged_param
-                    + self.scaling_factor * ties_merging_state_dict[param_name]
-                )
-                merged_model.get_parameter(param_name).data = new_param
+                if model_idx == 0:
+                    # if is the first model, the merged task vector is equal to the task vector
+                    ties_merging_state_dict = task_vector
+                else:
+                    # if is not the first model, we need to merge the task vector with the previous merged task vector
+                    merged_tv = state_dict_sub(
+                        merged_model.state_dict(),
+                        pretrained_model.state_dict(),
+                    )
+                    tv_flat_checks = torch.vstack(
+                        [
+                            state_dict_to_vector(
+                                merged_tv, remove_keys=self.remove_keys
+                            ),
+                            state_dict_to_vector(
+                                task_vector, remove_keys=self.remove_keys
+                            ),
+                        ]
+                    )
+                    # perform the TIES merging
+                    ties_merging_tv = ties_merging(
+                        tv_flat_checks,
+                        reset_thresh=self.threshold,
+                        merge_func=self.merge_func,
+                    )
+                    # convert the merged task vector back to a state dict
+                    ties_merging_state_dict = vector_to_state_dict(
+                        ties_merging_tv,
+                        merged_model.state_dict(),
+                        remove_keys=self.remove_keys,
+                    )
+
+                for param_name, param in task_model.named_parameters():
+                    if not param.requires_grad:
+                        continue
+
+                    merged_param = merged_model.get_parameter(param_name)
+                    new_param = (
+                        merged_param
+                        + self.scaling_factor * ties_merging_state_dict[param_name]
+                    )
+                    merged_model.get_parameter(param_name).data = new_param
 
             if self.save_on_every_step:
-                self.save_merged_model(merged_model, model_idx)
+                with self.profile("saving model"):
+                    self.save_merged_model(merged_model, model_idx)
 
             if self.evaluate_on_every_step:
-                self.taskpool._is_setup = False
-                self.taskpool._test_datasets = DictConfig(
-                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
-                )
-                report = self.taskpool.evaluate(deepcopy(merged_model))
-                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
-
+                with self.profile("evaluating model"):
+                    self.taskpool._is_setup = False
+                    self.taskpool._test_datasets = DictConfig(
+                        {
+                            n: self._test_datasets[n]
+                            for n in model_names[: model_idx + 1]
+                        }
+                    )
+                    report = self.taskpool.evaluate(deepcopy(merged_model))
+                    save_to_json(
+                        report, Path(self.log_dir) / f"report_{model_idx}.json"
+                    )
+
+        self.print_profile_summary()
         return merged_model
 
     def save_merged_model(self, merged_model: CLIPVisionModel, step: int):