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

fusion_bench/models/smile_moe/utils/__init__.py

@@ -0,0 +1,24 @@
+from typing import List
+
+import torch
+from torch import Tensor
+
+from .svd_utils import svd
+
+__all__ = ["svd_utils", "_is_all_zeros"]
+
+
+def _is_all_zeros(tensor: Tensor | List[Tensor]) -> bool:
+    """
+    Check if a tensor or a list of tensors are all zeros.
+
+    Args:
+        tensor (Tensor | List[Tensor]): A tensor or a list of tensors.
+
+    Returns:
+        bool: True if all elements are zeros, False otherwise.
+    """
+    if isinstance(tensor, Tensor):
+        return torch.allclose(tensor, torch.zeros_like(tensor))
+    else:
+        return all(_is_all_zeros(t) for t in tensor)

fusion_bench/models/smile_moe/utils/svd_utils.py

@@ -0,0 +1,46 @@
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+
+
+def _svd(w: Tensor, full_matrices: bool = True) -> Tuple[Tensor, Tensor, Tensor]:
+    """
+    Perform Singular Value Decomposition (SVD) on a tensor.
+
+    Args:
+        w (Tensor): The input tensor.
+        full_matrices (bool): Whether to compute the full-sized U and V matrices.
+
+    Returns:
+        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
+    """
+    u, s, vh = torch.linalg.svd(
+        w, full_matrices=full_matrices, driver="gesvd" if w.is_cuda else None
+    )
+    v = vh.T
+    return u, s, v
+
+
+def svd(
+    w: Tensor,
+    full_matrices: bool = True,
+    accelerator: Optional[Union[torch.device, str]] = None,
+) -> Tuple[Tensor, Tensor, Tensor]:
+    """
+    Perform SVD on a tensor, optionally using a specified accelerator.
+
+    Args:
+        w (Tensor): The input tensor.
+        full_matrices (bool): Whether to compute the full-sized U and V matrices.
+        accelerator (Optional[Union[torch.device, str]]): The device to perform the computation on.
+
+    Returns:
+        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
+    """
+    if accelerator is None:
+        return _svd(w, full_matrices=full_matrices)
+    original_device = w.device
+    w = w.to(accelerator)
+    u, s, v = _svd(w)
+    return u.to(original_device), s.to(original_device), v.to(original_device)

fusion_bench/scripts/nyuv2_mtl_train.py

@@ -1,5 +1,5 @@
 R"""
-This script is used to train a multi-task learning (MTL) model on the NYUv2 dataset. 
+This script is used to train a multi-task learning (MTL) model on the NYUv2 dataset.
 """
 
 import importlib

fusion_bench/taskpool/__init__.py

@@ -10,12 +10,14 @@ _import_structure = {
     "clip_vision": [
         "CLIPVisionModelTaskPool",
         "SparseWEMoECLIPVisionModelTaskPool",
-        "RankoneWEMoECLIPVisionModelTaskPool",
+        "RankoneMoECLIPVisionModelTaskPool",
     ],
     "dummy": ["DummyTaskPool"],
     "gpt2_text_classification": ["GPT2TextClassificationTaskPool"],
-    "nyuv2_taskpool": ["NYUv2TaskPool"],
     "llama": ["LlamaTestGenerationTaskPool"],
+    "lm_eval_harness": ["LMEvalHarnessTaskPool"],
+    "nyuv2_taskpool": ["NYUv2TaskPool"],
+    "openclip_vision": ["OpenCLIPVisionModelTaskPool"],
 }
 
 
@@ -23,13 +25,15 @@ if TYPE_CHECKING:
     from .base_pool import BaseTaskPool
     from .clip_vision import (
         CLIPVisionModelTaskPool,
-        RankoneWEMoECLIPVisionModelTaskPool,
+        RankoneMoECLIPVisionModelTaskPool,
         SparseWEMoECLIPVisionModelTaskPool,
     )
     from .dummy import DummyTaskPool
     from .gpt2_text_classification import GPT2TextClassificationTaskPool
     from .llama import LlamaTestGenerationTaskPool
+    from .lm_eval_harness import LMEvalHarnessTaskPool
     from .nyuv2_taskpool import NYUv2TaskPool
