fusion-bench 0.2.21__py3-none-any.whl → 0.2.23__py3-none-any.whl

fusion_bench/__init__.py

@@ -19,9 +19,32 @@ from . import (
     tasks,
     utils,
 )
+from .constants import RuntimeConstants
 from .method import BaseAlgorithm, BaseModelFusionAlgorithm
 from .mixins import auto_register_config
 from .modelpool import BaseModelPool
-from .models import separate_io
+from .models import (
+    create_default_model_card,
+    load_model_card_template,
+    save_pretrained_with_remote_code,
+    separate_io,
+)
+from .programs import BaseHydraProgram
 from .taskpool import BaseTaskPool
-from .utils import parse_dtype, print_parameters, timeit_context
+from .utils import (
+    BoolStateDictType,
+    LazyStateDict,
+    StateDictType,
+    TorchModelType,
+    cache_with_joblib,
+    get_rankzero_logger,
+    import_object,
+    instantiate,
+    parse_dtype,
+    print_parameters,
+    seed_everything_by_time,
+    set_default_cache_dir,
+    set_print_function_call,
+    set_print_function_call_permeanent,
+    timeit_context,
+)

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/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/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/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/constants/__init__.py

@@ -2,6 +2,7 @@
 import importlib.metadata
 
 from .paths import *
+from .runtime import RuntimeConstants
 
 # fusionbench version
 FUSION_BENCH_VERSION = importlib.metadata.version("fusion-bench")

fusion_bench/constants/runtime.py

@@ -0,0 +1,57 @@
+import threading
+from pathlib import Path
+from typing import Optional, Union
+
+
+class RuntimeConstants:
+    """
+    This class holds constants related to the runtime environment of the Fusion Bench framework.
+    It includes default values for cache directories and other runtime configurations.
+
+    Implemented as a thread-safe singleton to ensure consistent runtime configuration
+    across the entire application.
+    """
+
+    _instance: Optional["RuntimeConstants"] = None
+    _lock = threading.Lock()
+
+    def __new__(cls) -> "RuntimeConstants":
+        """Create a new instance using singleton pattern with thread safety."""
+        with cls._lock:
+            # Double-check locking pattern
+            if cls._instance is None:
+                cls._instance = super(RuntimeConstants, cls).__new__(cls)
+                cls._instance._initialized = False
+            return cls._instance
+
+    def __init__(self):
+        """Initialize the singleton instance only once."""
+        if not self._initialized:
+            # Add your runtime constants here
+            self._initialized = True
+
+    debug = False
+
+    @property
+    def cache_dir(self) -> Path:
+        from fusion_bench.utils.cache_utils import DEFAULT_CACHE_DIR
+
+        return DEFAULT_CACHE_DIR
+
+    @cache_dir.setter
+    def cache_dir(self, path: Union[str, Path]) -> None:
+        from fusion_bench.utils.cache_utils import set_default_cache_dir
+
+        set_default_cache_dir(path)
+
+    @property
+    def print_function_call(self) -> bool:
+        from fusion_bench.utils.instantiate_utils import PRINT_FUNCTION_CALL
+
+        return PRINT_FUNCTION_CALL
+
+    @print_function_call.setter
+    def print_function_call(self, enable: bool) -> None:
+        from fusion_bench.utils.instantiate_utils import set_print_function_call
+
+        set_print_function_call(enable)

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/method/__init__.py

@@ -30,7 +30,7 @@ _import_structure = {
         "TaskArithmeticForLlama",
         "LinearInterpolationAlgorithm",
     ],
-    "slerp": ["SlerpMergeAlgorithm"],
+    "slerp": ["SlerpMergeAlgorithm", "SlerpForCausalLM"],
     "simple_average": ["SimpleAverageAlgorithm"],
     "weighted_average": ["WeightedAverageAlgorithm", "WeightedAverageForLLama"],
     "task_arithmetic": ["TaskArithmeticAlgorithm"],
@@ -71,6 +71,7 @@ _import_structure = {
     ],
     "fw_merging": ["FrankWolfeHardAlgorithm", "FrankWolfeSoftAlgorithm"],
     "tall_mask": ["TallMaskTaskArithmeticAlgorithm"],
+    "model_stock": ["ModelStock"],
     # plug-and-play model merging methods
     "concrete_subspace": [
         "ConcreteTaskArithmeticAlgorithmForCLIP",
@@ -90,7 +91,10 @@ _import_structure = {
         "MixtralForCausalLMMergingAlgorithm",
     ],
     "dawe": ["DataAdaptiveWeightEnsemblingForCLIP"],
-    "we_moe": ["CLIPWeightEnsemblingMoEAlgorithm"],
+    "we_moe": [
+        "CLIPWeightEnsemblingMoEAlgorithm",
+        "FlanT5WeightEnsemblingMoEAlgorithm",
+    ],
     "rankone_moe": ["CLIPRankOneMoEAlgorithm", "RankOneMoEAlgorithm"],
     "sparse_we_moe": [
         "SparseWeightEnsemblingMoEAlgorithm",
@@ -191,6 +195,7 @@ if TYPE_CHECKING:
         MixtralUpscalingAlgorithm,
     )
     from .model_recombination import ModelRecombinationAlgorithm
+    from .model_stock import ModelStock
     from .opcm import OPCMForCLIP
     from .pruning import (
         MagnitudeDiffPruningAlgorithm,
@@ -210,7 +215,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,
@@ -228,7 +233,10 @@ if TYPE_CHECKING:
     from .task_arithmetic import TaskArithmeticAlgorithm
     from .task_singular_vector import TaskSingularVectorMerging
     from .ties_merging import TiesMergingAlgorithm
-    from .we_moe import CLIPWeightEnsemblingMoEAlgorithm
+    from .we_moe import (
+        CLIPWeightEnsemblingMoEAlgorithm,
+        FlanT5WeightEnsemblingMoEAlgorithm,
+    )
     from .weighted_average import WeightedAverageAlgorithm, WeightedAverageForLLama
 
 else:

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: