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

fusion_bench/method/slerp/slerp.py

@@ -1,16 +1,24 @@
 import logging
-from typing import Any, Dict
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 import torch
 from torch import nn
+from tqdm import tqdm
 from typing_extensions import override
 
+from fusion_bench import LazyStateDict, create_default_model_card, timeit_context
 from fusion_bench.method import BaseAlgorithm
-from fusion_bench.modelpool import BaseModelPool
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
+from fusion_bench.modelpool import BaseModelPool, CausalLMPool
 from fusion_bench.utils.type import StateDictType
 
 from .slerp_utils import slerp
 
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
 log = logging.getLogger(__name__)
 
 
@@ -21,6 +29,7 @@ def slerp_on_state_dicts(
     *,
     DOT_THRESHOLD: float = 0.9995,
     epsilon: float = 1e-8,
+    show_pbar: bool = False,
 ) -> StateDictType:
     """
     Perform spherical linear interpolation (slerp) on the state dictionaries of two models.
@@ -36,7 +45,8 @@ def slerp_on_state_dicts(
         dict: The interpolated state dictionary.
     """
     state_dict = {}
-    for key in secondary_state_dict:
+    pbar = secondary_state_dict if not show_pbar else tqdm(secondary_state_dict)
+    for key in pbar:
         v0 = primary_state_dict[key]
         v1 = secondary_state_dict[key]
         if v0.shape != v1.shape:
