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

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/modeling_smile_mistral/modeling_smile_mistral.py

@@ -27,6 +27,8 @@ from transformers.models.mistral.modeling_mistral import (
     MistralRotaryEmbedding,
 )
 
+from fusion_bench.models.smile_moe.linear_from_hf_config import SmileLinear
+
 from .configuration_smile_mistral import SmileMistralConfig
 
 logger = logging.getLogger(__name__)
@@ -80,209 +82,6 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
 
-class SmileGate(nn.Module):
-    __constants__ = ["in_features", "num_experts", "k"]
-    in_features: int
-    num_experts: int
-    k: int
-    weight: Tensor
-
-    def __init__(
-        self,
-        in_features: int,
-        num_experts: int,
-        k: int,
-        device=None,
-        dtype=None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.input_features = in_features
-        self.num_experts = num_experts
-        self.k = k
-
-        self.weight = nn.Parameter(
-            torch.empty(num_experts * k, in_features, **factory_kwargs)
-        )
-
-    def forward(self, x: Tensor):
-        batch_size = x.size(0)
-        if self.num_experts == 1:
-            return torch.ones(batch_size, 1, device=x.device, dtype=x.dtype)
-
-        routing_weights = F.linear(x, self.weight).view(
-            batch_size, self.num_experts, self.k
-        )
-        routing_weights = routing_weights.norm(p=2, dim=2)
-        return routing_weights
-
-
-class SmileLinearExpert(nn.Module):
-    __constants__ = ["in_features", "out_features", "k"]
-    in_features: int
-    out_features: int
-    k: int
-
-    def __init__(
-        self,
-        in_features,
-        out_features,
-        k: int,
-        bias: bool,
-        device=None,
-        dtype=None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.k = k
-
-        self.u = nn.Parameter(torch.empty(out_features, k, **factory_kwargs))
-        self.svh = nn.Parameter(torch.empty(k, in_features, **factory_kwargs))
-
-        if bias:
-            self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs))
-        else:
-            self.register_parameter("bias", None)
-
-    def forward(self, x):
-        x = F.linear(x, self.svh)
-        x = F.linear(x, self.u, self.bias)
-        return x
-
-
-class SmileLinear(nn.Module):
-    @torch.no_grad()
-    def __init__(
-        self,
-        config: SmileMistralConfig,
-        in_features,
-        out_features,
-        bias: bool,
-        device=None,
-        dtype=None,
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-        self.num_local_experts = config.num_local_experts
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.rank_of_expert = config.rank_of_expert
-        self.rank_of_router = config.rank_of_router
-        self.in_features = in_features
-        self.out_features = out_features
-
-        # construct the gate network
-        self.gate = SmileGate(
-            in_features=in_features,
-            num_experts=self.num_local_experts,
-            k=self.rank_of_router,
-            **factory_kwargs,
-        )
-
-        # the shared linear
-        self.shared_linear = nn.Linear(
-            in_features, out_features, bias=bias, **factory_kwargs
-        )
-
-        # construct experts
-        if self.rank_of_expert > 0:
-            self.experts = nn.ModuleList(
-                [
-                    SmileLinearExpert(
-                        in_features=in_features,
-                        out_features=out_features,
-                        bias=bias,
-                        k=self.rank_of_expert,
-                        **factory_kwargs,
-                    )
-                    for _ in range(self.num_local_experts)
-                ]
-            )
-        else:
-            self.experts = nn.ModuleList(
-                [
-                    nn.Linear(in_features, out_features, bias=bias, **factory_kwargs)
-                    for _ in range(self.num_local_experts)
-                ]
-            )
-
-    def forward(self, hidden_states: Tensor):
-        pretrained_out = self.shared_linear(hidden_states)
-
-        input_shape = hidden_states.size()
-        hidden_states = hidden_states.view(-1, self.in_features)
-
-        router_logits = self.gate(hidden_states)
-        routing_weights = F.softmax(router_logits, dim=1)
-        # sample the expert according to the routing weights
-        routing_weights, selected_experts = torch.topk(
-            routing_weights, self.num_experts_per_tok, dim=-1
-        )
-        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
-
-        final_hidden_states = torch.zeros(
-            (hidden_states.size(0), self.out_features),
-            dtype=hidden_states.dtype,
-            device=hidden_states.device,
-        )
-
-        # One hot encode the selected experts to create an expert mask
-        # this will be used to easily index which expert is going to be sollicitated
-        expert_mask = torch.nn.functional.one_hot(
-            selected_experts, num_classes=self.num_local_experts
-        ).permute(2, 1, 0)
-
-        # Loop over all available experts in the model and perform the computation on each expert
-        for expert_idx in range(self.num_local_experts):
-            expert_layer = self.experts[expert_idx]
-            idx, top_x = torch.where(expert_mask[expert_idx])
-
-            # Index the correct hidden states and compute the expert hidden state for
-            # the current expert. We need to make sure to multiply the output hidden
-            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
-            current_state = hidden_states[None, top_x].reshape(-1, self.in_features)
-            if current_state.numel() == 0:
-                continue
-            current_hidden_states = (
-                expert_layer(current_state) * routing_weights[top_x, idx, None]
-            )
-
-            # However `index_add_` only support torch tensors for indexing so we'll use
-            # the `top_x` tensor here.
-            final_hidden_states.index_add_(
-                0, top_x, current_hidden_states.to(hidden_states.dtype)
-            )
-        final_hidden_states = final_hidden_states.reshape(
-            *input_shape[:-1], self.out_features
-        )
-        final_hidden_states = pretrained_out + final_hidden_states
-        return final_hidden_states
-
-    @property
-    def weight(self):
-        """
-        Mimic linear layer. Bacause in some cases, user might indicate the device (or dtype of parameters) of the linear layer using `linear_layer.weight.device`
-        """
-        return self.shared_linear.weight
-
-    @property
-    def bias(self):
-        return self.shared_linear.bias
-
-    def __repr__(self):
-        return (
-            f"SingularMoELinear("
-            f"in_features={self.shared_linear.in_features}, "
-            f"out_features={self.shared_linear.out_features}, "
-            f"num_local_experts={self.num_local_experts}, "
-            f"num_experts_per_tok={self.num_experts_per_tok}, "
-            f"rank_of_router={self.rank_of_router}, "
-            f"rank_of_expert={self.rank_of_expert}"
-            f")"
-        )
-
-
 class SmileMistralAttention(nn.Module):
     """
     Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer

fusion_bench/models/modeling_smile_qwen2/__init__.py

@@ -0,0 +1,8 @@
+from . import register
+from .configuration_smile_qwen2 import SmileQwen2Config
+from .modeling_smile_qwen2 import (
+    SmileQwen2ForCausalLM,
+    SmileQwen2ForQuestionAnswering,
+    SmileQwen2ForSequenceClassification,
+    SmileQwen2Model,
+)

fusion_bench/models/modeling_smile_qwen2/configuration_smile_qwen2.py

@@ -0,0 +1,21 @@
+from transformers import PretrainedConfig
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+
+class SmileQwen2Config(Qwen2Config):
+    model_type = "smile_qwen2"
+
+    def __init__(
+        self,
+        num_experts_per_tok: int = 1,
+        rank_of_router: int = None,
+        rank_of_expert: int = None,
+        num_local_experts: int = None,
+        **kwargs,
+    ):
+        self.num_experts_per_tok = num_experts_per_tok
+        self.rank_of_router = rank_of_router
+        self.rank_of_expert = rank_of_expert
+        self.num_local_experts = num_local_experts
+
+        super().__init__(**kwargs)