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

fusion_bench/method/analysis/task_vector_violin_plot.py

@@ -11,30 +11,81 @@ from numpy.typing import NDArray
 from torch import nn
 from tqdm.auto import tqdm
 
-from fusion_bench import BaseAlgorithm, BaseModelPool
-from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
-from fusion_bench.utils import timeit_context
-from fusion_bench.utils.parameters import (
-    StateDictType,
-    state_dict_to_vector,
-    trainable_state_dict,
+from fusion_bench import BaseAlgorithm, BaseModelPool, StateDictType, timeit_context
+from fusion_bench.mixins import (
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+    auto_register_config,
 )
+from fusion_bench.utils import state_dict_to_vector, trainable_state_dict
 from fusion_bench.utils.state_dict_arithmetic import state_dict_sub
 
 log = logging.getLogger(__name__)
 
 
-class TaskVectorViolinPlot(BaseAlgorithm, LightningFabricMixin, SimpleProfilerMixin):
-    R"""
-    Plot violin plots of task vectors as in:
-    [L.Shen, A.Tang, E.Yang et al. Efficient and Effective Weight-Ensembling Mixture of Experts for Multi-Task Model Merging](https://arxiv.org/abs/2410.21804)
+@auto_register_config
+class TaskVectorViolinPlot(
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+    BaseAlgorithm,
+):
+    """
+    Creates violin plots to visualize the distribution of task vector values across models.
+
+    This class implements the task vector visualization technique described in:
+    "Efficient and Effective Weight-Ensembling Mixture of Experts for Multi-Task Model Merging"
+    by L. Shen, A. Tang, E. Yang et al. (https://arxiv.org/abs/2410.21804)
+
+    Task vectors represent the parameter differences between fine-tuned models and their
+    pretrained base model, computed as:
+        task_vector = finetuned_params - pretrained_params
+
+    The algorithm generates two types of violin plots:
+    1. Distribution of raw task vector values (positive and negative)
+    2. Distribution of absolute task vector values
+
+    Args:
+        trainable_only (bool): If True, only consider trainable parameters when computing
+            task vectors. If False, use all parameters.
+        max_points_per_model (int, optional): Maximum number of parameters to sample
+            per model for memory efficiency. If None or 0, uses all parameters.
+            Defaults to 1000.
+        fig_kwargs (dict, optional): Dictionary of keyword arguments to pass to
+            matplotlib.pyplot.subplots. Common options include:
+            - figsize: Tuple of (width, height) in inches
+            - dpi: Dots per inch for resolution
+            - facecolor: Figure background color
+            Defaults to None.
+        output_path (str, optional): Directory to save the violin plots. If None,
+            uses the fabric logger's log directory. Defaults to None.
+
+    Outputs:
+        - task_vector_violin.pdf: Violin plot of raw task vector value distributions
+        - task_vector_violin_abs.pdf: Violin plot of absolute task vector value distributions
+
+    Returns:
+        The pretrained model from the model pool.
+
+    Example:
+        ```python
+        plotter = TaskVectorViolinPlot(
+            trainable_only=True,
+            max_points_per_model=5000,
+            fig_kwargs={'figsize': (12, 8), 'dpi': 300},
+            output_path='./analysis_plots'
+        )
+        pretrained_model = plotter.run(modelpool)
+        ```
+
+    Note:
+        This visualization is particularly useful for understanding:
+        - How different tasks affect model parameters
+        - The magnitude and distribution of parameter changes
+        - Similarities and differences between task adaptations
     """
 
     # config_mapping is a mapping from the attributes to the key in the configuration files
     _config_mapping = BaseAlgorithm._config_mapping | {
-        "trainable_only": "trainable_only",
-        "max_points_per_model": "max_points_per_model",
-        "fig_kwargs": "fig_kwargs",
         "_output_path": "output_path",
     }
 
@@ -46,40 +97,34 @@ class TaskVectorViolinPlot(BaseAlgorithm, LightningFabricMixin, SimpleProfilerMi
         output_path: Optional[str] = None,
         **kwargs,
     ):
-        R"""
-        This class creates violin plots to visualize task vectors, which represent the differences
-        between fine-tuned models and their pretrained base model.
+        """
+        Initialize the TaskVectorViolinPlot analyzer.
 
         Args:
-            trainable_only (bool): If True, only consider trainable parameters when computing
-                task vectors. If False, use all parameters.
-            fig_kwargs (dict, optional): Dictionary of keyword arguments to pass to
-                `matplotlib.pyplot.subplots`. Common options include:
-                - figsize: Tuple of (width, height) in inches
-                - dpi: Dots per inch
-                - facecolor: Figure background color
-                Defaults to None.
-            output_path (str, optional): Path where the violin plot will be saved. If None,
-                uses the fabric logger's log directory. Defaults to None.
-            kwargs: Additional keyword arguments passed to the parent class(es).
-
-        Example:
-
-            ```python
-            plotter = TaskVectorViolinPlot(
-                trainable_only=True,
-                fig_kwargs={'figsize': (10, 6), 'dpi': 300},
-                output_path='./plots'
-            )
-
-            plotter.run(modelpool)
-            ```
+            trainable_only (bool): Whether to consider only trainable parameters when
+                computing task vectors. Set to True to focus on learnable parameters,
+                False to include all parameters including frozen ones.
+            max_points_per_model (int, optional): Maximum number of parameter values
+                to sample per model for visualization. Useful for large models to
+                manage memory usage and plot clarity. Set to None or 0 to use all
+                parameters. Defaults to 1000.
+            fig_kwargs (dict, optional): Keyword arguments passed to matplotlib's
+                subplots function for plot customization. Examples:
+                - {'figsize': (10, 6)} for plot dimensions
+                - {'dpi': 300} for high resolution
+                - {'facecolor': 'white'} for background color
+                Defaults to None (uses matplotlib defaults).
+            output_path (str, optional): Directory path where violin plots will be saved.
+                If None, uses the fabric logger's log directory. The directory will be
+                created if it doesn't exist. Defaults to None.
+            **kwargs: Additional keyword arguments passed to parent classes.
+
+        Note:
+            The parameter name 'fig_kwawrgs' appears to be a typo for 'fig_kwargs'.
+            This should be corrected in the parameter name for consistency.
         """
-        self.trainable_only = trainable_only
-        self.fig_kwargs = fig_kwawrgs
-        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):
@@ -89,20 +134,39 @@ class TaskVectorViolinPlot(BaseAlgorithm, LightningFabricMixin, SimpleProfilerMi
             return self._output_path
 
     def run(self, modelpool: BaseModelPool):
-        """Create violin plots of task vectors comparing different fine-tuned models against a pretrained model.
+        """
+        Execute the task vector violin plot analysis and visualization.
 
-        This method implements the visualization technique from the paper "Efficient and Effective
-        Weight-Ensembling Mixture of Experts for Multi-Task Model Merging". It:
+        This method implements the core algorithm that:
+        1. Loads the pretrained base model from the model pool
+        2. Computes task vectors for each fine-tuned model (parameter differences)
+        3. Creates two violin plots showing the distribution of task vector values:
+           - Raw values plot: Shows positive and negative parameter changes
+           - Absolute values plot: Shows magnitude of parameter changes
+        4. Saves both plots as PDF files in the output directory
 
-        1. Loads the pretrained model
-        2. Computes task vectors (differences between fine-tuned and pretrained models)
-        3. Creates violin plots showing the distribution of values in these task vectors
+        The visualization technique follows the approach described in:
+        "Efficient and Effective Weight-Ensembling Mixture of Experts for Multi-Task Model Merging"
 
         Args:
-            modelpool (BaseModelPool): Model pool containing the pretrained model and fine-tuned models
+            modelpool (BaseModelPool): Pool containing both a pretrained model and
+                fine-tuned models. Must have `has_pretrained=True`.
 
         Returns:
-            pretrained_model (nn.Model): The plot is saved to the specified output path.
+            nn.Module: The pretrained model loaded from the model pool.
+
+        Raises:
+            AssertionError: If the model pool doesn't contain a pretrained model.
+
+        Side Effects:
+            - Creates output directory if it doesn't exist
+            - Saves 'task_vector_violin.pdf' (raw values distribution)
+            - Saves 'task_vector_violin_abs.pdf' (absolute values distribution)
+            - Prints progress information during task vector computation
+
+        Example Output Files:
+            - task_vector_violin.pdf: Shows how parameters change (+ and -)
+            - task_vector_violin_abs.pdf: Shows magnitude of parameter changes
         """
         assert modelpool.has_pretrained
         pretrained_model = modelpool.load_pretrained_model()
@@ -175,6 +239,34 @@ class TaskVectorViolinPlot(BaseAlgorithm, LightningFabricMixin, SimpleProfilerMi
         return pretrained_model
 
     def get_task_vector(self, pretrained_model, finetuned_model):
+        """
+        Compute the task vector representing parameter changes from pretraining to fine-tuning.
+
+        The task vector quantifies how model parameters have changed during task-specific
+        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:
+            np.ndarray: Flattened numpy array containing parameter differences.
+                If max_points_per_model is set, the array may be randomly downsampled
+                for memory efficiency and visualization clarity.
+
+        Processing Steps:
+            1. Extract state dictionaries from both models
+            2. Compute parameter differences (subtraction)
+            3. Flatten to 1D vector
+            4. Convert to numpy array with float32 precision
+            5. Optionally downsample if max_points_per_model is specified
+
+        Note:
+            - Uses only trainable parameters if trainable_only=True
+            - Downsampling uses random sampling without replacement
+            - Preserves the relative distribution of parameter changes
+        """
         task_vector = state_dict_sub(
             self.get_state_dict(finetuned_model),
             self.get_state_dict(pretrained_model),
@@ -199,6 +291,22 @@ class TaskVectorViolinPlot(BaseAlgorithm, LightningFabricMixin, SimpleProfilerMi
         return task_vector
 
     def get_state_dict(self, model: nn.Module):
+        """
+        Extract the state dictionary from a model based on parameter filtering settings.
+
+        Args:
+            model (nn.Module): The PyTorch model to extract parameters from
+
+        Returns:
+            Dict[str, torch.Tensor]: State dictionary containing model parameters.
+                If trainable_only=True, returns only parameters with requires_grad=True.
+                If trainable_only=False, returns all parameters including frozen ones.
+
+        Note:
+            This method respects the trainable_only configuration to focus analysis
+            on either learnable parameters or the complete parameter set depending
+            on the research question being addressed.
+        """
         if self.trainable_only:
             return trainable_state_dict(model)
         else:

fusion_bench/method/bitdelta/__init__.py

@@ -1,4 +1,5 @@
 """
 Adapted from https://github.com/FasterDecoding/BitDelta
 """
+
 from .bitdelta import BitDeltaAlgorithm

fusion_bench/method/bitdelta/bitdelta.py

@@ -6,7 +6,11 @@ import torch.nn.functional as F
 from tqdm.auto import tqdm
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
-from fusion_bench.mixins import LightningFabricMixin, SimpleProfilerMixin
+from fusion_bench.mixins import (
+    LightningFabricMixin,
+    SimpleProfilerMixin,
+    auto_register_config,
+)
 from fusion_bench.modelpool import CausalLMPool
 
 from .bitdelta_utils.data import get_dataloader, get_dataset
@@ -15,23 +19,12 @@ from .bitdelta_utils.diff import compress_diff, save_diff, save_full_model
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class BitDeltaAlgorithm(
-    BaseAlgorithm,
     LightningFabricMixin,
     SimpleProfilerMixin,
+    BaseAlgorithm,
 ):
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "save_dir": "save_dir",
-        "save_full_model": "save_full_model",
-        "lr": "lr",
-        "batch_size": "batch_size",
-        "num_steps": "num_steps",
-        "dataset_name": "dataset_name",
-        "subset": "subset",
-        "split": "split",
-        "max_length": "max_length",
-    }
-
     def __init__(
         self,
         save_dir: str,
@@ -46,15 +39,6 @@ class BitDeltaAlgorithm(
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.save_dir = save_dir
-        self.save_full_model = save_full_model
-        self.lr = lr
-        self.batch_size = batch_size
-        self.num_steps = num_steps
-        self.dataset_name = dataset_name
-        self.subset = subset
-        self.split = split
-        self.max_length = max_length
 
     def run(self, modelpool: CausalLMPool):
         if self.save_dir is None:

fusion_bench/method/classification/clip_finetune.py

@@ -393,7 +393,7 @@ def convert_l_lora_state_dict_to_hf(
     base_model_name: Optional[str] = None,
 ):
     """
-    Convert a linearized Lora model's checkpoint to Hugggingface's format.
+    Convert a linearized Lora model's checkpoint to huggingface's format.
 
     Args:
         pretrained_path (str): The path to the pretrained model.

fusion_bench/method/expert_sparsity/mixtral/dynamic_skipping.py

@@ -23,6 +23,7 @@ from transformers import MixtralForCausalLM
 from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM
 
 import fusion_bench as fb
+from fusion_bench import auto_register_config
 from fusion_bench.method.expert_sparsity.utils.calibration_data import (
     build_calib_loader,
 )
@@ -97,6 +98,7 @@ def dynamic_skipping(
     return model, (res_median, res_mean)
 
 
+@auto_register_config
 class DynamicSkippingPruningForMixtral(
     fb.BaseAlgorithm,
     fb.mixins.LightningFabricMixin,

fusion_bench/method/expert_sparsity/mixtral/layer_wise_pruning.py

@@ -22,6 +22,7 @@ from transformers import MixtralForCausalLM
 from transformers.models.mixtral.modeling_mixtral import MixtralDecoderLayer
 
 import fusion_bench as fb
+from fusion_bench import auto_register_config
 from fusion_bench.method.expert_sparsity.utils.calibration_data import (
     build_calib_loader,
 )
@@ -81,6 +82,7 @@ def layerwise_pruning(
     return model, (global_loss_history,)
 
 
+@auto_register_config
 class LayerWisePruningForMixtral(
     fb.BaseAlgorithm,
     fb.mixins.LightningFabricMixin,

fusion_bench/method/expert_sparsity/mixtral/progressive_pruning.py

@@ -20,6 +20,7 @@ from tqdm import tqdm
 from transformers import MixtralForCausalLM
 
 import fusion_bench as fb
+from fusion_bench import auto_register_config
 from fusion_bench.method.expert_sparsity.utils.calibration_data import (
     build_calib_loader,
 )
@@ -95,6 +96,7 @@ def progressive_pruning(
     return model, (global_loss_history,)
 
 
+@auto_register_config
 class ProgressivePruningForMixtral(
     fb.BaseAlgorithm,
     fb.mixins.LightningFabricMixin,

fusion_bench/method/fisher_merging/clip_fisher_merging.py

@@ -32,7 +32,6 @@ class FisherMergingForCLIPVisionModel(
     zeroshot_weights = {}
 
     _config_mapping = FisherMergingAlgorithm._config_mapping | {
-        "zeroshot_weights_cache_dir": "zeroshot_weights_cache_dir",
         "_dataloader_kwargs": "dataloader_kwargs",
     }
 
@@ -44,7 +43,6 @@ class FisherMergingForCLIPVisionModel(
         minimal_fisher_weight,
         num_fisher_examples,
         dataloader_kwargs: DictConfig,
-        zeroshot_weights_cache_dir=None,
         **kwargs,
     ):
         """
@@ -56,7 +54,6 @@ class FisherMergingForCLIPVisionModel(
             minimal_fisher_weight (float): Minimal value for Fisher weights to avoid numerical issues.
             num_fisher_examples (int): Number of examples to compute Fisher weights.
             dataloader_kwargs (DictConfig): Configuration for the dataloader.
-            zeroshot_weights_cache_dir (str, optional): Directory to cache zero-shot weights. Defaults to None.
             **kwargs: Additional keyword arguments.
         """
         super().__init__(
@@ -66,7 +63,6 @@ class FisherMergingForCLIPVisionModel(
             num_fisher_examples=num_fisher_examples,
         )
         self.dataloader_kwargs = dataloader_kwargs
-        self.zeroshot_weights_cache_dir = zeroshot_weights_cache_dir
         for key, value in kwargs.items():
             log.warning(f"Unused argument: {key}={value}")
             setattr(self, key, value)

fusion_bench/method/fisher_merging/gpt2_fisher_merging.py

@@ -15,10 +15,10 @@ from transformers import GPT2ForSequenceClassification, GPT2Model
 from transformers.data import default_data_collator
 from transformers.models.gpt2.modeling_gpt2 import Conv1D
 
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import LightningFabricMixin, auto_register_config
 from fusion_bench.modelpool import GPT2ForSequenceClassificationPool
 from fusion_bench.utils import timeit_context
-from fusion_bench.mixins import auto_register_config
+
 from .fisher_merging import FisherMergingAlgorithm, get_param_squared_gradients
 
 

fusion_bench/method/linear/simple_average_for_llama.py

@@ -1,3 +1,4 @@
+import os
 from copy import deepcopy
 from typing import TYPE_CHECKING, Optional
 
@@ -7,13 +8,16 @@ from typing_extensions import override
 from fusion_bench import timeit_context
 from fusion_bench.method.base_algorithm import BaseAlgorithm
 from fusion_bench.method.simple_average import SimpleAverageAlgorithm
+from fusion_bench.mixins import auto_register_config
 from fusion_bench.modelpool import CausalLMBackbonePool, CausalLMPool
+from fusion_bench.models.hf_utils import create_default_model_card
 from fusion_bench.utils import instantiate
-from fusion_bench.utils.pylogger import getRankZeroLogger
+from fusion_bench.utils.pylogger import get_rankzero_logger
 
-log = getRankZeroLogger(__name__)
+log = get_rankzero_logger(__name__)
 
 
+@auto_register_config
 class SimpleAverageForLlama(BaseAlgorithm):
     R"""
     A simple averaging algorithm for LLama models. If `merge_backbone` is set to `True`, the backbone of the model will be averaged and the rest of the model will be loaded from the pre-trained model.
@@ -29,21 +33,14 @@ class SimpleAverageForLlama(BaseAlgorithm):
         ```
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "merge_backbone": "merge_backbone",
-        "show_pbar": "show_pbar",
-    }
-
     def __init__(
         self,
         merge_backbone: bool,
         model_save_path: Optional[str] = None,
         show_pbar: bool = False,
+        **kwargs,
     ):
-        super().__init__()
-        self.merge_backbone = merge_backbone
-        self.model_save_path = model_save_path
-        self.show_pbar = show_pbar
+        super().__init__(**kwargs)
 
     @override
     def run(self, modelpool: CausalLMPool):
@@ -75,4 +72,12 @@ class SimpleAverageForLlama(BaseAlgorithm):
             with timeit_context(f"Saving the model to {self.model_save_path}"):
                 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 simple averaging.",
+                    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/model_stock/__init__.py

@@ -0,0 +1 @@
+from .model_stock import ModelStock