fusion-bench 0.2.22__tar.gz → 0.2.24__tar.gz

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.22
+Version: 0.2.24
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 Project-URL: Repository, https://github.com/tanganke/fusion_bench
@@ -23,12 +23,19 @@ Requires-Dist: rich
 Requires-Dist: scipy
 Requires-Dist: h5py
 Requires-Dist: pytest
+Requires-Dist: joblib
+Requires-Dist: bidict
 Requires-Dist: transformers!=4.49
 Requires-Dist: pillow!=11.2.1
 Provides-Extra: lm-eval-harness
 Requires-Dist: lm-eval; extra == "lm-eval-harness"
 Requires-Dist: immutabledict; extra == "lm-eval-harness"
 Requires-Dist: langdetect; extra == "lm-eval-harness"
+Requires-Dist: rich-run; extra == "lm-eval-harness"
+Provides-Extra: docs
+Requires-Dist: mkdocs; extra == "docs"
+Requires-Dist: mkdocs-material; extra == "docs"
+Requires-Dist: mkdocstrings[python]; extra == "docs"
 Dynamic: license-file
 
 <div align='center'>
@@ -151,7 +158,7 @@ This will install the latest version of fusion-bench and the dependencies requir
 Documentation for using LM-Eval Harness within FusionBench framework can be found at [this online documentation](https://tanganke.github.io/fusion_bench/taskpool/lm_eval_harness) or in the [`docs/taskpool/lm_eval_harness.md`](docs/taskpool/lm_eval_harness.md) markdown file.
 
 > [!TIP]
-> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/causal_lm) or in the [`docs/modelpool/causal_lm.md`](docs/modelpool/causal_lm.md) markdown file.
+> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/llm) or in the [`docs/modelpool/llm/index.md`](docs/modelpool/llm/index.md) markdown file.
 
 ## Introduction to Deep Model Fusion
 
@@ -179,7 +186,7 @@ The project is structured as follows:
   - `taskpool`: configuration files for the task pool.
   - `model`: configuration files for the models.
   - `dataset`: configuration files for the datasets.
