fusion-bench 0.2.8__tar.gz → 0.2.10__tar.gz

{fusion_bench-0.2.8 → fusion_bench-0.2.10}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: fusion_bench
-Version: 0.2.8
+Version: 0.2.10
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License
@@ -70,22 +70,21 @@ 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>Hongling Zheng, Li Shen, Anke Tang, Yong Luo et al. Learn From Model Beyond Fine-Tuning: A Survey. has been accepted for publication in Nature Machine Intelligence. Nov, 2024. https://arxiv.org/abs/2310.08184</summary>
-
-  > Foundation models (FM) have demonstrated remarkable performance across a wide range of tasks (especially in the fields
-  of natural language processing and computer vision), primarily attributed to their ability to comprehend instructions and access
-  extensive, high-quality data. This not only showcases their current effectiveness but also sets a promising trajectory towards the
-  development of artificial general intelligence. Unfortunately, due to multiple constraints, the raw data of the model used for large model
-  training are often inaccessible, so the use of end-to-end models for downstream tasks has become a new research trend, which we call
-  Learn From Model (LFM) in this article. LFM focuses on the research, modification, and design of FM based on the model interface,
-  so as to better understand the model structure and weights (in a black box environment), and to generalize the model to downstream
-  tasks. The study of LFM techniques can be broadly categorized into five major areas: model tuning, model distillation, model reuse,
-  meta learning and model editing. Each category encompasses a repertoire of methods and strategies that aim to enhance the
-  capabilities and performance of FM. This paper gives a comprehensive review of the current methods based on FM from the
-  perspective of LFM, in order to help readers better understand the current research status and ideas. To conclude, we summarize the
-  survey by highlighting several critical areas for future exploration and addressing open issues that require further attention from the
-  research community. The relevant papers we investigated in this article can be accessed at
-  https://github.com/ruthless-man/Awesome-Learn-from-Model.
+  <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>
+
+Deep model merging represents an emerging research direction that combines multiple fine-tuned models to harness their specialized capabilities across different tasks and domains. Current model merging techniques focus on merging all available models simultaneously, with weight interpolation-based methods being the predominant approaches. However, these conventional approaches are not well-suited for scenarios where models become available sequentially, and they often suffer from high memory requirements and potential interference between tasks. In this study, we propose a training-free projection-based continual merging method that processes models sequentially through orthogonal projections of weight matrices and adaptive scaling mechanisms. Our method operates by projecting new parameter updates onto subspaces orthogonal to existing merged parameter updates while using an adaptive scaling mechanism to maintain stable parameter distances, enabling efficient sequential integration of task-specific knowledge. Our approach maintains constant memory complexity to the number of models, minimizes interference between tasks through orthogonal projections, and retains the performance of previously merged models through adaptive task vector scaling. Extensive experiments on CLIP-ViT models demonstrate that our method achieves a 5-8% average accuracy improvement while maintaining robust performance in different task orderings.
+</details>
+
+<details>
+  <summary>Yongxian Wei, et al. Modeling Multi-Task Model Merging as Adaptive Projective Gradient Descent. Jan 2025. https://arxiv.org/abs/2501.01230</summary>
+
+Merging multiple expert models offers a promising approach for performing multi-task learning without accessing their original data. Existing methods attempt to alleviate task conflicts by sparsifying task vectors or promoting orthogonality among them. However, they overlook the fundamental requirement of model merging: ensuring the merged model performs comparably to task-specific models on respective tasks. We find these methods inevitably discard task-specific information that, while causing conflicts, is crucial for performance. Based on our findings, we frame model merging as a constrained optimization problem (i.e., minimizing the gap between the merged model and individual models, subject to the constraint of retaining shared knowledge) and solve it via adaptive projective gradient descent. Specifically, we align the merged model with individual models by decomposing and reconstituting the loss function, alleviating conflicts through data-free optimization of task vectors. To retain shared knowledge, we optimize this objective by projecting gradients within a shared subspace spanning all tasks. Moreover, we view merging coefficients as adaptive learning rates and propose a task-aware, training-free strategy. Experiments show that our plug-andplay approach consistently outperforms previous methods, achieving state-of-the-art results across diverse architectures and tasks in both vision and NLP domains. Our code is available here.
+</details>
+
+<details>
+  <summary>Hongling Zheng, Li Shen, Anke Tang, Yong Luo et al. Learn From Model Beyond Fine-Tuning: A Survey. Nature Machine Intelligence. Jan, 2025. https://www.nature.com/articles/s42256-024-00961-0</summary>
+
+  > Foundation models (FM) have demonstrated remarkable performance across a wide range of tasks (especially in the fields of natural language processing and computer vision), primarily attributed to their ability to comprehend instructions and access extensive, high-quality data. This not only showcases their current effectiveness but also sets a promising trajectory towards the development of artificial general intelligence. Unfortunately, due to multiple constraints, the raw data of the model used for large model training are often inaccessible, so the use of end-to-end models for downstream tasks has become a new research trend, which we call Learn From Model (LFM) in this article. LFM focuses on the research, modification, and design of FM based on the model interface, so as to better understand the model structure and weights (in a black box environment), and to generalize the model to downstream tasks. The study of LFM techniques can be broadly categorized into five major areas: model tuning, model distillation, model reuse, meta learning and model editing. Each category encompasses a repertoire of methods and strategies that aim to enhance the capabilities and performance of FM. This paper gives a comprehensive review of the current methods based on FM from the perspective of LFM, in order to help readers better understand the current research status and ideas. To conclude, we summarize the survey by highlighting several critical areas for future exploration and addressing open issues that require further attention from the research community. The relevant papers we investigated in this article can be accessed at https://github.com/ruthless-man/Awesome-Learn-from-Model.
 </details>
 
 <details>