+    from .openclip_vision import OpenCLIPVisionModelTaskPool
 
 else:
     sys.modules[__name__] = LazyImporter(

fusion_bench/taskpool/clip_vision/__init__.py

@@ -1,4 +1,5 @@
 # flake8: noqa F401
 from .clip_rankone_moe_taskpool import RankoneMoECLIPVisionModelTaskPool
+from .clip_smile_taskpool import SmileCLIPVisionModelTaskPool
 from .clip_sparse_wemoe_taskpool import SparseWEMoECLIPVisionModelTaskPool
 from .taskpool import CLIPVisionModelTaskPool

fusion_bench/taskpool/clip_vision/clip_rankone_moe_taskpool.py

@@ -12,36 +12,7 @@ from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.models.rankone_moe import RankOneMoE
 
 from .taskpool import CLIPVisionModelTaskPool
-
-
-class LayerWiseRoutingWeightSaver:
-    def __init__(self, save_path: Path, max_num: Optional[int] = None):
-        self.save_path = save_path
-        self.max_num = max_num
-        self.routing_weights = []
-
-    def __call__(self, module, input: Tuple[Tensor], output: Tensor):
-        assert isinstance(output, Tensor), "Output is expected to be a Tensor"
-        # (batch_size, num_tokens, num_experts)
-        routing_weights = output.detach().cpu()
-        if self.max_num is not None and self.max_num > 0:
-            if len(self.routing_weights) > self.max_num:
-                return
-            elif routing_weights.size(0) + len(self.routing_weights) > self.max_num:
-                self.routing_weights.append(
-                    routing_weights[: self.max_num - len(self.routing_weights)]
-                )
-            else:
-                self.routing_weights.append(routing_weights)
-        else:
-            self.routing_weights.append(routing_weights)
-
-    def save_routing_weights(self):
-        routing_weights = torch.cat(self.routing_weights, dim=0)
-        if self.save_path is not None:
-            self.save_path.parent.mkdir(parents=True, exist_ok=True)
-            print(f"Saving routing weights to {self.save_path}")
-            torch.save(routing_weights, self.save_path)
+from .utils.routing_analysis_utils import LayerWiseRoutingWeightSaver
 
 
 class RankoneMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
@@ -109,4 +80,5 @@ class RankoneMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
             # remove hooks for saving layer-wise routing weights
             for i, handle in self._layer_wise_routing_weights_save_hook_handles.items():
                 self._layer_wise_routing_weights_save_hooks[i].save_routing_weights()
+                self._layer_wise_routing_weights_save_hook_handles.pop(i)
                 handle.remove()

fusion_bench/taskpool/clip_vision/clip_smile_taskpool.py

@@ -0,0 +1,102 @@
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor
+from torch.utils.hooks import RemovableHandle
+from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
+from transformers.models.clip.modeling_clip import CLIPVisionTransformer
+
+from fusion_bench.method.smile_upscaling import SmileMoELinear
+from fusion_bench.models.hf_clip import HFCLIPClassifier
+
+from .taskpool import CLIPVisionModelTaskPool
+from .utils.routing_analysis_utils import LayerWiseRoutingWeightSaver
+
+
+class SmileCLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
+
+    # hooks and handles for saving layer-wise routing weights
+    _layer_wise_routing_weights_save_hooks: Dict[Any, LayerWiseRoutingWeightSaver] = {}
+    _layer_wise_routing_weights_save_hook_handles: Dict[Any, RemovableHandle] = {}
+
+    def __init__(
+        self,
+        linear_module_names: Union[List[str], str],
+        layer_wise_routing_weights_save_path: Optional[str],
+        layer_wise_routing_weights_max_num: Optional[int] = None,
+        **kwargs,
+    ):
+        """
+        Initialize the SMILECLIPVisionModelTaskPool.
+
+        Args:
+            linear_module_names (Union[List[str], str]): The names of the linear modules to save the layer-wise routing weights for.
+            layer_wise_routing_weights_save_path (Optional[str]): The path to save the layer-wise routing weights.
+            layer_wise_routing_weights_max_num (Optional[int]): The maximum number of layer-wise routing weights to save.
+        """
+        # linear module names
+        assert linear_module_names is not None, "linear_module_names must be provided"
+        self.linear_module_names = (
+            [linear_module_names]
+            if isinstance(linear_module_names, str)
+            else list(linear_module_names)
+        )
+        # save path for layer-wise routing weights
+        self._layer_wise_routing_weights_save_path = (
+            layer_wise_routing_weights_save_path
+        )
+        self.layer_wise_routing_weights_save_path = (
+            Path(layer_wise_routing_weights_save_path)
+            if layer_wise_routing_weights_save_path is not None
+            else None
+        )
+        self.layer_wise_routing_weights_max_num = layer_wise_routing_weights_max_num
+        super().__init__(**kwargs)
+
+    def on_task_evaluation_begin(self, classifier: HFCLIPClassifier, task_name: str):
+        super().on_task_evaluation_begin(classifier, task_name)
+        if self.layer_wise_routing_weights_save_path is not None:
+            # setup hooks for saving layer-wise routing weights
+            assert isinstance(
+                classifier.clip_model.vision_model,
+                (CLIPVisionTransformer, CLIPVisionModel),
+            ), "Vision model is expected to be a CLIPVisionTransformer"
+            vision_model = classifier.clip_model.vision_model
+            if isinstance(vision_model, CLIPVisionModel):
+                vision_model = vision_model.vision_model
+                # assign forward hooks for each layer
+
+            for i, layer in enumerate(vision_model.encoder.layers):
+                for linear_module_name in self.linear_module_names:
+                    linear_module = layer.get_submodule(linear_module_name)
+                    assert isinstance(
+                        linear_module,
+                        (SmileMoELinear),
+                    ), f"Linear module is expected to be a SmileMoELinear, but got {type(linear_module)}"
+                    # layer-wise routing weights
+                    hook = LayerWiseRoutingWeightSaver(
+                        self.layer_wise_routing_weights_save_path
+                        / task_name
+                        / f"layer_{i}_{linear_module_name}.pt",
+                        max_num=self.layer_wise_routing_weights_max_num,
+                    )
+                    self._layer_wise_routing_weights_save_hooks[
+                        (i, linear_module_name)
+                    ] = hook
+                    self._layer_wise_routing_weights_save_hook_handles[
+                        (i, linear_module_name)
+                    ] = linear_module.gate.register_forward_hook(hook)
+
+    def on_task_evaluation_end(self):
+        super().on_task_evaluation_end()
+        if self.layer_wise_routing_weights_save_path is not None:
+            # remove hooks for saving layer-wise routing weights
+            for (
+                key,
+                handle,
+            ) in self._layer_wise_routing_weights_save_hook_handles.items():
+                self._layer_wise_routing_weights_save_hooks[key].save_routing_weights()
+                self._layer_wise_routing_weights_save_hook_handles.pop(key)
+                handle.remove()

fusion_bench/taskpool/clip_vision/clip_sparse_wemoe_taskpool.py

@@ -15,36 +15,7 @@ from fusion_bench.models.sparse_we_moe import (
 )
 
 from .taskpool import CLIPVisionModelTaskPool
-
-
-class LayerWiseRoutingWeightSaver:
-    def __init__(self, save_path: Path, max_num: Optional[int] = None):
-        self.save_path = save_path
-        self.max_num = max_num
-        self.routing_weights = []
-
-    def __call__(self, module, input: Tuple[Tensor], output: Tensor):
-        assert isinstance(output, Tensor), "Output is expected to be a Tensor"
-        # (batch_size, num_tokens, num_experts)
-        routing_weights = output.detach().cpu()
-        if self.max_num is not None and self.max_num > 0:
-            if len(self.routing_weights) > self.max_num:
-                return
-            elif routing_weights.size(0) + len(self.routing_weights) > self.max_num:
-                self.routing_weights.append(
-                    routing_weights[: self.max_num - len(self.routing_weights)]
-                )
-            else:
-                self.routing_weights.append(routing_weights)
-        else:
-            self.routing_weights.append(routing_weights)
-
-    def save_routing_weights(self):
-        routing_weights = torch.cat(self.routing_weights, dim=0)
-        if self.save_path is not None:
-            self.save_path.parent.mkdir(parents=True, exist_ok=True)
-            print(f"Saving routing weights to {self.save_path}")
-            torch.save(routing_weights, self.save_path)
+from .utils.routing_analysis_utils import LayerWiseRoutingWeightSaver
 
 
 class SparseWEMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
@@ -117,4 +88,5 @@ class SparseWEMoECLIPVisionModelTaskPool(CLIPVisionModelTaskPool):
             # remove hooks for saving layer-wise routing weights
             for i, handle in self._layer_wise_routing_weights_save_hook_handles.items():
                 self._layer_wise_routing_weights_save_hooks[i].save_routing_weights()
+                self._layer_wise_routing_weights_save_hook_handles.pop(i)
                 handle.remove()

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -32,8 +32,7 @@ from fusion_bench.mixins import LightningFabricMixin
 from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.taskpool import BaseTaskPool
 from fusion_bench.tasks.clip_classification import get_classnames_and_templates
-from fusion_bench.utils import instantiate
-from fusion_bench.utils.parameters import count_parameters
+from fusion_bench.utils import count_parameters, instantiate
 
 if TYPE_CHECKING:
     from fusion_bench.models.surgery.surgerymodelwrapper import SurgeryModelWrapper

fusion_bench/taskpool/clip_vision/utils/routing_analysis_utils.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+
+def _number_of_samples(routing_weights: List[Tensor]):
+    count = 0
+    for routing_weight in routing_weights:
+        count += routing_weight.size(0)
+    return count
+
+
+class LayerWiseRoutingWeightSaver:
+    """
+    A hook for saving layer-wise routing weights.
+    """
+
+    save_path: Path
+    "The path to save the layer-wise routing weights."
+    max_num: Optional[int]
+    "The maximum number of layer-wise routing weights to save. If None, all routing weights will be saved."
+    routing_weights: List[Tensor]
+    "The list of layer-wise routing weights."
+
+    def __init__(self, save_path: Path, max_num: Optional[int] = None):
+        """
+        Args:
+            save_path (Path): The path to save the layer-wise routing weights.
+            max_num (Optional[int]): The maximum number of layer-wise routing weights to save. If None, all routing weights will be saved.
+        """
+        self.save_path = save_path
+        self.max_num = max_num
+        self.routing_weights = []
+
+    def __call__(self, module, input: Tuple[Tensor], output: Tensor):
+        assert isinstance(output, Tensor), "Output is expected to be a Tensor"
+        # (batch_size, num_tokens, num_experts)
+        routing_weights = output.detach().cpu()
+        if self.max_num is not None and self.max_num > 0:
+            if _number_of_samples(self.routing_weights) > self.max_num:
+                return
+            elif (
+                routing_weights.size(0) + _number_of_samples(self.routing_weights)
+                > self.max_num
+            ):
+                self.routing_weights.append(
+                    routing_weights[
+                        : self.max_num - _number_of_samples(self.routing_weights)
+                    ]
+                )
+            else:
+                self.routing_weights.append(routing_weights)
+        else:
+            self.routing_weights.append(routing_weights)
+
+    def save_routing_weights(self):
+        routing_weights = torch.cat(self.routing_weights, dim=0)
+        if self.save_path is not None:
+            self.save_path.parent.mkdir(parents=True, exist_ok=True)
+            print(
+                f"Saving routing weights to {self.save_path}. Size: {routing_weights.size()}"
+            )
+            torch.save(routing_weights, self.save_path)

fusion_bench/taskpool/gpt2_text_classification.py

@@ -139,11 +139,40 @@ class GPT2TextClassificationTaskPool(BaseTaskPool, LightningFabricMixin):
         return dataloader
 
     @override
-    def evaluate(self, model: GPT2Model):
+    def evaluate(self, model: GPT2Model, name: str = None):
+        """Evaluate the model on the test datasets.
+
+        Args:
+            model (GPT2Model): The model to evaluate.
+            name (str, optional): The name of the model. Defaults to None. This is used to identify the model in the report.
+
+        Returns:
+            dict: A dictionary containing the evaluation results for each task.
+        """
         report = {}
+        if name is not None:
+            report["name"] = name
         for task_name in (pbar := tqdm(self._test_datasets, desc="Evaluating tasks")):
             pbar.set_description(f"Evaluating task {task_name}")
             dataloader = self.get_test_dataloader(task_name)
             result = self.evaluate_single_task(task_name, model, dataloader)
             report[task_name] = result
+
+        # calculate the average accuracy and loss
+        if "average" not in report:
+            report["average"] = {}
+            accuracies = [
+                value["accuracy"]
+                for key, value in report.items()
+                if isinstance(value, dict) and "accuracy" in value
+            ]
+            if len(accuracies) > 0:
+                average_accuracy = sum(accuracies) / len(accuracies)
+                report["average"]["accuracy"] = average_accuracy
+            losses = [value["loss"] for key, value in report.items() if "loss" in value]
+            if len(losses) > 0:
+                average_loss = sum(losses) / len(losses)
+                report["average"]["loss"] = average_loss
+
+        log.info(f"Evaluation Result: {report}")
         return report

fusion_bench/taskpool/lm_eval_harness/__init__.py

@@ -0,0 +1,3 @@
+from .taskpool import LMEvalHarnessTaskPool
+
+__all__ = ["LMEvalHarnessTaskPool"]

fusion_bench/taskpool/lm_eval_harness/taskpool.py

@@ -0,0 +1,87 @@
+import logging
+import os
+from typing import List, Literal, Optional, Union, TYPE_CHECKING
+
+import lightning.fabric
+import lm_eval
+import lm_eval.models
+from lm_eval.__main__ import check_argument_types, cli_evaluate, setup_parser
+from omegaconf import DictConfig, ListConfig
+
+from fusion_bench import BaseTaskPool
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.utils.strenum import _version
+
+
+log = logging.getLogger(__name__)
+
+
+class LMEvalHarnessTaskPool(BaseTaskPool, LightningFabricMixin):
+    def __init__(
+        self,
+        tasks: Union[str, List[str]],
+        apply_chat_template: bool = False,
+        include_path: Optional[str] = None,
+        batch_size: int = 1,
+        metadata: Optional[DictConfig] = None,
+        verbosity: Optional[
+            Literal["CRITICAL", "ERROR", "WARNING", "INFO", "DEBUG"]
+        ] = None,
+        output_path: Optional[str] = None,
+        log_samples: bool = False,
+        _usage_: Optional[str] = None,
+        _version_: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(_usage_=_usage_, _version_=_version_)
+        self.tasks = tasks
+        self.include_path = include_path
+        self.batch_size = batch_size
+        self.metadata = metadata
+        self.apply_chat_template = apply_chat_template
+        self.verbosity = verbosity
+        self.kwargs = kwargs
+        self.output_path = output_path
+        self.log_samples = log_samples
+
+    def evaluate(self, model, *command_line_args, **kwargs):
+        command_line_args = []
+        if self.include_path is not None:
+            command_line_args.extend(["--include_path", self.include_path])
+        if isinstance(self.tasks, (list, ListConfig)):
+            command_line_args.extend(["--tasks", ",".join(self.tasks)])
+        elif isinstance(self.tasks, str):
+            command_line_args.extend(["--tasks", self.tasks])
+        if self.apply_chat_template:
+            command_line_args.extend(
+                ["--apply_chat_template", str(self.apply_chat_template)]
+            )
+        if self.batch_size is not None:
+            command_line_args.extend(["--batch_size", str(self.batch_size)])
+        if self.verbosity is not None:
+            command_line_args.extend(["--verbosity", str(self.verbosity)])
+        if self.metadata is not None:
+            command_line_args.extend(["--metadata", str(self.metadata)])
+        if self.output_path is None:
+            command_line_args.extend(
+                [
+                    "--output_path",
+                    os.path.join(self.log_dir, "lm_eval_results"),
+                ]
+            )
+        else:
+            command_line_args.extend(["--output_path", self.output_path])
+        if self.log_samples:
+            command_line_args.extend(["--log_samples"])
+        for key, value in kwargs.items():
+            command_line_args.extend([f"--{key}", str(value)])
+
+        parser = setup_parser()
+        check_argument_types(parser)
+        args = parser.parse_args(args=command_line_args)
+        log.info("LM-Eval Harness arguments:\n%s", args)
+
+        if not lightning.fabric.is_wrapped(model):
+            model = self.fabric.setup(model)
+        args.model = lm_eval.models.huggingface.HFLM(pretrained=model)
+        cli_evaluate(args)

fusion_bench/taskpool/openclip_vision/__init__.py

@@ -0,0 +1 @@
+from .openclip_taskpool import OpenCLIPVisionModelTaskPool