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

fusion_bench/method/ensemble.py

@@ -1,11 +1,12 @@
 import logging
-from typing import List, Mapping, Union  # noqa: F401
+from typing import List, Mapping, Optional, Union  # noqa: F401
 
 import numpy as np
 import torch
 from torch import nn
 
 from fusion_bench.method import BaseAlgorithm
+from fusion_bench.mixins import auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.models.wrappers.ensemble import (
     EnsembleModule,
@@ -18,7 +19,7 @@ log = logging.getLogger(__name__)
 
 class SimpleEnsembleAlgorithm(BaseAlgorithm):
     @torch.no_grad()
-    def run(self, modelpool: BaseModelPool | List[nn.Module]):
+    def run(self, modelpool: BaseModelPool | List[nn.Module]) -> EnsembleModule:
         """
         Run the simple ensemble algorithm on the given model pool.
 
@@ -35,20 +36,19 @@ class SimpleEnsembleAlgorithm(BaseAlgorithm):
         return ensemble
 
 
+@auto_register_config
 class WeightedEnsembleAlgorithm(BaseAlgorithm):
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "normalize": "normalize",
-        "weights": "weights",
-    }
-
-    def __init__(self, normalize: bool, weights: List[float], **kwargs):
-        self.normalize = normalize
-        self.weights = weights
+    def __init__(
+        self,
+        normalize: bool = True,
+        weights: Optional[List[float]] = None,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
 
     @torch.no_grad()
-    def run(self, modelpool: BaseModelPool | List[nn.Module]):
+    def run(self, modelpool: BaseModelPool | List[nn.Module]) -> WeightedEnsembleModule:
         """
         Run the weighted ensemble algorithm on the given model pool.
 
@@ -78,7 +78,7 @@ class WeightedEnsembleAlgorithm(BaseAlgorithm):
 
 class MaxModelPredictorAlgorithm(BaseAlgorithm):
     @torch.no_grad()
-    def run(self, modelpool: BaseModelPool | List[nn.Module]):
+    def run(self, modelpool: BaseModelPool | List[nn.Module]) -> MaxModelPredictor:
         """
         Run the max model predictor algorithm on the given model pool.
 

fusion_bench/method/expert_sparsity/utils/calibration_data.py

@@ -12,9 +12,9 @@ import os
 import torch
 import transformers
 from datasets import load_dataset
+from huggingface_hub import hf_hub_download
 from transformers import PreTrainedTokenizer, default_data_collator
 from transformers.testing_utils import CaptureLogger
-from huggingface_hub import hf_hub_download
 
 logger = logging.getLogger(__name__)
 

fusion_bench/method/fisher_merging/clip_fisher_merging.py

@@ -65,7 +65,7 @@ class FisherMergingForCLIPVisionModel(
             minimal_fisher_weight=minimal_fisher_weight,
             num_fisher_examples=num_fisher_examples,
         )
-        self._dataloader_kwargs = dataloader_kwargs
+        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}")
@@ -127,7 +127,7 @@ class FisherMergingForCLIPVisionModel(
         """
         # setup dataloader
         train_dataset = CLIPDataset(train_dataset, self.clip_processor)
-        train_dataloader = DataLoader(train_dataset, **self._dataloader_kwargs)
+        train_dataloader = DataLoader(train_dataset, **self.dataloader_kwargs)
         if self.fabric is not None:
             train_dataloader = self.fabric.setup_dataloaders(train_dataloader)
             model = self.fabric.setup(model)

fusion_bench/method/fisher_merging/fisher_merging.py

@@ -5,14 +5,14 @@ This implementation is largely based on the implementation from  https://github.
 import logging
 import re
 from collections import defaultdict
-from typing import Dict, List
+from typing import Any, Dict, List
 
 import torch
 from torch import Tensor, nn
 from tqdm.autonotebook import tqdm
 
 from fusion_bench.method import BaseAlgorithm
-from fusion_bench.mixins import SimpleProfilerMixin
+from fusion_bench.mixins import SimpleProfilerMixin, auto_register_config
 from fusion_bench.modelpool import BaseModelPool
 
 log = logging.getLogger(__name__)
@@ -353,6 +353,7 @@ def filter_state_dict(
     return filtered_state_dict
 
 
+@auto_register_config
 class FisherMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
     """
     Implements the Fisher Merging Algorithm.
@@ -365,13 +366,6 @@ class FisherMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
             Executes the Fisher merging process on the model pool and returns the merged model.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "exclude_param_names_regex": "exclude_param_names_regex",
-        "normalize_fisher_weight": "normalize_fisher_weight",
-        "minimal_fisher_weight": "minimal_fisher_weight",
-        "num_fisher_examples": "num_fisher_examples",
-    }
-
     def __init__(
         self,
         *,
@@ -379,12 +373,9 @@ class FisherMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         normalize_fisher_weight: bool,
         minimal_fisher_weight: float,
         num_fisher_examples: int,
+        **kwargs,
     ):
-        super().__init__()
-        self.exclude_param_names_regex = exclude_param_names_regex
-        self.normalize_fisher_weight = normalize_fisher_weight
-        self.minimal_fisher_weight = minimal_fisher_weight
-        self.num_fisher_examples = num_fisher_examples
+        super().__init__(**kwargs)
 
     def run(self, modelpool: BaseModelPool) -> nn.Module:
         """
@@ -469,7 +460,7 @@ class FisherMergingAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         self,
         model_name: str,
         model: nn.Module,
-        train_dataset,
+        train_dataset: Any,
         param_names_to_merge: List[str],
     ) -> Dict[str, Tensor]:
         """

fusion_bench/method/fisher_merging/gpt2_fisher_merging.py

@@ -18,13 +18,14 @@ from transformers.models.gpt2.modeling_gpt2 import Conv1D
 from fusion_bench.mixins import LightningFabricMixin
 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
 
 
+@auto_register_config
 class FisherMergingAlgorithmForGPT2(
-    FisherMergingAlgorithm,
     LightningFabricMixin,
+    FisherMergingAlgorithm,
 ):
     """
     Implements the Fisher Merging Algorithm for GPT-2 models on text classification tasks.
@@ -42,11 +43,6 @@ class FisherMergingAlgorithmForGPT2(
 
     classifiers = {}
     modelpool: GPT2ForSequenceClassificationPool = None
-    _config_mapping = FisherMergingAlgorithm._config_mapping | {
-        "cache_dir": "cache_dir",
-        "batch_size": "batch_size",
-        "num_workers": "num_workers",
-    }
 
     def __init__(
         self,
@@ -64,9 +60,6 @@ class FisherMergingAlgorithmForGPT2(
             num_workers (int): Number of workers for data loading.
             **kwargs: Additional keyword arguments.
         """
-        self.cache_dir = cache_dir
-        self.batch_size = batch_size
-        self.num_workers = num_workers
         super().__init__(**kwargs)
 
     def on_fisher_merging_start(self):

fusion_bench/method/fw_merging/fw_hard.py

@@ -223,7 +223,7 @@ class FrankWolfeHardAlgorithm(
     def get_shuffled_loader_iter(self, task: str):
         if self.loss_fn == "cross_entropy":
             # get dataloader kwargs
-            dataloader_kwargs = self._dataloader_kwargs.copy()
+            dataloader_kwargs = self.dataloader_kwargs.copy()
             dataloader_kwargs["shuffle"] = True
             dataloader_kwargs["batch_size"] = 1
 

fusion_bench/method/fw_merging/fw_soft.py

@@ -193,7 +193,7 @@ class FrankWolfeSoftAlgorithm(
     @functools.cache
     def get_shuffled_train_loader_iter(self, task: str, batch_size: int = 1):
         # get dataloader kwargs
-        dataloader_kwargs = self._dataloader_kwargs.copy()
+        dataloader_kwargs = self.dataloader_kwargs.copy()
         dataloader_kwargs["shuffle"] = True
         dataloader_kwargs["batch_size"] = batch_size
 

fusion_bench/method/gossip/clip_layer_wise_gossip.py

@@ -3,13 +3,12 @@ Example Usage:
 
 ```bash
 fusion_bench \
-    method=adamerging \
+    path.log_dir=outputs/ViT-B-32/gossip_layer_wise_adamerging_adam \
+    method=adamerging/clip \
         method.name=clip_layer_wise_adamerging \
         method.save_merging_weights=merging_weights.pt \
-    modelpool=clip-vit-base-patch32_TA8 \
-    taskpool=clip-vit-classification_TA8 \
-    fabric_logger.root_dir=outputs/logs/ViT-B-32 \
-    fabric_logger.name=clip_layer_wise_adamerging_adam
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
 ```
 """
 

fusion_bench/method/linear/expo.py

@@ -7,6 +7,7 @@ Reference:
 
 import logging
 from copy import deepcopy
+from typing import Union
 
 import torch
 from torch import nn
@@ -79,7 +80,7 @@ class ExPOAlgorithm(BaseAlgorithm):
         self.extrapolation_factor = extrapolation_factor
         super().__init__(**kwargs)
 
-    def run(self, modelpool: BaseModelPool):
+    def run(self, modelpool: Union[BaseModelPool, list]) -> nn.Module:
         """
         Run the ExPO merge algorithm.
 

fusion_bench/method/linear/linear_interpolation.py

@@ -1,6 +1,8 @@
 import logging
+from typing import Any
 
 import torch
+from torch import nn
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
 from fusion_bench.utils.state_dict_arithmetic import state_dict_weighted_sum
@@ -10,7 +12,7 @@ log = logging.getLogger(__name__)
 
 class LinearInterpolationAlgorithm(BaseAlgorithm):
     R"""
-    LinearInterpolationAlgorithm performs linear interpolation between two models.
+    `LinearInterpolationAlgorithm` performs linear interpolation between two models.
     Returns a model with the state dict that is a linear interpolation of the state dicts of the two models.
     $\theta = (1-t) \theta_1 + t \theta_2$
     """
@@ -19,9 +21,9 @@ class LinearInterpolationAlgorithm(BaseAlgorithm):
         "t": "t",
     }
 
-    def __init__(self, t: float, **kwargs):
+    def __init__(self, t: float, **kwargs: Any):
         """
-        Initialize the LinearInterpolationAlgorithm with the given interpolation parameter.
+        Initialize the `LinearInterpolationAlgorithm` with the given interpolation parameter.
 
         Args:
             t (float): The interpolation parameter, should be in the range [0, 1].
@@ -31,7 +33,7 @@ class LinearInterpolationAlgorithm(BaseAlgorithm):
         self.t = t
         super().__init__(**kwargs)
 
-    def run(self, modelpool: BaseModelPool):
+    def run(self, modelpool: BaseModelPool) -> nn.Module:
         """
         Run the linear interpolation algorithm on the given model pool.
 

fusion_bench/method/linear/simple_average_for_llama.py

@@ -1,15 +1,15 @@
 from copy import deepcopy
 from typing import TYPE_CHECKING, Optional
 
+from omegaconf import flag_override
 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.modelpool import CausalLMBackbonePool, CausalLMPool
-from fusion_bench.utils.pylogger import getRankZeroLogger
-from omegaconf import flag_override
 from fusion_bench.utils import instantiate
+from fusion_bench.utils.pylogger import getRankZeroLogger
 
 log = getRankZeroLogger(__name__)
 
@@ -19,7 +19,6 @@ class SimpleAverageForLlama(BaseAlgorithm):
     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.
 
     Examples:
-
         The following example demonstrates how to use the `SimpleAverageForLlama` algorithm to merge Mistral models.
 
         ```bash

fusion_bench/method/lm_finetune/bradley_terry_rm.py

@@ -31,7 +31,7 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer
 from fusion_bench.dataset.llama.collate import bradley_terry_rm_collate
 from fusion_bench.method import BaseAlgorithm
 from fusion_bench.mixins import FabricTrainingMixin
-from fusion_bench.modelpool import SeqenceClassificationModelPool
+from fusion_bench.modelpool import SequenceClassificationModelPool
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.dtype import get_dtype
 
@@ -121,7 +121,7 @@ class BradleyTerryRewardModeling(BaseAlgorithm, FabricTrainingMixin):
         self.fix_token_embedding = fix_token_embedding
         super().__init__(**kwargs)
 
-    def run(self, modelpool: SeqenceClassificationModelPool):
+    def run(self, modelpool: SequenceClassificationModelPool):
         self.modelpool = modelpool
         self.setup()
         self.train(self.model, self.optimizer, self.lr_scheduler)

fusion_bench/method/mixture_of_experts/mixtral_upcycling.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Any, Optional
 
 import torch
 from tqdm.autonotebook import tqdm
@@ -23,8 +23,7 @@ from transformers.models.mixtral.modeling_mixtral import (
 )
 from transformers.utils import ContextManagers
 
-from fusion_bench.method import BaseAlgorithm
-from fusion_bench.modelpool import BaseModelPool
+from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
 
 log = logging.getLogger(__name__)
 
@@ -114,7 +113,7 @@ def _upscale_decoder_layer(
 
 def upscale_to_mixtral_model(
     input_model: LlamaModel | MistralModel, output_model: MixtralModel
-):
+) -> None:
     """
     A helper function.
 
@@ -140,7 +139,7 @@ def upscale_to_mixtral_model(
 
 def upscale_to_mixtral_for_causal_lm(
     input_model: LlamaForCausalLM | MistralForCausalLM, output_model: MixtralForCausalLM
-):
+) -> None:
     """
     A helper function.
 
@@ -157,24 +156,19 @@ def upscale_to_mixtral_for_causal_lm(
     upscale_to_mixtral_model(input_model.model, output_model.model)
 
 
+@auto_register_config
 class MixtralUpscalingAlgorithm(BaseAlgorithm):
     """
     This class is responsible for upscaling a model to a MixtralModel.
     It inherits from the ModelFusionAlgorithm class.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "num_experts": "num_experts",
-        "experts_per_token": "experts_per_token",
-        "save_checkpoint": "save_checkpoint",
-    }
-
     def __init__(
         self,
         num_experts: int,
         experts_per_token: int,
         save_checkpoint: str,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Initialize the MixtralUpscalingAlgorithm.
@@ -185,9 +179,6 @@ class MixtralUpscalingAlgorithm(BaseAlgorithm):
             save_checkpoint (str): The path to save the checkpoint.
             **kwargs: Additional keyword arguments.
         """
-        self.num_experts = num_experts
-        self.experts_per_token = experts_per_token
-        self.save_checkpoint = save_checkpoint
         super().__init__(**kwargs)
 
     @torch.no_grad()
@@ -242,24 +233,19 @@ class MixtralUpscalingAlgorithm(BaseAlgorithm):
         return mixtral_model
 
 
+@auto_register_config
 class MixtralForCausalLMUpscalingAlgorithm(BaseAlgorithm):
     """
     This class is responsible for upscaling a model to a MixtralForCausalLM.
     It inherits from the ModelFusionAlgorithm class.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "num_experts": "num_experts",
-        "experts_per_token": "experts_per_token",
-        "save_checkpoint": "save_checkpoint",
-    }
-
     def __init__(
         self,
         num_experts: int,
         experts_per_token: int,
         save_checkpoint: str,
-        **kwargs,
+        **kwargs: Any,
     ):
         """
         Initialize the MixtralForCausalLMUpscalingAlgorithm.
@@ -270,9 +256,6 @@ class MixtralForCausalLMUpscalingAlgorithm(BaseAlgorithm):
             save_checkpoint (str): The path to save the checkpoint.
             **kwargs: Additional keyword arguments.
         """
-        self.num_experts = num_experts
-        self.experts_per_token = experts_per_token
-        self.save_checkpoint = save_checkpoint
         super().__init__(**kwargs)
 
     @torch.no_grad()
@@ -302,7 +285,7 @@ class MixtralForCausalLMUpscalingAlgorithm(BaseAlgorithm):
             self.config.experts_per_token,
         )
 
-        with ContextManagers([no_init_weights(True)]):
+        with ContextManagers([no_init_weights()]):
             for _ in tqdm(range(1), desc="Initializing Mixtral model"):
                 mixtral_model = MixtralForCausalLM(mixtral_config)
         upscale_to_mixtral_for_causal_lm(pretrained_model, mixtral_model)

fusion_bench/method/model_recombination.py

@@ -5,6 +5,7 @@ from typing import List, Mapping, Union  # noqa: F401
 import torch
 from torch import nn
 
+from fusion_bench import auto_register_config
 from fusion_bench.method import BaseAlgorithm
 from fusion_bench.modelpool import BaseModelPool
 
@@ -52,17 +53,13 @@ def recombine_state_dict(models: List[nn.Module]):
     return models
 
 
+@auto_register_config
 class ModelRecombinationAlgorithm(BaseAlgorithm):
     """
     Model recombination recombinates the layers of the given models, to create a new set of models.
     """
 
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "return_modelpool": "return_modelpool",
-    }
-
     def __init__(self, return_modelpool: bool, **kwargs):
-        self.return_modelpool = return_modelpool
         super().__init__(**kwargs)
 
     @torch.no_grad()

fusion_bench/method/moe_pruner/hooks/__init__.py

@@ -1,6 +1,5 @@
-from .hook import BaseHookFn
 from .deepseek_v2 import (
     MoEPrunerHookFnForDeepseekV2Gate,
     MoEPrunerHookFnForDeepseekV2Linear,
 )
-
+from .hook import BaseHookFn

fusion_bench/method/moe_pruner/utils/data.py

@@ -1,8 +1,9 @@
 # Code adapted from https://github.com/IST-DASLab/sparsegpt/blob/master/datautils.py
 
+import os
 import random
 from typing import List, Optional, Tuple, cast  # noqa: F401
-import os
+
 from datasets import load_dataset
 from torch import Tensor
 from tqdm.auto import tqdm

fusion_bench/method/moe_pruner/utils/prune.py

@@ -107,7 +107,12 @@ def prepare_calibration_input(
         device=device,
         requires_grad=False,
     )
-    cache = {"i": 0, "attention_mask": None, "position_ids": None, 'position_embeddings': None}
+    cache = {
+        "i": 0,
+        "attention_mask": None,
+        "position_ids": None,
+        "position_embeddings": None,
+    }
 
     class Catcher(nn.Module):
         def __init__(self, module):

fusion_bench/method/pruning/llama_magnitude_prune.py

@@ -167,7 +167,7 @@ class MagnitudePruningForLlama(BaseAlgorithm, SimpleProfilerMixin):
         super().__init__(**kwargs)
 
     @torch.no_grad()
-    def run(self, modelpool: CausalLMPool):
+    def run(self, modelpool: CausalLMPool) -> LlamaForCausalLM:
         """
         Execute the pruning process on the first model from the given model pool.
 

fusion_bench/method/pruning/wanda_utils/data.py

@@ -4,12 +4,11 @@ import os
 import random
 from typing import List, Optional, Tuple, cast  # noqa: F401
 
+from datasets import load_dataset
 from torch import Tensor
 from tqdm.auto import tqdm
 from transformers import PreTrainedTokenizer
 
-from datasets import load_dataset
-
 
 # Wrapper for tokenized input IDs
 class TokenizerWrapper:

fusion_bench/method/pwe_moe/clip_pwe_moe.py

@@ -16,14 +16,18 @@ from transformers import CLIPVisionModel
 from transformers.models.clip.modeling_clip import CLIPEncoderLayer
 from typing_extensions import override
 
-from fusion_bench.method.base_algorithm import BaseAlgorithm
+from fusion_bench import (
+    BaseAlgorithm,
+    auto_register_config,
+    print_parameters,
+    timeit_context,
+)
+from fusion_bench.dataset import CLIPDataset
 from fusion_bench.method.task_arithmetic import task_arithmetic_merge
 from fusion_bench.mixins.clip_classification import CLIPClassificationMixin
 from fusion_bench.mixins.simple_profiler import SimpleProfilerMixin
 from fusion_bench.modelpool import CLIPVisionModelPool
-from fusion_bench.utils import timeit_context
 from fusion_bench.utils.data import InfiniteDataLoader
-from fusion_bench.utils.parameters import print_parameters
 
 from .module import ParetoWeightEnsemblingModule
 from .utils import generate_simplex_grid
@@ -31,27 +35,13 @@ from .utils import generate_simplex_grid
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class PWEMoEAlgorithmForCLIP(
     BaseAlgorithm,
     SimpleProfilerMixin,
     CLIPClassificationMixin,
 ):
     modelpool: CLIPVisionModelPool = None
-    _config_mapping = BaseAlgorithm._config_mapping | {
-        "upscale_mlp": "upscale_mlp",
-        "upscale_attn": "upscale_attn",
-        "init_lambda": "init_lambda",
-        "router_hidden_layers": "router_hidden_layers",
-        "lr": "lr",
-        "num_steps": "num_steps",
-        "save_interval": "save_interval",
-        "alpha": "alpha",
-        "checkpoint_path": "checkpoint_path",
-        "eval_grid": "eval_grid",
-        "eval_grid_n": "eval_grid_n",
-        "eval_grid_m": "eval_grid_m",
-        "_dataloader_kwargs": "dataloader_kwargs",
-    }
 
     def __init__(
         self,
@@ -72,19 +62,6 @@ class PWEMoEAlgorithmForCLIP(
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.upscale_mlp = upscale_mlp
-        self.upscale_attn = upscale_attn
-        self.init_lambda = init_lambda
-        self.router_hidden_layers = router_hidden_layers
-        self.lr = lr
-        self.num_steps = num_steps
-        self.save_interval = save_interval
-        self.alpha = alpha
-        self.checkpoint_path = checkpoint_path
-        self.eval_grid = eval_grid
-        self.eval_grid_n = eval_grid_n
-        self.eval_grid_m = eval_grid_m
-        self._dataloader_kwargs = dataloader_kwargs
 
     @override
     def run(self, modelpool: CLIPVisionModelPool):
@@ -193,13 +170,14 @@ class PWEMoEAlgorithmForCLIP(
         Loads the datasets specified in the configuration.
         """
         train_datasets = {
-            dataset_name: self.modelpool.load_train_dataset(
-                dataset_name, self.clip_processor
+            dataset_name: CLIPDataset(
+                self.modelpool.load_train_dataset(dataset_name),
+                processor=self.clip_processor,
             )
             for dataset_name in self.modelpool.model_names
         }
         train_loaders = {
-            dataset_name: DataLoader(dataset, shuffle=True, **self._dataloader_kwargs)
+            dataset_name: DataLoader(dataset, shuffle=True, **self.dataloader_kwargs)
             for dataset_name, dataset in train_datasets.items()
         }
         train_loaders = {

fusion_bench/method/randes/modelsoup.py

@@ -5,9 +5,7 @@ import torch
 
 from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.parameters import count_parameters
-from fusion_bench.utils.state_dict_arithmetic import (
-    state_dict_mul,
-)
+from fusion_bench.utils.state_dict_arithmetic import state_dict_mul
 
 from .base_algorithm import SuperposedAlgorithmBase, compare_models
 

fusion_bench/method/regmean/clip_regmean.py

@@ -27,7 +27,7 @@ class RegMeanAlgorithmForCLIP(
 
     def __init__(self, *, dataloader_kwargs: DictConfig, **kwargs):
         super().__init__(**kwargs)
-        self._dataloader_kwargs = dataloader_kwargs
+        self.dataloader_kwargs = dataloader_kwargs
 
     def on_regmean_start(self):
         self.setup_zero_shot_classification_head()
@@ -60,7 +60,7 @@ class RegMeanAlgorithmForCLIP(
         # setup dataloader
         train_dataset = CLIPDataset(train_dataset, self.clip_processor)
         train_dataloader = DataLoader(
-            train_dataset, shuffle=True, **self._dataloader_kwargs
+            train_dataset, shuffle=True, **self.dataloader_kwargs
         )
         train_dataloader = self.fabric.setup_dataloaders(train_dataloader)
         model = self.fabric.setup(model)