fusion-bench 0.2.12__py3-none-any.whl → 0.2.13__py3-none-any.whl

fusion_bench/method/ties_merging/ties_merging.py

@@ -16,6 +16,7 @@ from torch import Tensor, nn
 
 from fusion_bench.compat.modelpool import to_modelpool
 from fusion_bench.method import BaseAlgorithm
+from fusion_bench.mixins import SimpleProfilerMixin
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.type import StateDictType
 
@@ -24,7 +25,7 @@ from .ties_merging_utils import state_dict_to_vector, ties_merging, vector_to_st
 log = logging.getLogger(__name__)
 
 
-class TiesMergingAlgorithm(BaseAlgorithm):
+class TiesMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
     """
     TiesMergingAlgorithm is a class for fusing multiple models using the TIES merging technique.
 
@@ -84,34 +85,38 @@ class TiesMergingAlgorithm(BaseAlgorithm):
         scaling_factor = self.scaling_factor
         threshold = self.threshold
 
-        # Load the pretrained model
-        pretrained_model = modelpool.load_model("_pretrained_")
-
-        # Load the state dicts of the models
-        ft_checks: List[StateDictType] = [
-            modelpool.load_model(model_name).state_dict(keep_vars=True)
-            for model_name in modelpool.model_names
-        ]
-        ptm_check: StateDictType = pretrained_model.state_dict(keep_vars=True)
-
-        # Compute the task vectors
-        flat_ft: Tensor = torch.vstack(
-            [state_dict_to_vector(check, remove_keys) for check in ft_checks]
-        )
-        flat_ptm: Tensor = state_dict_to_vector(ptm_check, remove_keys)
-        tv_flat_checks = flat_ft - flat_ptm
-
-        # Perform TIES Merging
-        merged_tv = ties_merging(
-            tv_flat_checks,
-            reset_thresh=threshold,
-            merge_func=merge_func,
-        )
-        merged_check = flat_ptm + scaling_factor * merged_tv
-        merged_state_dict = vector_to_state_dict(
-            merged_check, ptm_check, remove_keys=remove_keys
-        )
-
-        # Load the merged state dict into the pretrained model
-        pretrained_model.load_state_dict(merged_state_dict)
+        with self.profile("loading models"):
+            # Load the pretrained model
+            pretrained_model = modelpool.load_model("_pretrained_")
+
+            # Load the state dicts of the models
+            ft_checks: List[StateDictType] = [
+                modelpool.load_model(model_name).state_dict(keep_vars=True)
+                for model_name in modelpool.model_names
+            ]
+            ptm_check: StateDictType = pretrained_model.state_dict(keep_vars=True)
+
+        with self.profile("merging models"):
+            # Compute the task vectors
+            flat_ft: Tensor = torch.vstack(
+                [state_dict_to_vector(check, remove_keys) for check in ft_checks]
+            )
+            flat_ptm: Tensor = state_dict_to_vector(ptm_check, remove_keys)
+            tv_flat_checks = flat_ft - flat_ptm
+
+            # Perform TIES Merging
+            merged_tv = ties_merging(
+                tv_flat_checks,
+                reset_thresh=threshold,
+                merge_func=merge_func,
+            )
+            merged_check = flat_ptm + scaling_factor * merged_tv
+            merged_state_dict = vector_to_state_dict(
+                merged_check, ptm_check, remove_keys=remove_keys
+            )
+
+            # Load the merged state dict into the pretrained model
+            pretrained_model.load_state_dict(merged_state_dict)
+
+        self.print_profile_summary()
         return pretrained_model

fusion_bench/method/we_moe/we_moe.py

@@ -5,7 +5,6 @@ from typing import cast  # noqa: F401
 import lightning as L
 import lightning.fabric.wrappers
 import torch
-from lightning.pytorch.profilers import SimpleProfiler
 from omegaconf import DictConfig
 from torch import Tensor
 from torch.utils.data import DataLoader
@@ -13,6 +12,7 @@ from tqdm.autonotebook import tqdm
 
 from fusion_bench.compat.method.base_algorithm import ModelFusionAlgorithm
 from fusion_bench.compat.modelpool import ModelPool
+from fusion_bench.mixins import SimpleProfilerMixin
 from fusion_bench.models.we_moe import WeightEnsemblingMoE
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.parameters import print_parameters
@@ -34,7 +34,10 @@ def entropy_loss(logits: Tensor) -> Tensor:
     return -torch.sum(probs * torch.log(probs + 1e-8), dim=-1).mean()
 
 
-class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
+class WeightEnsemblingMoEAlgorithm(
+    ModelFusionAlgorithm,
+    SimpleProfilerMixin,
+):
     """
     Algorithm for fusing models using Weight Ensembling Mixture of Experts (MoE).
 
@@ -44,7 +47,6 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
     Attributes:
         _fabric (L.Fabric): The fabric for distributed training.
         modelpool (ModelPool): The pool of models to be fused.
-        profiler (SimpleProfiler): The profiler for measuring performance.
     """
 
     _fabric: L.Fabric = None
@@ -66,9 +68,6 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
             self._fabric.launch()
         else:
             assert "No CUDA device available."
-        self.profiler = SimpleProfiler(
-            self.config.get("cache_dir", "outputs"), "we_moe_profiler.txt"
-        )
 
     @abstractmethod
     def load_checkpoint(self, model, checkpoint):
@@ -177,9 +176,9 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
         for step_idx in pbar:
             if self.config.use_grad_accumulate:
                 for task in self.modelpool.model_names:
-                    with self.profiler.profile("data time"):
+                    with self.profile("data time"):
                         batch = next(self.get_shuffled_test_loader_iter(task))
-                    with self.profiler.profile("forward pass"):
+                    with self.profile("forward pass"):
                         logits = self.compute_logits(module, batch, task)
                         assert (
                             logits.dim() == 2
@@ -187,23 +186,23 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
                         loss = entropy_loss(logits)
                     # .backward() accumulates when .zero_grad() wasn't called
                     # this can save memory
-                    with self.profiler.profile("backward pass"):
+                    with self.profile("backward pass"):
                         self._fabric.backward(loss, retain_graph=True)
             else:
                 loss = 0
                 for task in self.modelpool.model_names:
-                    with self.profiler.profile("data time"):
+                    with self.profile("data time"):
                         batch = next(self.get_shuffled_test_loader_iter(task))
-                    with self.profiler.profile("forward pass"):
+                    with self.profile("forward pass"):
                         logits = self.compute_logits(module, batch, task)
                         assert (
                             logits.dim() == 2
                         ), f"Expected logits to be 2D, got {logits.dim()}"
                         loss = loss + entropy_loss(logits)
-                with self.profiler.profile("backward pass"):
+                with self.profile("backward pass"):
                     self._fabric.backward(loss, retain_graph=True)
 
-            with self.profiler.profile("optimizer step"):
+            with self.profile("optimizer step"):
                 optimizer.step()
                 optimizer.zero_grad()
 
@@ -232,7 +231,7 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
             )
             self.load_checkpoint(moe_model, self.config.checkpoint)
         else:
-            with self.profiler.profile("test-time adaptation"):
+            with self.profile("test-time adaptation"):
                 moe_model = self.test_time_adaptation(moe_model)
             if self.config.get("save_checkpoint", False):
                 log.info(f"save checkpoint to {self.config.save_checkpoint}")
@@ -243,5 +242,5 @@ class WeightEnsemblingMoEAlgorithm(ModelFusionAlgorithm):
 
         # enable sample-wise adaptation
         moe_model.batch_reduce = False
-        print(self.profiler.summary())
+        self.print_profile_summary()
         return moe_model

fusion_bench/mixins/__init__.py

@@ -6,20 +6,23 @@ from typing_extensions import TYPE_CHECKING
 from fusion_bench.utils.lazy_imports import LazyImporter
 
 _import_structure = {
+    "clip_classification": ["CLIPClassificationMixin"],
+    "fabric_training": ["FabricTrainingMixin"],
+    "hydra_config": ["HydraConfigMixin"],
     "lightning_fabric": ["LightningFabricMixin"],
+    "openclip_classification": ["OpenCLIPClassificationMixin"],
     "serialization": ["YAMLSerializationMixin", "BaseYAMLSerializableModel"],
     "simple_profiler": ["SimpleProfilerMixin"],
-    "clip_classification": ["CLIPClassificationMixin"],
-    "fabric_training": ["FabricTrainingMixin"],
 }
 
 if TYPE_CHECKING:
     from .clip_classification import CLIPClassificationMixin
     from .fabric_training import FabricTrainingMixin
+    from .hydra_config import HydraConfigMixin
     from .lightning_fabric import LightningFabricMixin
+    from .openclip_classification import OpenCLIPClassificationMixin
     from .serialization import BaseYAMLSerializableModel, YAMLSerializationMixin
     from .simple_profiler import SimpleProfilerMixin
-
 else:
     sys.modules[__name__] = LazyImporter(
         __name__,

fusion_bench/mixins/hydra_config.py

@@ -0,0 +1,49 @@
+import logging
+import os
+from copy import deepcopy
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import hydra.core.global_hydra
+from hydra import compose, initialize
+from omegaconf import DictConfig, OmegaConf
+
+from fusion_bench.utils import import_object, instantiate
+from fusion_bench.utils.instantiate import set_print_function_call
+
+log = logging.getLogger(__name__)
+
+
+class HydraConfigMixin:
+    """
+    A mixin for classes that need to be instantiated from a config file.
+    """
+
+    @classmethod
+    def from_config(
+        cls,
+        config_name: Union[str, Path],
+        overrides: Optional[List[str]] = None,
+    ):
+        if not hydra.core.global_hydra.GlobalHydra.instance().is_initialized():
+            raise RuntimeError("Hydra is not initialized.")
+        else:
+            cfg = compose(config_name=config_name, overrides=overrides)
+
+        config_groups = config_name.split("/")[:-1]
+        for config_group in config_groups:
+            cfg = cfg[config_group]
+
+        if "_target_" in cfg:
+            # if the config has a _target_ key, check if it is equal to the class name
+            target_cls = import_object(cfg["_target_"])
+            if target_cls != cls:
+                log.warning(
+                    f"The _target_ key in the config is {cfg['_target_']}, but the class name is {cls.__name__}."
+                )
+            with set_print_function_call(False):
+                obj = instantiate(cfg)
+        else:
+            obj = cls(**cfg)
+
+        return obj

fusion_bench/mixins/openclip_classification.py

@@ -0,0 +1,11 @@
+import logging
+
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.models.open_clip import ImageClassifier, ImageEncoder
+
+log = logging.getLogger(__name__)
+
+
+class OpenCLIPClassificationMixin(LightningFabricMixin):
+    _train_processor = None
+    _test_processor = None

fusion_bench/mixins/simple_profiler.py

@@ -1,5 +1,5 @@
 from contextlib import contextmanager
-from typing import Generator
+from typing import Generator, Optional
 
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 from lightning.pytorch.profilers import SimpleProfiler
@@ -70,7 +70,9 @@ class SimpleProfilerMixin:
         self.profiler.stop(action_name)
 
     @rank_zero_only
-    def print_profile_summary(self):
+    def print_profile_summary(self, title: Optional[str] = None):
+        if title is not None:
+            print(title)
         print(self.profiler.summary())
 
     def __del__(self):

fusion_bench/modelpool/__init__.py

@@ -6,12 +6,13 @@ from fusion_bench.utils.lazy_imports import LazyImporter
 
 _import_structure = {
     "base_pool": ["BaseModelPool"],
+    "causal_lm": ["CausalLMPool", "CausalLMBackbonePool"],
     "clip_vision": ["CLIPVisionModelPool"],
     "nyuv2_modelpool": ["NYUv2ModelPool"],
     "huggingface_automodel": ["AutoModelPool"],
-    "causal_lm": ["CausalLMPool", "CausalLMBackbonePool"],
     "seq2seq_lm": ["Seq2SeqLMPool"],
     "PeftModelForSeq2SeqLM": ["PeftModelForSeq2SeqLMPool"],
+    "openclip_vision": ["OpenCLIPVisionModelPool"],
     "huggingface_gpt2_classification": [
         "HuggingFaceGPT2ClassificationPool",
         "GPT2ForSequenceClassificationPool",
@@ -30,6 +31,7 @@ if TYPE_CHECKING:
         HuggingFaceGPT2ClassificationPool,
     )
     from .nyuv2_modelpool import NYUv2ModelPool
+    from .openclip_vision import OpenCLIPVisionModelPool
     from .PeftModelForSeq2SeqLM import PeftModelForSeq2SeqLMPool
     from .seq2seq_lm import Seq2SeqLMPool
     from .seq_classification_lm import SeqenceClassificationModelPool

fusion_bench/modelpool/base_pool.py

@@ -7,7 +7,7 @@ from omegaconf import DictConfig
 from torch import nn
 from torch.utils.data import Dataset
 
-from fusion_bench.mixins import BaseYAMLSerializableModel
+from fusion_bench.mixins import BaseYAMLSerializableModel, HydraConfigMixin
 from fusion_bench.utils import instantiate, timeit_context
 
 __all__ = ["BaseModelPool"]
@@ -15,7 +15,7 @@ __all__ = ["BaseModelPool"]
 log = logging.getLogger(__name__)
 
 
-class BaseModelPool(BaseYAMLSerializableModel):
+class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
     """
     A class for managing and interacting with a pool of models along with their associated datasets or other specifications. For example, a model pool may contain multiple models, each with its own training, validation, and testing datasets. As for the specifications, a vision model pool may contain image preprocessor, and a language model pool may contain a tokenizer.
 

fusion_bench/modelpool/openclip_vision/__init__.py

@@ -0,0 +1 @@
+from .modelpool import OpenCLIPVisionModelPool

fusion_bench/modelpool/openclip_vision/modelpool.py

@@ -0,0 +1,255 @@
+import logging
+import pickle
+import sys
+from typing import Callable, Optional, Union, cast
+
+import torch
+from datasets import load_dataset
+from omegaconf import DictConfig, OmegaConf
+from torch import nn
+
+from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.models.open_clip import ClassificationHead, ImageEncoder
+from fusion_bench.utils import instantiate
+from fusion_bench.utils.expr import is_expr_match
+from fusion_bench.utils.packages import _get_package_version, compare_versions
+
+log = logging.getLogger(__name__)
+
+# Add flag to track if warning has been shown
+_openclip_version_warning_shown = False
+
+
+def _check_and_redirect_open_clip_modeling():
+    global _openclip_version_warning_shown
+    if compare_versions(_get_package_version("open-clip-torch").__str__(), "2.0.2") > 0:
+        if not _openclip_version_warning_shown:
+            log.warning(
+                "OpenCLIP version is greater than 2.0.2. This may cause issues with the modelpool."
+            )
+            _openclip_version_warning_shown = True
+        import open_clip.model
+        import open_clip.transformer
+
+        if not hasattr(open_clip.model, "VisualTransformer"):
+            open_clip.model.VisualTransformer = open_clip.model.VisionTransformer
+        if not hasattr(open_clip.model, "Transformer"):
+            open_clip.model.Transformer = open_clip.transformer.Transformer
+        if not hasattr(open_clip.model, "ResidualAttentionBlock"):
+            open_clip.model.ResidualAttentionBlock = (
+                open_clip.transformer.ResidualAttentionBlock
+            )
+
+    try:
+        import src
+        import src.modeling
+    except ImportError:
+        if "src" not in sys.modules:
+            # redirect the import of `src` to `fusion_bench.models.open_clip`
+            import fusion_bench.models.open_clip as open_clip
+
+            sys.modules["src"] = open_clip
+            log.warning(
+                "`src` is not imported."
+                "Redirecting the import to `fusion_bench.models.open_clip`"
+            )
+        if "src.modeling" not in sys.modules:
+            # redirect the import of `src.modeling` to `fusion_bench.models.open_clip.modeling`
+            import fusion_bench.models.open_clip.modeling as open_clip_modeling
+
+            sys.modules["src.modeling"] = open_clip_modeling
+            log.warning(
+                "`src.modeling` is not imported."
+                "Redirecting the import to `fusion_bench.models.open_clip.modeling`"
+            )
+
+
+def load_classifier_head(model_config: Union[str, DictConfig], *args, **kwargs):
+    if isinstance(model_config, str):
+        _check_and_redirect_open_clip_modeling()
+        log.info(f"Loading `ClassificationHead` from {model_config}")
+        weights_only = kwargs["weights_only"] if "weights_only" in kwargs else False
+        head = torch.load(model_config, weights_only=weights_only, *args, **kwargs)
+    elif isinstance(model_config, nn.Module):
+        log.info(f"Returning existing model: {model_config}")
+        head = model_config
+    else:
+        head = instantiate(model_config, *args, **kwargs)
+    head = cast(ClassificationHead, head)
+    return head
+
+
+class OpenCLIPVisionModelPool(BaseModelPool):
+    """
+    A model pool for managing OpenCLIP Vision models (models from task vector paper).
+    """
+
+    _train_processor = None
+    _test_processor = None
+
+    def __init__(
+        self,
+        models: DictConfig,
+        classification_heads: Optional[DictConfig] = None,
+        **kwargs,
+    ):
+        super().__init__(models, **kwargs)
+        self._classification_heads = classification_heads
+
+    @property
+    def train_processor(self):
+        if self._train_processor is None:
+            encoder: ImageEncoder = self.load_pretrained_or_first_model()
+            self._train_processor = encoder.train_preprocess
+            if self._test_processor is None:
+                self._test_processor = encoder.val_preprocess
+        return self._train_processor
+
+    @property
+    def test_processor(self):
+        if self._test_processor is None:
+            encoder: ImageEncoder = self.load_pretrained_or_first_model()
+            if self._train_processor is None:
+                self._train_processor = encoder.train_preprocess
+            self._test_processor = encoder.val_preprocess
+        return self._test_processor
+
+    def load_model(
+        self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
+    ) -> ImageEncoder:
+        R"""
+        The model config can be:
+
+        - A string, which is the path to the model checkpoint in pickle format. Load directly using `torch.load`.
+        - {"model_name": str, "pickle_path": str}, load the model from the binary file (pickle format). This will first construct the model using `ImageEncoder(model_name)`, and then load the state dict from model located in the pickle file.
+        - {"model_name": str, "state_dict_path": str}, load the model from the state dict file. This will first construct the model using `ImageEncoder(model_name)`, and then load the state dict from the file.
+        - Default, load the model using `instantiate` from hydra.
+        """
+        if (
+            isinstance(model_name_or_config, str)
+            and model_name_or_config in self._models
+        ):
+            model_config = self._models[model_name_or_config]
+        else:
+            model_config = model_name_or_config
+        if isinstance(model_config, DictConfig):
+            model_config = OmegaConf.to_container(model_config, resolve=True)
+
+        if isinstance(model_config, str):
+            # the model config is a string, which is the path to the model checkpoint in pickle format
+            # load the model using `torch.load`
+            # this is the original usage in the task arithmetic codebase
+            _check_and_redirect_open_clip_modeling()
+            log.info(f"loading ImageEncoder from {model_config}")
+            weights_only = kwargs["weights_only"] if "weights_only" in kwargs else False
+            try:
+                encoder = torch.load(
+                    model_config, weights_only=weights_only, *args, **kwargs
+                )
+            except RuntimeError as e:
+                encoder = pickle.load(open(model_config, "rb"))
+        elif is_expr_match({"model_name": str, "pickle_path": str}, model_config):
+            # the model config is a dictionary with the following keys:
+            # - model_name: str, the name of the model
+            # - pickle_path: str, the path to the binary file (pickle format)
+            # load the model from the binary file (pickle format)
+            # this is useful when you use a newer version of torchvision
+            _check_and_redirect_open_clip_modeling()
+            log.info(
+                f"loading ImageEncoder of {model_config['model_name']} from {model_config['pickle_path']}"
+            )
+            weights_only = kwargs["weights_only"] if "weights_only" in kwargs else False
+            try:
+                encoder = torch.load(
+                    model_config["pickle_path"],
+                    weights_only=weights_only,
+                    *args,
+                    **kwargs,
+                )
+            except RuntimeError as e:
+                encoder = pickle.load(open(model_config["pickle_path"], "rb"))
+            _encoder = ImageEncoder(model_config["model_name"])
+            _encoder.load_state_dict(encoder.state_dict())
+            encoder = _encoder
+        elif is_expr_match({"model_name": str, "state_dict_path": str}, model_config):
+            # the model config is a dictionary with the following keys:
+            # - model_name: str, the name of the model
+            # - state_dict_path: str, the path to the state dict file
+            # load the model from the state dict file
+            log.info(
+                f"loading ImageEncoder of {model_config['model_name']} from {model_config['state_dict_path']}"
+            )
+            encoder = ImageEncoder(model_config["model_name"])
+            encoder.load_state_dict(
+                torch.load(
+                    model_config["state_dict_path"], weights_only=True, *args, **kwargs
+                )
+            )
+        elif isinstance(model_config, nn.Module):
+            # the model config is an existing model
+            log.info(f"Returning existing model: {model_config}")
+            encoder = model_config
+        else:
+            encoder = super().load_model(model_name_or_config, *args, **kwargs)
+        encoder = cast(ImageEncoder, encoder)
+
+        # setup the train and test processors
+        if self._train_processor is None and hasattr(encoder, "train_preprocess"):
+            self._train_processor = encoder.train_preprocess
+        if self._test_processor is None and hasattr(encoder, "val_preprocess"):
+            self._test_processor = encoder.val_preprocess
+
+        return encoder
+
+    def load_classification_head(
+        self, model_name_or_config: Union[str, DictConfig], *args, **kwargs
+    ) -> ClassificationHead:
+        R"""
+        The model config can be:
+
+        - A string, which is the path to the model checkpoint in pickle format. Load directly using `torch.load`.
+        - Default, load the model using `instantiate` from hydra.
+        """
+        if (
+            isinstance(model_name_or_config, str)
+            and model_name_or_config in self._classification_heads
+        ):
+            model_config = self._classification_heads[model_name_or_config]
+        else:
+            model_config = model_name_or_config
+
+        head = load_classifier_head(model_config, *args, **kwargs)
+        return head
+
+    def load_train_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._train_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading train dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="train")
+        else:
+            dataset = super().load_train_dataset(dataset_name, *args, **kwargs)
+        return dataset
+
+    def load_val_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._val_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading validation dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="validation")
+        else:
+            dataset = super().load_val_dataset(dataset_name, *args, **kwargs)
+        return dataset
+
+    def load_test_dataset(self, dataset_name: str, *args, **kwargs):
+        dataset_config = self._test_datasets[dataset_name]
+        if isinstance(dataset_config, str):
+            log.info(
+                f"Loading test dataset using `datasets.load_dataset`: {dataset_config}"
+            )
+            dataset = load_dataset(dataset_config, split="test")
+        else:
+            dataset = super().load_test_dataset(dataset_name, *args, **kwargs)
+        return dataset

fusion_bench/models/open_clip/__init__.py

@@ -0,0 +1,6 @@
+"""
+This module contains the support for the open_clip model.
+Modified from https://github.com/nik-dim/tall_masks/
+"""
+
+from .modeling import ClassificationHead, ImageClassifier, ImageEncoder