{fusion_bench-0.2.8 → fusion_bench-0.2.10}/README.md

@@ -22,22 +22,21 @@ 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>Hongling Zheng, Li Shen, Anke Tang, Yong Luo et al. Learn From Model Beyond Fine-Tuning: A Survey. has been accepted for publication in Nature Machine Intelligence. Nov, 2024. https://arxiv.org/abs/2310.08184</summary>
-
-  > Foundation models (FM) have demonstrated remarkable performance across a wide range of tasks (especially in the fields
-  of natural language processing and computer vision), primarily attributed to their ability to comprehend instructions and access
-  extensive, high-quality data. This not only showcases their current effectiveness but also sets a promising trajectory towards the
-  development of artificial general intelligence. Unfortunately, due to multiple constraints, the raw data of the model used for large model
-  training are often inaccessible, so the use of end-to-end models for downstream tasks has become a new research trend, which we call
-  Learn From Model (LFM) in this article. LFM focuses on the research, modification, and design of FM based on the model interface,
-  so as to better understand the model structure and weights (in a black box environment), and to generalize the model to downstream
-  tasks. The study of LFM techniques can be broadly categorized into five major areas: model tuning, model distillation, model reuse,
-  meta learning and model editing. Each category encompasses a repertoire of methods and strategies that aim to enhance the
-  capabilities and performance of FM. This paper gives a comprehensive review of the current methods based on FM from the
-  perspective of LFM, in order to help readers better understand the current research status and ideas. To conclude, we summarize the
-  survey by highlighting several critical areas for future exploration and addressing open issues that require further attention from the
-  research community. The relevant papers we investigated in this article can be accessed at
-  https://github.com/ruthless-man/Awesome-Learn-from-Model.
+  <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>
+
+Deep model merging represents an emerging research direction that combines multiple fine-tuned models to harness their specialized capabilities across different tasks and domains. Current model merging techniques focus on merging all available models simultaneously, with weight interpolation-based methods being the predominant approaches. However, these conventional approaches are not well-suited for scenarios where models become available sequentially, and they often suffer from high memory requirements and potential interference between tasks. In this study, we propose a training-free projection-based continual merging method that processes models sequentially through orthogonal projections of weight matrices and adaptive scaling mechanisms. Our method operates by projecting new parameter updates onto subspaces orthogonal to existing merged parameter updates while using an adaptive scaling mechanism to maintain stable parameter distances, enabling efficient sequential integration of task-specific knowledge. Our approach maintains constant memory complexity to the number of models, minimizes interference between tasks through orthogonal projections, and retains the performance of previously merged models through adaptive task vector scaling. Extensive experiments on CLIP-ViT models demonstrate that our method achieves a 5-8% average accuracy improvement while maintaining robust performance in different task orderings.
+</details>
+
+<details>
+  <summary>Yongxian Wei, et al. Modeling Multi-Task Model Merging as Adaptive Projective Gradient Descent. Jan 2025. https://arxiv.org/abs/2501.01230</summary>
+
+Merging multiple expert models offers a promising approach for performing multi-task learning without accessing their original data. Existing methods attempt to alleviate task conflicts by sparsifying task vectors or promoting orthogonality among them. However, they overlook the fundamental requirement of model merging: ensuring the merged model performs comparably to task-specific models on respective tasks. We find these methods inevitably discard task-specific information that, while causing conflicts, is crucial for performance. Based on our findings, we frame model merging as a constrained optimization problem (i.e., minimizing the gap between the merged model and individual models, subject to the constraint of retaining shared knowledge) and solve it via adaptive projective gradient descent. Specifically, we align the merged model with individual models by decomposing and reconstituting the loss function, alleviating conflicts through data-free optimization of task vectors. To retain shared knowledge, we optimize this objective by projecting gradients within a shared subspace spanning all tasks. Moreover, we view merging coefficients as adaptive learning rates and propose a task-aware, training-free strategy. Experiments show that our plug-andplay approach consistently outperforms previous methods, achieving state-of-the-art results across diverse architectures and tasks in both vision and NLP domains. Our code is available here.
+</details>
+
+<details>
+  <summary>Hongling Zheng, Li Shen, Anke Tang, Yong Luo et al. Learn From Model Beyond Fine-Tuning: A Survey. Nature Machine Intelligence. Jan, 2025. https://www.nature.com/articles/s42256-024-00961-0</summary>
+
+  > Foundation models (FM) have demonstrated remarkable performance across a wide range of tasks (especially in the fields of natural language processing and computer vision), primarily attributed to their ability to comprehend instructions and access extensive, high-quality data. This not only showcases their current effectiveness but also sets a promising trajectory towards the development of artificial general intelligence. Unfortunately, due to multiple constraints, the raw data of the model used for large model training are often inaccessible, so the use of end-to-end models for downstream tasks has become a new research trend, which we call Learn From Model (LFM) in this article. LFM focuses on the research, modification, and design of FM based on the model interface, so as to better understand the model structure and weights (in a black box environment), and to generalize the model to downstream tasks. The study of LFM techniques can be broadly categorized into five major areas: model tuning, model distillation, model reuse, meta learning and model editing. Each category encompasses a repertoire of methods and strategies that aim to enhance the capabilities and performance of FM. This paper gives a comprehensive review of the current methods based on FM from the perspective of LFM, in order to help readers better understand the current research status and ideas. To conclude, we summarize the survey by highlighting several critical areas for future exploration and addressing open issues that require further attention from the research community. The relevant papers we investigated in this article can be accessed at https://github.com/ruthless-man/Awesome-Learn-from-Model.
 </details>
 
 <details>

