fusion-bench 0.2.16__tar.gz → 0.2.17__tar.gz

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.16
+Version: 0.2.17
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
 License: MIT License
@@ -171,6 +171,8 @@ It can be used to improve the performance and robustness of model or to combine
 For a more detailed introduction to deep model fusion, you can refer to [W. Li, 2023, 'Deep Model Fusion: A Survey'](https://arxiv.org/abs/2309.15698). We also provide a brief overview of deep model fusion in [our documentation](https://tanganke.github.io/fusion_bench/).
 In this benchmark, we evaluate the performance of different fusion methods on a variety of datasets and tasks.
 
+A comprehensive list of papers about model merging can be found at [this repository](https://github.com/EnnengYang/Awesome-Model-Merging-Methods-Theories-Applications), and [the arXiv paper](https://arxiv.org/abs/2408.07666) is also available.
+
 ## Project Structure
 
 The project is structured as follows:

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/README.md

@@ -120,6 +120,8 @@ It can be used to improve the performance and robustness of model or to combine
 For a more detailed introduction to deep model fusion, you can refer to [W. Li, 2023, 'Deep Model Fusion: A Survey'](https://arxiv.org/abs/2309.15698). We also provide a brief overview of deep model fusion in [our documentation](https://tanganke.github.io/fusion_bench/).
 In this benchmark, we evaluate the performance of different fusion methods on a variety of datasets and tasks.
 
+A comprehensive list of papers about model merging can be found at [this repository](https://github.com/EnnengYang/Awesome-Model-Merging-Methods-Theories-Applications), and [the arXiv paper](https://arxiv.org/abs/2408.07666) is also available.
+
 ## Project Structure
 
 The project is structured as follows:

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/adamerging/flan_t5_layer_wise_adamerging.py

@@ -29,7 +29,7 @@ from fusion_bench.models.wrappers.layer_wise_fusion import (
     get_layer_wise_weights,
 )
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .entropy_loss import entropy_loss
 from .min_norm_solvers import MinNormSolver

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/adamerging/gpt2_layer_wise_adamerging.py

@@ -29,7 +29,7 @@ from fusion_bench.models.wrappers.layer_wise_fusion import (
     get_layer_wise_weights,
 )
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .entropy_loss import entropy_loss
 from .min_norm_solvers import MinNormSolver

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/base_algorithm.py

@@ -19,6 +19,7 @@ class BaseAlgorithm(BaseYAMLSerializableModel):
     """
 
     _program = None
+    _config_key = "method"
 
     @abstractmethod
     def run(self, modelpool: BaseModelPool):

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/dawe/dawe_for_clip.py

@@ -23,7 +23,7 @@ from fusion_bench.mixins import CLIPClassificationMixin
 from fusion_bench.modelpool import CLIPVisionModelPool
 from fusion_bench.utils import timeit_context
 from fusion_bench.utils.data import InfiniteDataLoader
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .warppers.dawe_model import DataAdaptiveWeightEnsemblingCLIPVisionModel
 

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/depth_upscaling/depth_upscaling_for_llama.py

@@ -1,9 +1,10 @@
 import os
 from typing import Optional
 
+from transformers import PreTrainedModel
 from typing_extensions import override
 
-from fusion_bench.modelpool.causal_lm.causal_lm import CausalLM, CausalLMPool
+from fusion_bench.modelpool.causal_lm.causal_lm import CausalLMPool
 from fusion_bench.utils import timeit_context
 
 from .depth_upscaling import DepthUpscalingAlgorithm
@@ -46,7 +47,7 @@ class DepthUpscalingForLlama(DepthUpscalingAlgorithm):
         if self.model_save_path is not None:
             tokenizer = modelpool.load_tokenizer()
 
-        model: CausalLM = modelpool.load_pretrained_or_first_model()
+        model: PreTrainedModel = modelpool.load_pretrained_or_first_model()
         model.model.layers = super().run(model.model.layers)
         model.config.num_hidden_layers = len(model.model.layers)
 

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/gossip/flan_t5_layer_wise_gossip.py

@@ -32,7 +32,7 @@ from fusion_bench.models.wrappers.layer_wise_fusion import (
     get_layer_wise_weights,
 )
 from fusion_bench.utils.data import InfiniteDataLoader, load_tensor_from_file
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 from .entropy_loss import entropy_loss
 from .layer_wise_gossip import ModelScheduler

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/pwe_moe/module.py

@@ -13,14 +13,9 @@ import torch.func
 from torch import Tensor, nn
 from torch.nn import functional as F
 
-log = logging.getLogger(__name__)
-
+from fusion_bench.utils import join_list
 
-def join_list(list_of_list: List[List]):
-    ans = []
-    for item in list_of_list:
-        ans.extend(item)
-    return ans
+log = logging.getLogger(__name__)
 
 
 class PWEMoEGate(nn.Module):

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/simple_average.py

@@ -39,12 +39,13 @@ def simple_average(
         >>> import torch.nn as nn
         >>> model1 = nn.Linear(10, 10)
         >>> model2 = nn.Linear(10, 10)
-        >>> averaged_model = simple_averageing([model1, model2])
+        >>> averaged_model = simple_average([model1, model2])
 
         >>> state_dict1 = model1.state_dict()
         >>> state_dict2 = model2.state_dict()
-        >>> averaged_state_dict = simple_averageing([state_dict1, state_dict2])
+        >>> averaged_state_dict = simple_average([state_dict1, state_dict2])
     """
+    assert len(modules) > 0, "modules must be a non-empty list"
     if isinstance(modules[0], nn.Module):
         if base_module is None:
             new_module = deepcopy(modules[0])

fusion_bench-0.2.17/fusion_bench/method/task_singular_vector/TSVM.py

@@ -0,0 +1,296 @@
+R"""
+# Task Singular Vector Merging (TSVM) Algorithm Implementation
+
+This module implements the Task Singular Vector Merging algorithm for combining multiple fine-tuned models
+into a single unified model.
+
+## Example Usage:
+
+Merge 8 CLIP-ViT-B/32 models with TSVM:
+
+```bash
+fusion_bench \
+    method=task_singular_vector/TaskSingularVectorMerging \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8_model_only \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+```
+
+Merge 8 CLIP-ViT-B/32 models with TSVM and return individual transformed models:
+
+```bash
+fusion_bench \
+    method=task_singular_vector/TaskSingularVectorMerging \
+    method.return_single_task_models=true \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8_model_only \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+```
+
+Merge 20 CLIP-VIT-B/32 models with TSVM:
+
+```bash
+fusion_bench \
+    method=task_singular_vector/TaskSingularVectorMerging \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TALL20_model_only \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TALL20
+```
+
+## References:
+
+- Gargiulo, et al. Task Singular Vectors: Reducing Task Interference in Model Merging. 
+    https://arxiv.org/abs/2412.00081
+- See `docs/algorithms/task_singular_vector.md` for more details.
+"""
+
+from copy import deepcopy
+from typing import Iterable, List, Optional, Union
+
+import torch
+from omegaconf import ListConfig
+from torch import Tensor, nn
+
+import fusion_bench as fb
+from fusion_bench import BaseAlgorithm
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.utils import timeit_context
+from fusion_bench.utils.state_dict_arithmetic import (
+    state_dict_add,
+    state_dict_mul,
+    state_dict_sub,
+)
+from fusion_bench.utils.type import StateDictType
+
+from .utils import (
+    TSVM_utils,
+    check_parameterNamesMatch,
+    check_state_dicts_equal,
+    state_dict_to_vector,
+    vector_to_state_dict,
+)
+
+
+class TaskSingularVectorMerging(BaseAlgorithm, LightningFabricMixin):
+    """
+    Task Singular Vector Merging (TSVM) Algorithm
+
+    This class implements a model merging technique that leverages Singular Value
+    Decomposition (SVD) to identify and combine the most important directions in the task vector
+    space. The algorithm is particularly effective for merging multiple models fine-tuned on
+    different tasks while preserving their essential capabilities.
+
+    Key Concepts:
+    - Task Vector: The difference between a fine-tuned model and its pretrained base model,
+      representing the knowledge gained during fine-tuning for a specific task.
+    - Singular Value Decomposition: A matrix factorization technique used to find the principal
+      components (most important directions) in the space of task vectors.
+    - Model Merging: The process of combining multiple models into a single unified model that
+      retains capabilities from all constituent models.
+
+    Algorithm Steps:
+    1. Extract task vectors from all fine-tuned models by subtracting the pretrained model
+    2. Apply SVD to the matrix of task vectors to find principal components
+    3. Reconstruct task vectors using only the most significant singular vectors
+    4. Merge the reconstructed task vectors (either individually scaled or as a sum)
+    5. Add the final merged task vector to the pretrained model to create the unified model
+
+    see `docs/algorithms/task_singular_vector.md` for comprehensive algorithmic details.
+    """
+
+    def __init__(
+        self,
+        alpha: Optional[Union[float, Iterable[float]]] = None,
+        exclude_keys: Optional[List[str]] = None,
+        return_single_task_models: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize the Task Singular Vector Merging algorithm.
+
+        Args:
+            alpha (Union[float, Iterable[float]], optional): Scaling factor(s) for task vectors.
+                This parameter controls the strength of the task-specific adaptations in the final model.
+
+                - If a single float: Applied to the final merged task vector after SVD reconstruction.
+                  This uniformly scales the entire merged adaptation.
+
+                - If an iterable of floats: Applied to individual task vectors before SVD and merging.
+                  Must have the same length as the number of models in the modelpool.
+                  This allows for task-specific weighting (e.g., giving more importance to certain tasks).
+
+                - If None: No scaling is applied (equivalent to alpha=1.0).
+
+                Example: alpha=[0.8, 1.2, 0.5] would apply different weights to three different task vectors.
+
+            exclude_keys (Optional[List[str]], optional): List of parameter names to exclude from TSVM.
+                These parameters will not participate in the SVD computation and merging process.
+                Useful for excluding certain layers (e.g., task-specific heads, normalization layers)
+                that should not be merged across tasks. Defaults to an empty list.
+
+                Example: exclude_keys=['classifier.weight', 'classifier.bias'] to skip classification heads.
+
+            return_single_task_models (bool, optional): Whether to return individual transformed models.
+
+                - If True: Returns a dictionary containing both individual models with their transformed
+                  task vectors applied AND the final merged model. The dictionary has the structure:
+
+                  >>> {'model_name_1': transformed_model_1, ..., 'merged': final_merged_model}
+
+                - If False: Returns only the final merged model.
+
+                This is useful for analysis or when you need access to intermediate results.
+                Defaults to False.
+
+            **kwargs: Additional arguments passed to the parent BaseAlgorithm class.
+
+        Note:
+            The choice between single alpha vs. list of alphas affects the merging strategy:
+            - Single alpha: SVD is applied first, then the result is scaled
+            - List of alphas: Individual task vectors are scaled first, then SVD is applied
+        """
+        self.alpha = alpha
+        self.exclude_keys = exclude_keys if exclude_keys is not None else []
+        self.return_single_task_models = return_single_task_models
+        super().__init__(**kwargs)
+
+    def load_pretrained_model_and_task_vectors(self, modelpool: fb.BaseModelPool):
+        """
+        Load the pretrained base model and compute task vectors from all fine-tuned models.
+
+        This method performs the initial step of the TSVM algorithm by:
+        1. Loading the original pretrained model (before any task-specific fine-tuning)
+        2. For each fine-tuned model in the pool:
+           - Load the fine-tuned model
+           - Compute the task vector (fine-tuned params - pretrained params)
+           - Optionally apply individual scaling if alpha is provided as a list
+
+        Task vectors represent the knowledge gained during fine-tuning and are the core
+        data structure that TSVM operates on.
+
+        Args:
+            modelpool (fb.BaseModelPool): Pool containing the pretrained model and all
+                fine-tuned models to be merged.
+
+        Returns:
+            tuple: A tuple containing:
+                - pretrained_model: The original pretrained model (torch.nn.Module)
+                - task_vectors: List of task vectors (List[StateDictType]), where each
+                  task vector is a state dictionary representing the parameter differences
+                  for one specific task
+        """
+        # Load the original pretrained model that serves as the base for all fine-tuned variants
+        pretrained_model = modelpool.load_pretrained_model()
+
+        # Initialize list to store computed task vectors
+        task_vectors = []
+
+        # Process each fine-tuned model in the modelpool
+        for model_idx, model_name in enumerate(modelpool.model_names):
+            # Load the current fine-tuned model
+            finetuned_model = modelpool.load_model(model_name)
+
+            # Compute task vector: difference between fine-tuned and pretrained parameters
+            # This captures the task-specific adaptations learned during fine-tuning
+            task_vector = state_dict_sub(
+                finetuned_model.state_dict(), pretrained_model.state_dict()
+            )
+            task_vectors.append(task_vector)
+
+            # Apply individual scaling to task vectors if alpha is provided as a list
+            # This allows for task-specific weighting before the SVD computation
+            if self.alpha is not None and isinstance(self.alpha, Iterable):
+                # Ensure the number of alpha values matches the number of models
+                assert len(self.alpha) == len(
+                    modelpool.model_names
+                ), f"Alpha list length ({len(self.alpha)}) must match number of models ({len(modelpool.model_names)})"
+
+                # Scale the current task vector by its corresponding alpha value
+                task_vectors[-1] = state_dict_mul(
+                    state_dict=task_vectors[-1], scalar=self.alpha[model_idx]
+                )
+
+        return pretrained_model, task_vectors
+
+    def run(self, modelpool: fb.BaseModelPool):
+        """
+        Execute the complete Task Singular Vector Merging algorithm.
+
+        This is the main entry point that orchestrates the entire TSVM process:
+
+        The algorithm leverages the mathematical insight that task vectors often lie in a
+        lower-dimensional subspace, and SVD helps identify the most important directions
+        in this subspace while filtering out noise and interference.
+
+        Args:
+            modelpool (fb.BaseModelPool): Pool of models to merge, including:
+                - The pretrained base model
+                - Multiple fine-tuned models (one per task)
+                All models must have compatible architectures.
+
+        Returns:
+            Union[torch.nn.Module, Dict[str, torch.nn.Module]]:
+                - If return_single_task_models=False: Returns the merged model
+                - If return_single_task_models=True: Returns a dictionary with:
+                  * Individual transformed models keyed by their original names
+                  * Final merged model under the key 'merged'
+
+        Raises:
+            AssertionError: If alpha list length doesn't match the number of models
+        """
+        # Determine the compute device for SVD operations (GPU if available for faster computation)
+        accelerator = self.fabric.device
+
+        # Phase 1: Load pretrained model and compute task vectors from all fine-tuned models
+        pretrained_model, task_vectors = self.load_pretrained_model_and_task_vectors(
+            modelpool
+        )
+
+        # Phase 2: Apply SVD-based merging to the task vectors
+        # This is the core of the TSVM algorithm where:
+        # - Task vectors are organized into a matrix
+        # - SVD finds the principal components (most important directions)
+        # - Task vectors are reconstructed using only the most significant components
+        # - The reconstructed vectors are merged (summed) to create a unified task vector
+        new_merged_tv = TSVM_utils.compute_and_sum_svd_mem_reduction(
+            task_vectors,
+            exclude_keys=self.exclude_keys,  # Skip certain parameters from SVD
+            accelerator=accelerator,  # Use GPU if available
+            return_single_task_models=self.return_single_task_models,
+        )
+
+        # Handle the case where individual transformed task vectors are also returned
+        if self.return_single_task_models:
+            new_merged_tv, single_task_models = new_merged_tv
+
+        # Phase 3: Apply global scaling to the merged task vector (if alpha is a single value)
+        # This is different from individual scaling applied earlier - here we scale the
+        # final merged result, which affects the overall strength of all merged adaptations
+        if self.alpha is not None and isinstance(self.alpha, (float, int)):
+            print(f"Scaling new merged task vector by alpha: {self.alpha}")
+            new_merged_tv = state_dict_mul(state_dict=new_merged_tv, scalar=self.alpha)
+
+        # Phase 4: Prepare individual transformed models if requested
+        if self.return_single_task_models:
+            models = {}
+            # Create individual models by adding each transformed task vector to the pretrained base
+            for model_idx, model_name in enumerate(modelpool.model_names):
+                # Create a deep copy to avoid modifying the original pretrained model
+                model = deepcopy(pretrained_model)
+                # Apply the transformed task vector to get the individual model
+                model.load_state_dict(
+                    state_dict_add(model.state_dict(), single_task_models[model_idx])
+                )
+                models[model_name] = model
+
+        # Phase 5: Create the final merged model by adding the merged task vector to pretrained model
+        # This produces a single model that combines capabilities from all input models
+        pretrained_model.load_state_dict(
+            state_dict_add(new_merged_tv, pretrained_model.state_dict())
+        )
+
+        # Phase 6: Return results based on the requested output format
+        if self.return_single_task_models:
+            # Include the final merged model in the dictionary of results
+            models["merged"] = pretrained_model
+            return models
+        else:
+            # Return only the merged model
+            return pretrained_model

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/method/task_singular_vector/utils/TSVM_utils.py

@@ -1,7 +1,9 @@
+import collections
 import math
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 import torch
+from torch import Tensor, nn
 
 from fusion_bench.utils.type import StateDictType
 
@@ -314,7 +316,8 @@ def compute_and_sum_svd_mem_reduction(
     task_vectors: List[StateDictType],
     exclude_keys: Optional[List[str]] = None,
     accelerator: torch.device = "cuda" if torch.cuda.is_available() else "cpu",
-) -> StateDictType:
+    return_single_task_models: bool = False,
+):
     """
     Computes the Singular Value Decomposition (SVD) for each vector in the task_vectors,
     reduces the dimensionality of the vectors based on the sv_reduction factor, and concatenate
@@ -326,23 +329,25 @@ def compute_and_sum_svd_mem_reduction(
                             dictionary of vectors.
         exclude_keys (list): A list of keys to exclude from the TSVM.
         accelerator (torch.device): The device to use for the computation.
+        return_single_task_models (bool): Whether to return the single task models after the TSVM.
 
     Returns:
         dict: A dictionary containing the new vectors after SVD computation and merging.
     """
     if exclude_keys is None:
         exclude_keys = []
-    sv_reduction = 1 / len(task_vectors)
+    num_tasks = len(task_vectors)
+    sv_reduction = 1 / num_tasks
 
-    new_vector = {}
+    new_vector: Dict[str, Tensor] = {}
+    if return_single_task_models:
+        single_task_models = [{} for _ in range(num_tasks)]
     for key in task_vectors[0]:
         original_device = task_vectors[0][key].device
         original_dtype = task_vectors[0][key].dtype
 
-        new_vector[key] = {}
         for i, task_vector in enumerate(task_vectors):
-            vec = task_vector[key].to(accelerator)
-
+            vec = task_vector[key].to(device=accelerator, non_blocking=True)
             if len(task_vector[key].shape) == 2 and key not in exclude_keys:
                 # at current, the SVD is not supported for half precision, so we need to convert to float32
                 if not (
@@ -350,13 +355,14 @@ def compute_and_sum_svd_mem_reduction(
                 ):
                     vec = vec.to(dtype=torch.float32)
 
-                u, s, v = torch.linalg.svd(vec, full_matrices=False)
+                # vec = u @ torch.diag(s) @ vh
+                u, s, vh = torch.linalg.svd(vec, full_matrices=False)
 
                 if i == 0:
                     print(f"Computed SVD for {key}...")
                     sum_u = torch.zeros_like(u, device=accelerator)
                     sum_s = torch.zeros_like(s, device=accelerator)
-                    sum_v = torch.zeros_like(v, device=accelerator)
+                    sum_vh = torch.zeros_like(vh, device=accelerator)
                 reduced_index_s = int(s.shape[0] * sv_reduction)
 
                 # select only the first reduced_index_s columns of u and place them
@@ -367,10 +373,9 @@ def compute_and_sum_svd_mem_reduction(
                     :reduced_index_s
                 ]
                 # select only the first reduced_index_s rows of v and place them
-                sum_v[i * reduced_index_s : (i + 1) * reduced_index_s, :] = v[
+                sum_vh[i * reduced_index_s : (i + 1) * reduced_index_s, :] = vh[
                     :reduced_index_s, :
                 ]
-
             else:
                 # if the vector is not a 2D tensor or is in exclude_keys, compute the mean
                 if i == 0:
@@ -379,22 +384,49 @@ def compute_and_sum_svd_mem_reduction(
                     new_vector[key] += (vec - new_vector[key]) / (i + 1)
 
         if len(task_vector[key].shape) == 2 and key not in exclude_keys:
-            u_u, s_u, v_u = torch.linalg.svd(sum_u, full_matrices=False)
-            u_v, s_v, v_v = torch.linalg.svd(sum_v, full_matrices=False)
+            u_u, s_u, vh_u = torch.linalg.svd(sum_u, full_matrices=False)
+            u_vh, s_vh, vh_vh = torch.linalg.svd(sum_vh, full_matrices=False)
 
             new_vector[key] = torch.linalg.multi_dot(
                 (
                     u_u,
-                    v_u,
+                    vh_u,
                     torch.diag(sum_s),
-                    u_v,
-                    v_v,
+                    u_vh,
+                    vh_vh,
                 )
             )
-        new_vector[key] = new_vector[key].to(
-            device=original_device, dtype=original_dtype, non_blocking=True
-        )
-    return new_vector
+            new_vector[key] = new_vector[key].to(
+                device=original_device, dtype=original_dtype, non_blocking=True
+            )
+            if return_single_task_models:
+                reduced_index_s = int(sum_s.shape[0] * sv_reduction)
+                new_u = u_u @ vh_u
+                new_vh = u_vh @ vh_vh
+                for i in range(num_tasks):
+                    single_task_models[i][key] = torch.linalg.multi_dot(
+                        (
+                            new_u[:, i * reduced_index_s : (i + 1) * reduced_index_s],
+                            torch.diag(
+                                sum_s[i * reduced_index_s : (i + 1) * reduced_index_s]
+                            ),
+                            new_vh[i * reduced_index_s : (i + 1) * reduced_index_s, :],
+                        )
+                    ).to(
+                        device=original_device, dtype=original_dtype, non_blocking=True
+                    )
+        else:
+            new_vector[key] = new_vector[key].to(
+                device=original_device, dtype=original_dtype, non_blocking=True
+            )
+            if return_single_task_models:
+                for i in range(num_tasks):
+                    single_task_models[i][key] = new_vector[key].clone()
+
+    if not return_single_task_models:
+        return new_vector
+    else:
+        return new_vector, single_task_models
 
 
 ###############

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/mixins/hydra_config.py

@@ -9,7 +9,7 @@ from hydra import compose, initialize
 from omegaconf import DictConfig, OmegaConf
 
 from fusion_bench.utils import import_object, instantiate
-from fusion_bench.utils.instantiate import set_print_function_call
+from fusion_bench.utils.instantiate_utils import set_print_function_call
 
 log = logging.getLogger(__name__)
 

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/mixins/lightning_fabric.py

@@ -11,7 +11,7 @@ from lightning.fabric.utilities.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
 
 from fusion_bench.utils import import_object
-from fusion_bench.utils.instantiate import instantiate
+from fusion_bench.utils.instantiate_utils import instantiate
 
 if TYPE_CHECKING:
     import lightning.fabric.loggers.tensorboard
@@ -172,3 +172,27 @@ class LightningFabricMixin:
             return True
         else:
             return False
+
+    def log(self, name: str, value: Any, step: Optional[int] = None):
+        """
+        Logs the metric to the fabric's logger.
+        """
+        self.fabric.log(name, value, step=step)
+
+    def log_dict(self, metrics: dict, step: Optional[int] = None):
+        """
+        Logs the metrics to the fabric's logger.
+        """
+        self.fabric.log_dict(metrics, step=step)
+
+    def log_optimizer_lr(
+        self,
+        optimizer: torch.optim.Optimizer,
+        step: Optional[int] = None,
+        name_template: str = "train/lr_group_{0}",
+    ):
+        """
+        Logs the learning rate of the optimizer to the fabric's logger.
+        """
+        for i, param_group in enumerate(optimizer.param_groups):
+            self.fabric.log(name_template.format(i), param_group["lr"], step=step)

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/mixins/serialization.py

@@ -4,13 +4,14 @@ from typing import Dict, Optional, Union
 
 from omegaconf import OmegaConf
 
-from fusion_bench.utils import instantiate
+from fusion_bench.utils import import_object, instantiate
 
 log = logging.getLogger(__name__)
 
 
 class YAMLSerializationMixin:
     _recursive_: bool = False
+    _config_key: Optional[str] = None
     _config_mapping: Dict[str, str] = {
         "_recursive_": "_recursive_",
     }
@@ -99,7 +100,22 @@ class YAMLSerializationMixin:
             BaseModelPool: The loaded model pool.
         """
         config = OmegaConf.load(path)
-        return instantiate(config, _recursive_=cls._recursive_)
+        if cls._config_key is not None and cls._config_key in config:
+            config = config[cls._config_key]
+        target_cls = import_object(config["_target_"])
+        if target_cls != cls:
+            log.warning(
+                f"The class {target_cls.__name__} is not the same as the class {cls.__name__}. "
+                f"Instantiating the class {target_cls.__name__} instead."
+            )
+        return instantiate(
+            config,
+            _recursive_=(
+                cls._recursive_
+                if config.get("_recursive_") is None
+                else config.get("_recursive_")
+            ),
+        )
 
     def to_config(self):
         """

{fusion_bench-0.2.16 → fusion_bench-0.2.17}/fusion_bench/modelpool/base_pool.py

@@ -29,6 +29,7 @@ class BaseModelPool(BaseYAMLSerializableModel, HydraConfigMixin):
     """
 
     _program = None
+    _config_key = "modelpool"
     _models: Union[DictConfig, Dict[str, nn.Module]]
     _config_mapping = BaseYAMLSerializableModel._config_mapping | {
         "_models": "models",