fusion-bench 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

fusion_bench/method/task_singular_vector/TSVM.py

@@ -1,5 +1,31 @@
-"""
-Example:
+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 \
@@ -7,14 +33,22 @@ fusion_bench \
     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
@@ -35,49 +69,228 @@ from .utils import (
 
 
 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: Union[float, Iterable[float]] = None,
-        remove_keys: Optional[List[str]] = None,
+        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.remove_keys = remove_keys if remove_keys is not None else []
+        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 run(self, modelpool):
-        # Load the pre-trained model and the fine-tuned models
+    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()
-        finetuned_models = list(modelpool.models())
 
-        ptm_check = pretrained_model.state_dict()
-        ft_checks = [model.state_dict() for model in finetuned_models]
-        check_parameterNamesMatch(ft_checks + [ptm_check])
+        # 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)
 
-        with timeit_context("Flattening out Checkpoints"):
-            task_vectors = [state_dict_sub(check, ptm_check) for check in ft_checks]
-            if isinstance(self.alpha, Iterable):
+            # 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(
-                    task_vectors
-                ), "Alpha and task vectors must have the same length"
-                task_vectors = [
-                    state_dict_mul(state_dict=tv, scalar=alpha)
-                    for alpha, tv in zip(self.alpha, task_vectors)
-                ]
+                    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.remove_keys,
-            accelerator=self.fabric.device,
+            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,
         )
 
-        # If alpha is a float, we need to scale the new merged task vector by alpha
-        if self.alpha is not None and isinstance(self.alpha, float):
+        # 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())
         )
-        return pretrained_model
+
+        # 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/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/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/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/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/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",

fusion_bench/modelpool/clip_vision/modelpool.py

@@ -3,6 +3,7 @@ from copy import deepcopy
 from typing import Optional, Union
 
 from datasets import load_dataset
+from lightning.fabric.utilities import rank_zero_only
 from omegaconf import DictConfig, open_dict
 from torch import nn
 from torch.utils.data import Dataset
@@ -40,7 +41,8 @@ class CLIPVisionModelPool(BaseModelPool):
     def load_processor(self, *args, **kwargs) -> CLIPProcessor:
         assert self._processor is not None, "Processor is not defined in the config"
         if isinstance(self._processor, str):
-            log.info(f"Loading `transformers.CLIPProcessor`: {self._processor}")
+            if rank_zero_only.rank == 0:
+                log.info(f"Loading `transformers.CLIPProcessor`: {self._processor}")
             processor = CLIPProcessor.from_pretrained(self._processor)
         else:
             processor = instantiate(self._processor, *args, **kwargs)
@@ -50,7 +52,8 @@ class CLIPVisionModelPool(BaseModelPool):
         model_config = self._models[model_name]
 
         if isinstance(model_config, str):
-            log.info(f"Loading `transformers.CLIPModel`: {model_config}")
+            if rank_zero_only.rank == 0:
+                log.info(f"Loading `transformers.CLIPModel`: {model_config}")
             clip_model = CLIPModel.from_pretrained(model_config, *args, **kwargs)
             return clip_model
         else:
@@ -102,10 +105,12 @@ class CLIPVisionModelPool(BaseModelPool):
         ):
             model = self._models[model_name_or_config]
             if isinstance(model, str):
-                log.info(f"Loading `transformers.CLIPVisionModel`: {model}")
+                if rank_zero_only.rank == 0:
+                    log.info(f"Loading `transformers.CLIPVisionModel`: {model}")
                 return CLIPVisionModel.from_pretrained(model, *args, **kwargs)
             if isinstance(model, nn.Module):
-                log.info(f"Returning existing model: {model}")
+                if rank_zero_only.rank == 0:
+                    log.info(f"Returning existing model: {model}")
                 return model
 
         # If the model is not a string, we use the default load_model method
@@ -114,9 +119,10 @@ class CLIPVisionModelPool(BaseModelPool):
     def load_train_dataset(self, dataset_name: str, *args, **kwargs):
         dataset_config = self._train_datasets[dataset_name]
         if isinstance(dataset_config, str):
-            log.info(
-                f"Loading train dataset using `datasets.load_dataset`: {dataset_config}"
-            )
+            if rank_zero_only.rank == 0:
+                log.info(
+                    f"Loading train dataset using `datasets.load_dataset`: {dataset_config}"
+                )
             dataset = load_dataset(dataset_config, split="train")
         else:
             dataset = super().load_train_dataset(dataset_name, *args, **kwargs)
@@ -125,9 +131,10 @@ class CLIPVisionModelPool(BaseModelPool):
     def load_val_dataset(self, dataset_name: str, *args, **kwargs):
         dataset_config = self._val_datasets[dataset_name]
         if isinstance(dataset_config, str):
-            log.info(
-                f"Loading validation dataset using `datasets.load_dataset`: {dataset_config}"
-            )
+            if rank_zero_only.rank == 0:
+                log.info(
+                    f"Loading validation dataset using `datasets.load_dataset`: {dataset_config}"
+                )
             dataset = load_dataset(dataset_config, split="validation")
         else:
             dataset = super().load_val_dataset(dataset_name, *args, **kwargs)
@@ -136,9 +143,10 @@ class CLIPVisionModelPool(BaseModelPool):
     def load_test_dataset(self, dataset_name: str, *args, **kwargs):
         dataset_config = self._test_datasets[dataset_name]
         if isinstance(dataset_config, str):
-            log.info(
-                f"Loading test dataset using `datasets.load_dataset`: {dataset_config}"
-            )
+            if rank_zero_only.rank == 0:
+                log.info(
+                    f"Loading test dataset using `datasets.load_dataset`: {dataset_config}"
+                )
             dataset = load_dataset(dataset_config, split="test")
         else:
             dataset = super().load_test_dataset(dataset_name, *args, **kwargs)

fusion_bench/modelpool/lazy_state_dict_pool.py

@@ -0,0 +1,15 @@
+from fusion_bench import BaseModelPool
+from fusion_bench.utils import instantiate
+from fusion_bench.utils.lazy_state_dict import LazyStateDict
+
+
+class LazyStateDictPool(BaseModelPool):
+    def load_model(self, model_name_or_config: str, *args, **kwargs) -> LazyStateDict:
+        if model_name_or_config in self._models:
+            checkpoint_config = self._models[model_name_or_config]
+        else:
+            checkpoint_config = model_name_or_config
+        if isinstance(checkpoint_config, str):
+            return LazyStateDict(checkpoint_config, *args, **kwargs)
+        else:
+            return instantiate(checkpoint_config, *args, **kwargs)

fusion_bench/models/modeling_deepseek_v2/__init__.py

@@ -0,0 +1,15 @@
+"""
+This is a direct copy of the DeepSeek-V2-Lite model from HuggingFace https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite/tree/main
+"""
+
+from .configuration_deepseek import DeepseekV2Config
+from .modeling_deepseek import (
+    DeepseekV2ForCausalLM,
+    DeepseekV2ForSequenceClassification,
+    DeepseekV2MLP,
+    DeepseekV2Model,
+    DeepseekV2MoE,
+    DeepseekV2DecoderLayer,
+)
+from .modeling_deepseek import MoEGate as DeepseekV2MoEGate
+from .tokenization_deepseek_fast import DeepseekTokenizerFast