fusion_bench-0.2.10/fusion_bench/__main__.py

@@ -0,0 +1,4 @@
+from fusion_bench.scripts.cli import main
+
+if __name__ == "__main__":
+    main()

{fusion_bench-0.2.8 → fusion_bench-0.2.10}/fusion_bench/dataset/fer2013.py

@@ -7,6 +7,7 @@ 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.8 → fusion_bench-0.2.10}/fusion_bench/method/__init__.py

@@ -9,7 +9,10 @@ _import_structure = {
     "base_algorithm": ["BaseModelFusionAlgorithm", "BaseAlgorithm"],
     "dummy": ["DummyAlgorithm"],
     # single task learning (fine-tuning)
-    "classification": ["ImageClassificationFineTuningForCLIP"],
+    "classification": [
+        "ImageClassificationFineTuningForCLIP",
+        "ContinualImageClassificationFineTuningForCLIP",
+    ],
     "lm_finetune": ["FullFinetuneSFT", "PeftFinetuneSFT", "BradleyTerryRewardModeling"],
     # analysis
     "analysis": ["TaskVectorCosSimilarity", "TaskVectorViolinPlot"],
@@ -27,11 +30,12 @@ _import_structure = {
         "TaskArithmeticForLlama",
         "LinearInterpolationAlgorithm",
     ],
+    "slerp": ["SlerpMergeAlgorithm"],
     "simple_average": ["SimpleAverageAlgorithm"],
     "weighted_average": ["WeightedAverageAlgorithm", "WeightedAverageForLLama"],
     "task_arithmetic": ["TaskArithmeticAlgorithm"],
     "ties_merging": ["TiesMergingAlgorithm"],
-    "dare": ["DareSimpleAverage", "DareTaskArithmetic"],
+    "dare": ["DareSimpleAverage", "DareTaskArithmetic", "DareTiesMerging"],
     "fisher_merging": [
         "FisherMergingForCLIPVisionModel",
         "FisherMergingAlgorithmForGPT2",
@@ -50,6 +54,13 @@ _import_structure = {
     ],
     "ada_svd": ["AdaSVDMergingForCLIPVisionModel"],
     "task_singular_vector": ["TaskSingularVectorMerging"],
+    "isotropic_merging": [
+        "ISO_C_Merge",  # alias
+        "ISO_CTS_Merge",  # alias
+        "IsotropicMergingInCommonAndTaskSubspace",
+        "IsotropicMergingInCommonSubspace",
+    ],
+    "opcm": ["OPCMForCLIP"],
     # plug-and-play model merging methods
     "concrete_subspace": [
         "ConcreteTaskArithmeticAlgorithmForCLIP",
@@ -96,13 +107,16 @@ if TYPE_CHECKING:
     from .adamerging import *
     from .analysis import TaskVectorCosSimilarity, TaskVectorViolinPlot
     from .base_algorithm import BaseAlgorithm, BaseModelFusionAlgorithm
-    from .classification import ImageClassificationFineTuningForCLIP
+    from .classification import (
+        ContinualImageClassificationFineTuningForCLIP,
+        ImageClassificationFineTuningForCLIP,
+    )
     from .concrete_subspace import (
         ConcreteLayerWiseAdaMergingForCLIP,
         ConcreteTaskArithmeticAlgorithmForCLIP,
         ConcreteTaskWiseAdaMergingForCLIP,
     )
-    from .dare import DareSimpleAverage, DareTaskArithmetic
+    from .dare import DareSimpleAverage, DareTaskArithmetic, DareTiesMerging
     from .dawe import DataAdaptiveWeightEnsemblingForCLIP
     from .depth_upscaling import DepthUpscalingAlgorithm, DepthUpscalingForLlama
     from .dummy import DummyAlgorithm
@@ -112,6 +126,12 @@ if TYPE_CHECKING:
         WeightedEnsembleAlgorithm,
     )
     from .fisher_merging import FisherMergingForCLIPVisionModel
+    from .isotropic_merging import (
+        ISO_C_Merge,
+        ISO_CTS_Merge,
+        IsotropicMergingInCommonAndTaskSubspace,
+        IsotropicMergingInCommonSubspace,
+    )
     from .linear import (
         ExPOAlgorithm,
         ExPOAlgorithmForLlama,
@@ -127,6 +147,7 @@ if TYPE_CHECKING:
         MixtralUpscalingAlgorithm,
     )
     from .model_recombination import ModelRecombinationAlgorithm
+    from .opcm import OPCMForCLIP
     from .pruning import (
         MagnitudeDiffPruningAlgorithm,
         MagnitudePruningForLlama,
@@ -140,6 +161,7 @@ if TYPE_CHECKING:
     from .rankone_moe import CLIPRankOneMoEAlgorithm, RankOneMoEAlgorithm
     from .regmean import RegMeanAlgorithmForCLIP, RegMeanAlgorithmForGPT2
     from .simple_average import SimpleAverageAlgorithm
+    from .slerp import SlerpMergeAlgorithm
     from .smile_upscaling import (
         SingularProjectionMergingAlgorithm,
         SmileUpscalingAlgorithm,

{fusion_bench-0.2.8 → fusion_bench-0.2.10}/fusion_bench/method/classification/__init__.py

@@ -1,2 +1,3 @@
 # flake8: noqa F401
 from .clip_finetune import ImageClassificationFineTuningForCLIP
+from .continual_clip_finetune import ContinualImageClassificationFineTuningForCLIP

{fusion_bench-0.2.8 → fusion_bench-0.2.10}/fusion_bench/method/classification/clip_finetune.py

@@ -184,9 +184,7 @@ class ImageClassificationFineTuningForCLIP(
                 self.save_model(classifier, save_path)
 
         if config.state_dict_save_path is not None:
-            self.save_model(
-                classifier, config.state_dict_save_path, trainable_only=True
-            )
+            self.save_model(classifier, config.state_dict_save_path)
         self.print_profile_summary()
         return classifier.clip_model.vision_model
 

fusion_bench-0.2.10/fusion_bench/method/classification/continual_clip_finetune.py

@@ -0,0 +1,297 @@
+import os
+import random
+import time
+from copy import deepcopy
+from typing import Optional, Tuple, cast
+
+import lightning as L
+import torch
+from omegaconf import DictConfig, OmegaConf
+from peft import LoraConfig, PeftModel, get_peft_model
+from peft.tuners.lora import LoraLayer
+from safetensors.torch import save_file
+from torch import nn
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
+from transformers.models.clip.modeling_clip import CLIPVisionTransformer
+
+from fusion_bench import BaseAlgorithm, print_parameters
+from fusion_bench.compat.modelpool import to_modelpool
+from fusion_bench.dataset.clip_dataset import CLIPDataset
+from fusion_bench.mixins import CLIPClassificationMixin
+from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.modelpool import CLIPVisionModelPool
+from fusion_bench.models.hf_clip import HFCLIPClassifier
+from fusion_bench.models.linearized.linearized_model_utils import LinearizedModelWraper
+from fusion_bench.taskpool import CLIPVisionModelTaskPool
+from fusion_bench.utils.data import InfiniteDataLoader
+from fusion_bench.utils.fabric import seed_everything_by_time
+from fusion_bench.utils.json import load_from_json, save_to_json
+
+
+class ContinualImageClassificationFineTuningForCLIP(
+    CLIPClassificationMixin,
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
+    # attributes to configuration keys mapping
+    _config_mapping = BaseAlgorithm._config_mapping | {
+        "seed": "seed",
+        "shuffle_order": "shuffle_order",
+        "learning_rate": "learning_rate",
+        "weight_decay": "weight_decay",
+        "num_steps": "num_steps",
+        "batch_size": "batch_size",
+        "num_workers": "num_workers",
+        "save_interval": "save_interval",
+        "state_dict_load_path": "state_dict_load_path",
+        "state_dict_save_path": "state_dict_save_path",
+        "skip_training": "skip_training",
+        "use_lora": "use_lora",
+        "lora_config": "lora_config",
+    }
+
+    def __init__(
+        self,
+        seed: int = 42,
+        shuffle_order: bool = True,
+        learning_rate: float = 1e-5,
+        weight_decay: float = 0,
+        num_steps: int = 4000,
+        batch_size: int = 128,
+        num_workers: int = 16,
+        save_interval: int = 500,
+        state_dict_load_path: Optional[str] = None,
+        state_dict_save_path: Optional[str] = None,
+        skip_training: bool = False,
+        use_lora: bool = False,
+        lora_config: Optional[LoraConfig] = None,
+    ):
+        self.seed = seed
+        self.shuffle_order = shuffle_order
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.num_steps = num_steps
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.save_interval = save_interval
+        self.state_dict_load_path = state_dict_load_path
+        self.state_dict_save_path = state_dict_save_path
+        self.skip_training = skip_training
+        self.use_lora = use_lora
+        self.lora_config = lora_config
+
+    def run(self, modelpool: CLIPVisionModelPool):
+        self.modelpool = to_modelpool(modelpool)
+        config = self.config
+        self.log_hyperparams(config, filename="method_config.yaml")
+        self.finetune_method = "fine-tune"
+
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+        else:
+            seed_everything_by_time(self.fabric)
+
+        task_names = list(modelpool.train_dataset_names)
+        if self.shuffle_order:
+            random.shuffle(task_names)
+        if self.fabric.is_global_zero:
+            save_to_json(task_names, os.path.join(self.log_dir, "task_names.json"))
+
+        if self._program.taskpool is not None and isinstance(
+            self._program.taskpool, CLIPVisionModelTaskPool
+        ):
+            has_taskpool = True
+            taskpool = cast(CLIPVisionModelTaskPool, self._program.taskpool)
+            test_datasets = taskpool._test_datasets
+        else:
+            has_taskpool = False
+
+        with self.profile("setup model and optimizer"):
+            processor, classifier, optimizer, lr_scheduler = self.setup_model()
+
+            if self.state_dict_load_path is not None:
+                self.fabric.load(
+                    self.state_dict_load_path,
+                    {"vision_model": classifier.clip_model.vision_model},
+                )
+                if self.skip_training:
+                    return classifier.clip_model.vision_model
+
+            self.setup_zero_shot_classification_head(
+                clip_processor=processor,
+                clip_model=classifier.clip_model,
+                task_names=task_names,
+            )
+
+            init_optimizer_state_dict = optimizer.state_dict()
+            init_lr_scheduler_state_dict = lr_scheduler.state_dict()
+            self.fabric.setup(classifier, optimizer)
+
+        with self.profile("setup data"):
+            train_datasets = [
+                CLIPDataset(modelpool.load_train_dataset(task_name), processor)
+                for task_name in task_names
+            ]
+            train_dataloaders = [
+                DataLoader(
+                    dataset,
+                    shuffle=True,
+                    batch_size=self.batch_size,
+                    num_workers=self.num_workers,
+                )
+                for dataset in train_datasets
+            ]
+            train_dataloaders = self.fabric.setup_dataloaders(*train_dataloaders)
+            if not isinstance(train_dataloaders, (list, tuple)):
+                train_dataloaders = [train_dataloaders]
+            train_dataloader_iters = [
+                iter(InfiniteDataLoader(loader)) for loader in train_dataloaders
+            ]
+
+        # continual train
+        for task_idx, task_name in tqdm(
+            enumerate(task_names),
+            dynamic_ncols=True,
+            disable=not self.fabric.is_global_zero,
+        ):
+            train_dataloader_iter = train_dataloader_iters[task_idx]
+
+            # reset optimizer and lr scheduler
+            print("reset optimizer and lr scheduler")
+            optimizer.load_state_dict(init_optimizer_state_dict)
+            lr_scheduler.load_state_dict(init_lr_scheduler_state_dict)
+
+            for step_idx in tqdm(
+                range(self.num_steps),
+                desc=f"continual fine-tune on {task_name}",
+                disable=not self.fabric.is_global_zero,
+                dynamic_ncols=True,
+                leave=False,
+            ):
+                optimizer.zero_grad()
+                loss = 0
+                with self.profile("data loading"):
+                    batch = next(train_dataloader_iter)
+                    images, labels = batch
+                with self.profile("forward"):
+                    classifier.zeroshot_weights = self.zeroshot_weights[task_name]
+                    logits = classifier(images)
+                    assert (
+                        labels.max() < logits.shape[1]
+                    ), f"for task {task_name}, labels.max() = {labels.max()}, logits.shape[1] = {logits.shape[1]}"
+                loss = loss + nn.functional.cross_entropy(logits, labels)
+
+                with self.profile("backward"):
+                    self.fabric.backward(loss)
+                with self.profile("optimizer step"):
+                    optimizer.step()
+                    lr_scheduler.step()
+
+                metrics = {"train/loss": loss}
+                self.fabric.log_dict(metrics, step=step_idx)
+
+                if (step_idx + 1) % self.save_interval == 0:
+                    save_path = os.path.join(
+                        self.log_dir,
+                        "checkpoints",
+                        f"task={task_idx}_step={step_idx}.ckpt",
+                    )
+                    self.save_model(classifier, save_path)
+
+            if has_taskpool:
+                taskpool._is_setup = False
+                taskpool._test_datasets = DictConfig(
+                    {t: test_datasets[t] for t in task_names[: task_idx + 1]}
+                )
+                eval_report = taskpool.evaluate(
+                    deepcopy(classifier.clip_model.vision_model),
+                    name=task_name,
+                )
+                if self.fabric.is_global_zero:
+                    save_to_json(
+                        eval_report,
+                        os.path.join(self.log_dir, f"results_{task_idx}.json"),
+                    )
+
+        if self.state_dict_save_path is not None:
+            self.save_model(classifier, self.state_dict_save_path)
+        self.print_profile_summary()
+        return classifier.clip_model.vision_model
+
+    def save_model(
+        self,
+        model: HFCLIPClassifier | CLIPModel | CLIPVisionModel | CLIPVisionTransformer,
+        save_path: str,
+    ):
+        """
+        Save the vision model to the specified path.
+
+        Args:
+            model (Union[HFCLIPClassifier, CLIPModel, CLIPVisionModel, CLIPVisionTransformer]): The model to save.
+            save_path (str): The path to save the model.
+        """
+        if isinstance(model, HFCLIPClassifier):
+            vision_model = model.clip_model.vision_model
+        elif isinstance(model, CLIPModel):
+            vision_model = model.vision_model
+        elif isinstance(model, CLIPVisionModel):
+            vision_model = model.vision_model
+        elif isinstance(model, CLIPVisionTransformer):
+            vision_model = model
+        else:
+            raise ValueError(f"Unsupported model type: {type(model)}")
+
+        save_dir = os.path.dirname(save_path)
+        if save_dir and not os.path.exists(save_dir):
+            os.makedirs(save_dir, exist_ok=True)
+        self.fabric.save(save_path, {"vision_model": vision_model})
+
+    def setup_model(self):
+        """
+        Sets up the model, optimizer, and learning rate scheduler.
+
+        This method initializes the CLIP model, applies LoRA if specified, and configures the optimizer and learning rate scheduler.
+
+        Returns:
+            Tuple: A tuple containing the processor, classifier, optimizer, and learning rate scheduler.
+        """
+        config = self.config
+        modelpool = self.modelpool
+
+        clip_model: CLIPModel = modelpool.load_clip_model("_pretrained_")
+        processor = modelpool.load_processor()
+
+        self.finetune_method = "full fine-tune"
+        if self.use_lora:
+            self.finetune_method = "lora fine-tune"
+            lora_config = LoraConfig(
+                **OmegaConf.to_container(
+                    self.lora_config, resolve=True, enum_to_str=True
+                )
+            )
+            clip_model.vision_model = get_peft_model(
+                clip_model.vision_model, lora_config
+            )
+
+        classifier = HFCLIPClassifier(clip_model, processor=processor)
+
+        if self.fabric.is_global_zero:
+            print("=== Model Summary (For Vision Model Only) ===")
+            print_parameters(classifier.clip_model.vision_model)
+        # configure optimizers
+        optimizer = torch.optim.Adam(
+            [
+                p
+                for p in classifier.clip_model.vision_model.parameters()
+                if p.requires_grad
+            ],
+            lr=self.learning_rate,
+            weight_decay=self.weight_decay,
+        )
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer=optimizer, T_max=self.num_steps
+        )
+
+        return processor, classifier, optimizer, lr_scheduler

{fusion_bench-0.2.8 → fusion_bench-0.2.10}/fusion_bench/method/dare/__init__.py

@@ -1,3 +1,4 @@
 # flake8: noqa F401
 from .simple_average import DareSimpleAverage
 from .task_arithmetic import DareTaskArithmetic
+from .ties_merging import DareTiesMerging

{fusion_bench-0.2.8 → fusion_bench-0.2.10}/fusion_bench/method/dare/task_arithmetic.py

@@ -33,21 +33,28 @@ class DareTaskArithmetic(BaseAlgorithm):
         self.rescale = rescale
         super().__init__(**kwargs)
 
+    def _load_task_vector(
+        self,
+        modelpool: BaseModelPool,
+        model_name: str,
+        pretrained_model: nn.Module,
+    ):
+        finetuned_model = modelpool.load_model(model_name)
+        task_vector = module_sub_(finetuned_model, pretrained_model)
+        return task_vector
+
     @torch.no_grad()
     def run(self, modelpool: BaseModelPool):
         assert (
             self.sparsity_ratio >= 0 and self.sparsity_ratio <= 1
         ), "Sparsity ratio must be between 0 and 1"
         pretrained_model = modelpool.load_pretrained_model()
-        finetuned_models = {
-            model_name: modelpool.load_model(model_name)
-            for model_name in modelpool.model_names
-        }
+
+        # load task vectors
         task_vectors = {
-            model_name: module_sub_(finetuned_models[model_name], pretrained_model)
-            for model_name in finetuned_models
+            model_name: self._load_task_vector(modelpool, model_name, pretrained_model)
+            for model_name in modelpool.model_names
         }
-        del finetuned_models
 
         # drop and rescale task vectors
         for model_name, tv in task_vectors.items():