@@ -49,18 +59,19 @@ def slerp_on_state_dicts(
     return state_dict
 
 
+@auto_register_config
 class SlerpMergeAlgorithm(BaseAlgorithm):
     """
     General purpose implementation of Slerp (Spherical Linear Interpolation) for PyTorch models.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "t": "t",
-        "DOT_THRESHOLD": "DOT_THRESHOLD",
-        "epsilon": "epsilon",
-    }
-
-    def __init__(self, t: float, DOT_THRESHOLD: float = 0.9995, epsilon: float = 1e-8):
+    def __init__(
+        self,
+        t: float,
+        DOT_THRESHOLD: float = 0.9995,
+        epsilon: float = 1e-8,
+        **kwargs,
+    ):
         """
         Initialize the SlerpMergeAlgorithm.
 
@@ -69,10 +80,7 @@ class SlerpMergeAlgorithm(BaseAlgorithm):
             DOT_THRESHOLD (float, optional): The threshold for the dot product of the two vectors. Defaults to 0.9995.
             epsilon (float, optional): The epsilon value for numerical stability. Defaults to 1e-8.
         """
-        self.t = t
-        self.DOT_THRESHOLD = DOT_THRESHOLD
-        self.epsilon = epsilon
-        super().__init__()
+        super().__init__(**kwargs)
 
     @override
     def run(self, modelpool: BaseModelPool) -> nn.Module:
@@ -102,3 +110,91 @@ class SlerpMergeAlgorithm(BaseAlgorithm):
 
         primary_model.load_state_dict(state_dict)
         return primary_model
+
+
+@auto_register_config
+class SlerpForCausalLM(
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
+    """
+    Slerp (Spherical Linear Interpolation) for Causal Language Models.
+    """
+
+    def __init__(
+        self,
+        t: float,
+        DOT_THRESHOLD: float = 0.9995,
+        epsilon: float = 1e-8,
+        model_save_path: Optional[str] = None,
+        show_pbar: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize the SlerpForCausalLM algorithm.
+
+        Args:
+            t (float): The interpolation parameter. Must be in the range [0, 1].
+                      t=0 returns the first model, t=1 returns the second model,
+                      t=0.5 provides balanced interpolation.
+            DOT_THRESHOLD (float, optional): The threshold for the dot product of normalized vectors.
+                                           When the absolute dot product exceeds this threshold,
+                                           vectors are considered nearly collinear and linear
+                                           interpolation (LERP) is used instead of SLERP for
+                                           numerical stability. Defaults to 0.9995.
+            epsilon (float, optional): Small value used for numerical stability to avoid
+                                     division by zero during vector normalization.
+                                     Defaults to 1e-8.
+            model_save_path (Optional[str], optional): Path where the merged model should be saved.
+                                                     If None, the model is not saved to disk.
+                                                     Defaults to None.
+            show_pbar (bool, optional): Whether to display a progress bar during the interpolation
+                                      process. Useful for debugging or monitoring progress with
+                                      large models. Defaults to False.
+            **kwargs: Additional keyword arguments passed to the parent BaseAlgorithm class.
+        """
+        super().__init__(**kwargs)
+
+    @override
+    def run(self, modelpool: CausalLMPool):
+        assert len(modelpool.all_model_names) == 2, "Slerp expect exactly 2 models"
+        primary_model = modelpool.load_model(modelpool.all_model_names[0])
+        secondary_model = modelpool.load_model(modelpool.all_model_names[1])
+
+        with torch.no_grad():
+            primary_state_dict = primary_model.state_dict()
+            secondary_state_dict = secondary_model.state_dict()
+            state_dict = slerp_on_state_dicts(
+                self.t,
+                primary_state_dict,
+                secondary_state_dict,
+                DOT_THRESHOLD=self.DOT_THRESHOLD,
+                epsilon=self.epsilon,
+            )
+
+        if isinstance(primary_model, nn.Module):
+            model = primary_model
+            model.load_state_dict(state_dict)
+        elif isinstance(primary_model, LazyStateDict):
+            model: "PreTrainedModel" = deepcopy(primary_model.meta_module)
+            model.to(device=primary_model._device)
+            model.load_state_dict(state_dict)
+        else:
+            raise TypeError(
+                f"Unsupported model type: {type(primary_model)}. "
+                "Expected nn.Module or LazyStateDict."
+            )
+        if self.model_save_path is not None:
+            with timeit_context(f"Saving the model to {self.model_save_path}"):
+                tokenizer = modelpool.load_tokenizer()
+                tokenizer.save_pretrained(self.model_save_path)
+                model.save_pretrained(self.model_save_path)
+                model_card_str = create_default_model_card(
+                    models=[modelpool.get_model_path(m) for m in modelpool.model_names],
+                    description="Merged model using Slerp.",
+                    algorithm_config=self.config,
+                    modelpool_config=modelpool.config,
+                )
+                with open(os.path.join(self.model_save_path, "README.md"), "w") as f:
+                    f.write(model_card_str)
+        return model

fusion_bench/method/we_moe/flan_t5_we_moe.py

@@ -16,11 +16,14 @@ from transformers.data import default_data_collator
 
 from fusion_bench.method import BaseAlgorithm
 from fusion_bench.method.task_arithmetic.task_arithmetic import task_arithmetic_merge
-from fusion_bench.mixins.lightning_fabric import LightningFabricMixin
-from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
+from fusion_bench.mixins import (
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+    auto_register_config,
+)
 from fusion_bench.modelpool import Seq2SeqLMPool
 from fusion_bench.models.we_moe import WeightEnsemblingMoE
-from fusion_bench.utils import timeit_context
+from fusion_bench.utils import print_parameters, timeit_context
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
 from fusion_bench.utils.instantiate_utils import instantiate
 from fusion_bench.utils.parameters import print_parameters
@@ -31,10 +34,11 @@ from .utils import get_memory_usage
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class FlanT5WeightEnsemblingMoEAlgorithm(
-    BaseAlgorithm,
     LightningFabricMixin,
     SimpleProfilerMixin,
+    BaseAlgorithm,
 ):
     """
     FlanT5WeightEnsemblingMoEAlgorithm is a class that implements the WeightEnsemblingMoEAlgorithm
@@ -60,7 +64,6 @@ class FlanT5WeightEnsemblingMoEAlgorithm(
         num_workers: int = 0,
         max_steps: int = 1000,
         use_grad_accumulate: bool = True,
-        cache_dir: bool = "outputs",
         fast_dev_run: bool = False,
         **kwargs,
     ):
@@ -70,23 +73,9 @@ class FlanT5WeightEnsemblingMoEAlgorithm(
         Args:
             algorithm_config (DictConfig): The configuration for the algorithm.
         """
-        self.checkpoint = checkpoint
-        self.save_checkpoint = save_checkpoint
-        self.router_hidden_layers = router_hidden_layers
-        self.init_lambda = init_lambda
-        self.batch_reduce = batch_reduce
-        self.lr = lr
-        self.optimizer = optimizer
-        self.devices = devices
-        self.batch_size = batch_size
-        self.num_workers = num_workers
-        self.max_steps = max_steps
-        self.use_grad_accumulate = use_grad_accumulate
-        self.cache_dir = cache_dir
-        self.fast_dev_run = fast_dev_run
         super().__init__(**kwargs)
 
-    def construct_moe_model(self) -> WeightEnsemblingMoE:
+    def construct_moe_model(self):
         """
         Construct the Mixture of Experts (MoE) model using the models in the model pool.
 

fusion_bench/mixins/clip_classification.py

@@ -113,11 +113,27 @@ class CLIPClassificationMixin(LightningFabricMixin):
         clip_model: Optional[CLIPModel] = None,
         task_names: Optional[List[str]] = None,
     ):
+        """
+        Initializes a zero-shot classification head.
+
+        This method constructs a zero-shot classification head by generating text embeddings for each class name using a set of templates.
+        These embeddings function as the weights of the classification layer. The method also extracts the `visual_projection` and `logit_scale`
+        from the provided CLIP model, which are necessary for calculating the final logits.
+
+        Args:
+            clip_processor (Optional[CLIPProcessor]): The processor for the CLIP model. If not provided, it is loaded from the model pool.
+            clip_model (Optional[CLIPModel]): The CLIP model to use. If not provided, a pretrained model is loaded from the model pool.
+            task_names (Optional[List[str]]): A list of task names to set up the classification head for. If not provided, all models in the model pool will be used.
+        """
         self.whether_setup_zero_shot_classification_head = True
+        # load clip model if not provided
         if clip_model is None:
             if self.modelpool.has_pretrained:
                 clip_model = self.modelpool.load_clip_model("_pretrained_")
             else:
+                log.warning(
+                    f"No pretrained CLIP model found, using the model from the model pool: {self.modelpool.model_names[0]}."
+                )
                 clip_model = self.modelpool.load_clip_model(
                     self.modelpool.model_names[0]
                 )
@@ -166,16 +182,20 @@ class CLIPClassificationMixin(LightningFabricMixin):
         image_embeds: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """
-        Compute the logits of the images for a given task.
+        Computes the classification logits for a batch of images for a specific task.
+
+        This method performs zero-shot classification by calculating the cosine similarity between image and text embeddings.
+        The image embeddings are obtained from the provided vision model, and the text embeddings (zero-shot weights) are pre-computed for the task.
+        The similarity scores are then scaled by the CLIP model's `logit_scale` to produce the final logits.
 
         Args:
-            module (Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"]): The module to compute the logits.
-            images (torch.Tensor): The images to compute the logits.
-            task (str): The task to compute the logits.
-            image_embeds (Optional[torch.Tensor]): The precomputed image embeddings. If None, the image embeddings will be computed.
+            module (Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"]): The vision encoder part of the CLIP model.
+            images (torch.Tensor): A batch of images to classify.
+            task (str): The name of the classification task.
+            image_embeds (Optional[torch.Tensor]): Pre-computed image embeddings. If provided, the method skips the image encoding step.
 
         Returns:
-            torch.Tensor: The logits of the images.
+            torch.Tensor: A tensor of logits for each image, with shape (batch_size, num_classes).
         """
         text_embeds = self.zeroshot_weights[task]
 

fusion_bench/mixins/serialization.py

@@ -4,7 +4,7 @@ from copy import deepcopy
 from functools import wraps
 from inspect import Parameter, _ParameterKind
 from pathlib import Path
-from typing import Dict, Optional, Union
+from typing import Dict, Mapping, Optional, Union
 
 from omegaconf import DictConfig, OmegaConf
 
@@ -21,6 +21,20 @@ __all__ = [
 ]
 
 
+def _get_attr_name(config_mapping: Mapping[str, str], param_name):
+    for attr_name, p in config_mapping.items():
+        if p == param_name:
+            return attr_name
+    else:
+        raise ValueError(f"Parameter {param_name} not found in config mapping.")
+
+
+def _set_attr(self, param_name: str, value):
+    attr_name = _get_attr_name(self._config_mapping, param_name)
+    log.debug(f"set {attr_name} to {value}. Parameter name: {param_name}")
+    setattr(self, attr_name, value)
+
+
 def auto_register_config(cls):
     """
     Decorator to automatically register __init__ parameters in _config_mapping.
@@ -56,8 +70,8 @@ def auto_register_config(cls):
         ```python
         @auto_register_config
         class MyAlgorithm(BaseYAMLSerializable):
-            def __init__(self, learning_rate: float = 0.001, batch_size: int = 32, model_name: str = "default"):
-                super().__init__()
+            def __init__(self, learning_rate: float = 0.001, batch_size: int = 32, model_name: str = "default", **kwargs):
+                super().__init__(**kwargs)
 
         # All instantiation methods work automatically:
         algo1 = MyAlgorithm(0.01, 64)  # positional args
@@ -103,8 +117,7 @@ def auto_register_config(cls):
             cls._config_mapping[param_name] = param_name
 
     def __init__(self, *args, **kwargs):
-        nonlocal original_init, registered_parameters
-
+        log.debug(f"set attributes for {self.__class__.__name__} in {cls.__name__}")
         # auto-register the attributes based on the signature
         sig = inspect.signature(original_init)
         param_names = list(sig.parameters.keys())[1:]  # Skip 'self'
@@ -117,29 +130,26 @@ def auto_register_config(cls):
                     _ParameterKind.VAR_POSITIONAL,
                     _ParameterKind.VAR_KEYWORD,
                 ]:
-                    setattr(self, param_name, arg_value)
+                    _set_attr(self, param_name, arg_value)
 
         # Handle keyword arguments and defaults
         for param_name in param_names:
-            if (
-                sig.parameters[param_name].kind
-                not in [
-                    _ParameterKind.VAR_POSITIONAL,
-                    _ParameterKind.VAR_KEYWORD,
-                ]
-            ) and (param_name not in registered_parameters):
+            if sig.parameters[param_name].kind not in [
+                _ParameterKind.VAR_POSITIONAL,
+                _ParameterKind.VAR_KEYWORD,
+            ]:
                 # Skip if already set by positional argument
                 param_index = param_names.index(param_name)
                 if param_index >= 0 and param_index < len(args):
                     continue
 
                 if param_name in kwargs:
-                    setattr(self, param_name, kwargs[param_name])
+                    _set_attr(self, param_name, kwargs[param_name])
                 else:
                     # Set default value if available and attribute doesn't exist
                     default_value = sig.parameters[param_name].default
                     if default_value is not Parameter.empty:
-                        setattr(self, param_name, default_value)
+                        _set_attr(self, param_name, default_value)
 
         # Call the original __init__
         result = original_init(self, *args, **kwargs)

fusion_bench/modelpool/base_pool.py

@@ -277,7 +277,7 @@ class BaseModelPool(
         for dataset_name in self.test_dataset_names:
             yield self.load_test_dataset(dataset_name)
 
-    def save_model(self, model: nn.Module, path: str):
+    def save_model(self, model: nn.Module, path: str, *args, **kwargs):
         """
         Save the state dictionary of the model to the specified path.