-- `docs/`: documentation for the benchmark. We use [mkdocs](https://www.mkdocs.org/) to generate the documentation. Start the documentation server locally with `mkdocs serve`. The required packages can be installed with `pip install -r mkdocs-requirements.txt`.
+- `docs/`: documentation for the benchmark. We use [mkdocs](https://www.mkdocs.org/) to generate the documentation. Start the documentation server locally with `mkdocs serve`. The required packages can be installed with `pip install -e ".[docs]"`.
 - `examples/`: example scripts for running some of the experiments.
   > **naming convention**: `examples/{method_name}/` contains the files such as bash scripts and jupyter notebooks for the specific method.
 - `tests/`: unit tests for the benchmark.

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/README.md

@@ -118,7 +118,7 @@ This will install the latest version of fusion-bench and the dependencies requir
 Documentation for using LM-Eval Harness within FusionBench framework can be found at [this online documentation](https://tanganke.github.io/fusion_bench/taskpool/lm_eval_harness) or in the [`docs/taskpool/lm_eval_harness.md`](docs/taskpool/lm_eval_harness.md) markdown file.
 
 > [!TIP]
-> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/causal_lm) or in the [`docs/modelpool/causal_lm.md`](docs/modelpool/causal_lm.md) markdown file.
+> Documentation for merging large language models using FusionBench can be found at [this online documentation](https://tanganke.github.io/fusion_bench/modelpool/llm) or in the [`docs/modelpool/llm/index.md`](docs/modelpool/llm/index.md) markdown file.
 
 ## Introduction to Deep Model Fusion
 
@@ -146,7 +146,7 @@ The project is structured as follows:
   - `taskpool`: configuration files for the task pool.
   - `model`: configuration files for the models.
   - `dataset`: configuration files for the datasets.
-- `docs/`: documentation for the benchmark. We use [mkdocs](https://www.mkdocs.org/) to generate the documentation. Start the documentation server locally with `mkdocs serve`. The required packages can be installed with `pip install -r mkdocs-requirements.txt`.
+- `docs/`: documentation for the benchmark. We use [mkdocs](https://www.mkdocs.org/) to generate the documentation. Start the documentation server locally with `mkdocs serve`. The required packages can be installed with `pip install -e ".[docs]"`.
 - `examples/`: example scripts for running some of the experiments.
   > **naming convention**: `examples/{method_name}/` contains the files such as bash scripts and jupyter notebooks for the specific method.
 - `tests/`: unit tests for the benchmark.

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/__init__.py

@@ -32,6 +32,10 @@ from .models import (
 from .programs import BaseHydraProgram
 from .taskpool import BaseTaskPool
 from .utils import (
+    BoolStateDictType,
+    LazyStateDict,
+    StateDictType,
+    TorchModelType,
     cache_with_joblib,
     get_rankzero_logger,
     import_object,

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/compat/method/__init__.py

@@ -1,4 +1,5 @@
 import warnings
+from typing import Any, List, Type
 
 from omegaconf import DictConfig
 
@@ -76,7 +77,9 @@ class AlgorithmFactory:
         return algorithm_cls(method_config)
 
     @staticmethod
-    def register_algorithm(name: str, algorithm_cls):
+    def register_algorithm(
+        name: str, algorithm_cls: Type[ModelFusionAlgorithm]
+    ) -> None:
         """
         Register a new algorithm with the factory.
 
@@ -87,7 +90,7 @@ class AlgorithmFactory:
         AlgorithmFactory._aglorithms[name] = algorithm_cls
 
     @classmethod
-    def available_algorithms(cls):
+    def available_algorithms(cls) -> List[str]:
         """
         Get a list of available algorithms.
 

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/compat/method/base_algorithm.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from omegaconf import DictConfig
 
 if TYPE_CHECKING:
+    from fusion_bench import BaseModelPool
     from fusion_bench.programs.base_program import BaseHydraProgram
 
 __all__ = ["ModelFusionAlgorithm"]
@@ -51,7 +52,7 @@ class ModelFusionAlgorithm(ABC):
         pass
 
     @abstractmethod
-    def run(self, modelpool):
+    def run(self, modelpool: "BaseModelPool") -> Any:
         """
         Fuse the models in the given model pool.
 

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/compat/modelpool/base_pool.py

@@ -42,7 +42,7 @@ class ModelPool(ABC):
             ), "Duplicate model names found in model pool"
             self._model_names = model_names
 
-    def __len__(self):
+    def __len__(self) -> int:
         """
         Return the number of models in the model pool, exclude special models such as `_pretrained_`.
 
@@ -66,7 +66,7 @@ class ModelPool(ABC):
         return names
 
     @property
-    def has_pretrained(self):
+    def has_pretrained(self) -> bool:
         """
         Check if the pretrained model is available in the model pool.
 
@@ -78,7 +78,7 @@ class ModelPool(ABC):
                 return True
         return False
 
-    def get_model_config(self, model_name: str):
+    def get_model_config(self, model_name: str) -> Dict:
         """
         Retrieves the configuration for a specific model from the model pool.
 

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/compat/taskpool/clip_image_classification.py

@@ -169,7 +169,7 @@ class CLIPImageClassificationTaskPool(TaskPool):
             self._fabric = L.Fabric(devices=1)
             self._fabric.launch()
 
-        # CLIPVisionModel works the same with CLIPVisonTransformer, so we can use it directly
+        # CLIPVisionModel works the same with CLIPVisionTransformer, so we can use it directly
         self.clip_model.vision_model = model
         report = {}
         training_params, all_params = count_parameters(model)

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/dataset/gpt2_glue.py

@@ -121,7 +121,7 @@ class TokenizedGLUE:
 
     def load_dataset(
         self, name: Literal["mrpc", "mnli", "cola", "sst2", "qnli", "qqp", "rte"]
-    ):
+    ) -> Dataset:
         """
         Load and tokenize a GLUE dataset.
 

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/method/__init__.py

@@ -26,11 +26,14 @@ _import_structure = {
     "linear": [
         "ExPOAlgorithm",
         "ExPOAlgorithmForLlama",
+        "SimpleAverageForCausalLM",
         "SimpleAverageForLlama",
+        "TaskArithmeticForCausalLM",
         "TaskArithmeticForLlama",
         "LinearInterpolationAlgorithm",
+        "TiesMergingForCausalLM",
     ],
-    "slerp": ["SlerpMergeAlgorithm"],
+    "slerp": ["SlerpMergeAlgorithm", "SlerpForCausalLM"],
     "simple_average": ["SimpleAverageAlgorithm"],
     "weighted_average": ["WeightedAverageAlgorithm", "WeightedAverageForLLama"],
     "task_arithmetic": ["TaskArithmeticAlgorithm"],
@@ -71,6 +74,8 @@ _import_structure = {
     ],
     "fw_merging": ["FrankWolfeHardAlgorithm", "FrankWolfeSoftAlgorithm"],
     "tall_mask": ["TallMaskTaskArithmeticAlgorithm"],
+    "model_stock": ["ModelStock"],
+    "wudi": ["wudi_merging", "WUDIMerging"],
     # plug-and-play model merging methods
     "concrete_subspace": [
         "ConcreteTaskArithmeticAlgorithmForCLIP",
@@ -183,8 +188,11 @@ if TYPE_CHECKING:
         ExPOAlgorithm,
         ExPOAlgorithmForLlama,
         LinearInterpolationAlgorithm,
+        SimpleAverageForCausalLM,
         SimpleAverageForLlama,
+        TaskArithmeticForCausalLM,
         TaskArithmeticForLlama,
+        TiesMergingForCausalLM,
     )
     from .lm_finetune import *
     from .mixture_of_experts import (
@@ -194,6 +202,7 @@ if TYPE_CHECKING:
         MixtralUpscalingAlgorithm,
     )
     from .model_recombination import ModelRecombinationAlgorithm
+    from .model_stock import ModelStock
     from .opcm import OPCMForCLIP
     from .pruning import (
         MagnitudeDiffPruningAlgorithm,
@@ -213,7 +222,7 @@ if TYPE_CHECKING:
         RegMeanAlgorithmPlusPlus,
     )
     from .simple_average import SimpleAverageAlgorithm
-    from .slerp import SlerpMergeAlgorithm
+    from .slerp import SlerpForCausalLM, SlerpMergeAlgorithm
     from .smile_upscaling import (
         SingularProjectionMergingAlgorithm,
         SmileUpscalingAlgorithm,
@@ -236,6 +245,7 @@ if TYPE_CHECKING:
         FlanT5WeightEnsemblingMoEAlgorithm,
     )
     from .weighted_average import WeightedAverageAlgorithm, WeightedAverageForLLama
+    from .wudi import WUDIMerging, wudi_merging
 
 else:
     sys.modules[__name__] = LazyImporter(

{fusion_bench-0.2.22 → fusion_bench-0.2.24}/fusion_bench/method/analysis/task_vector_cos_similarity.py

@@ -11,7 +11,7 @@ from torch import nn
 from tqdm.auto import tqdm
 
 from fusion_bench.method import BaseAlgorithm
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.parameters import (
     StateDictType,
@@ -23,14 +23,50 @@ from fusion_bench.utils.state_dict_arithmetic import state_dict_sub
 log = logging.getLogger(__name__)
 
 
-class TaskVectorCosSimilarity(BaseAlgorithm, LightningFabricMixin):
+@auto_register_config
+class TaskVectorCosSimilarity(
+    LightningFabricMixin,
+    BaseAlgorithm,
+):
     """
-    This class is similar to the Dummy algorithm,
-    but it also print (or save) the cosine similarity matrix between the task vectors of the models in the model pool.
+    Computes and analyzes cosine similarity between task vectors of models in a model pool.
+
+    This algorithm extracts task vectors from fine-tuned models by computing the difference
+    between their parameters and a pretrained base model. It then calculates the pairwise
+    cosine similarity between all task vectors to understand the relationships and overlap
+    between different tasks.
+
+    The task vector for a model is defined as:
+        task_vector = finetuned_model_params - pretrained_model_params
+
+    Args:
+        plot_heatmap (bool): Whether to generate and save a heatmap visualization
+        trainable_only (bool, optional): If True, only consider trainable parameters
+            when computing task vectors. Defaults to True.
+        max_points_per_model (int, optional): Maximum number of parameters to sample
+            per model for memory efficiency. If None, uses all parameters.
+        output_path (str, optional): Directory to save outputs. If None, uses the
+            fabric logger directory.
+
+    Outputs:
+        - task_vector_cos_similarity.csv: Pairwise cosine similarity matrix
+        - task_vector_cos_similarity.pdf: Heatmap visualization (if plot_heatmap=True)
+
+    Returns:
+        The pretrained model from the model pool.
+
+    Example:
+        ```python
+        >>> algorithm = TaskVectorCosSimilarity(
+        ...     plot_heatmap=True,
+        ...     trainable_only=True,
+        ...     output_path="/path/to/outputs"
+        ... )
+        >>> result = algorithm.run(modelpool)
+        ```
     """
 
     _config_mapping = BaseAlgorithm._config_mapping | {
-        "plot_heatmap": "plot_heatmap",
         "_output_path": "output_path",
     }
 
@@ -42,11 +78,8 @@ class TaskVectorCosSimilarity(BaseAlgorithm, LightningFabricMixin):
         output_path: Optional[str] = None,
         **kwargs,
     ):
-        self.plot_heatmap = plot_heatmap
-        self.trainable_only = trainable_only
-        self.max_points_per_model = max_points_per_model
-        self._output_path = output_path
         super().__init__(**kwargs)
+        self._output_path = output_path
 
     @property
     def output_path(self):
@@ -57,6 +90,22 @@ class TaskVectorCosSimilarity(BaseAlgorithm, LightningFabricMixin):
 
     @torch.no_grad()
     def run(self, modelpool: BaseModelPool):
+        """
+        Execute the task vector cosine similarity analysis.
+
+        This method:
+        1. Loads the pretrained base model from the model pool
+        2. Computes task vectors for each fine-tuned model
+        3. Calculates pairwise cosine similarities between all task vectors
+        4. Saves the similarity matrix as a CSV file
+        5. Optionally generates and saves a heatmap visualization
+
+        Args:
+            modelpool (BaseModelPool): Pool containing pretrained and fine-tuned models
+
+        Returns:
+            nn.Module: The pretrained model from the model pool
+        """
         pretrained_model = modelpool.load_pretrained_model()
 
         task_vectors = []
@@ -103,11 +152,14 @@ class TaskVectorCosSimilarity(BaseAlgorithm, LightningFabricMixin):
 
     def _plot_heatmap(self, data: pd.DataFrame):
         """
-        This function plots a heatmap of the provided data using seaborn.
+        Generate and save a heatmap visualization of the cosine similarity matrix.
+
+        Creates a color-coded heatmap showing pairwise cosine similarities between
+        task vectors. The heatmap is saved as a PDF file in the output directory.
 
         Args:
-            data (pd.DataFrame): A pandas DataFrame containing the data to be plotted.
-            figsize (tuple): A tuple specifying the size of the figure. Default is (4, 3).
+            data (pd.DataFrame): Symmetric matrix of cosine similarities between
+                task vectors, with model names as both index and columns.
 
         Returns:
             None
@@ -141,6 +193,26 @@ class TaskVectorCosSimilarity(BaseAlgorithm, LightningFabricMixin):
     def get_task_vector(
         self, pretrained_model: nn.Module, finetuned_model: nn.Module
     ) -> torch.Tensor:
+        """
+        Compute the task vector for a fine-tuned model.
+
+        The task vector represents the parameter changes from pretraining to
+        fine-tuning and is computed as:
+            task_vector = finetuned_params - pretrained_params
+
+        Args:
+            pretrained_model (nn.Module): The base pretrained model
+            finetuned_model (nn.Module): The fine-tuned model for a specific task
+
+        Returns:
+            torch.Tensor: Flattened task vector containing parameter differences.
+                If max_points_per_model is set, the vector may be downsampled.
+
+        Note:
+            - Converts parameters to float64 for numerical precision
+            - Supports optional downsampling for memory efficiency
+            - Uses only trainable parameters if trainable_only=True
+        """
         task_vector = state_dict_sub(
             self.get_state_dict(finetuned_model),
             self.get_state_dict(pretrained_model),
@@ -166,6 +238,17 @@ class TaskVectorCosSimilarity(BaseAlgorithm, LightningFabricMixin):
         return task_vector
 
     def get_state_dict(self, model: nn.Module):
+        """
+        Extract the state dictionary from a model.
+
+        Args:
+            model (nn.Module): The model to extract parameters from
+
+        Returns:
+            Dict[str, torch.Tensor]: State dictionary containing model parameters.
+                Returns only trainable parameters if trainable_only=True,
+                otherwise returns all parameters.
+        """
         if self.trainable_only:
             return trainable_state_dict(model)
         else: