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

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)

fusion_bench/method/regmean/gpt2_regmean.py

@@ -15,7 +15,7 @@ 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.utils import timeit_context
 
 from .regmean import RegMeanAlgorithm
@@ -23,22 +23,15 @@ from .regmean import RegMeanAlgorithm
 log = logging.getLogger(__name__)
 
 
+@auto_register_config
 class RegMeanAlgorithmForGPT2(
-    RegMeanAlgorithm,
     LightningFabricMixin,
+    RegMeanAlgorithm,
 ):
     _include_module_type = [Conv1D]
     classifiers = {}
-    _config_mapping = RegMeanAlgorithm._config_mapping | {
-        "cache_dir": "cache_dir",
-        "batch_size": "batch_size",
-        "num_workers": "num_workers",
-    }
 
     def __init__(self, cache_dir: str, batch_size: int, num_workers: int, **kwargs):
-        self.cache_dir = cache_dir
-        self.batch_size = batch_size
-        self.num_workers = num_workers
         super().__init__(**kwargs)
 
     def on_regmean_start(self):

fusion_bench/method/regmean/regmean.py

@@ -13,7 +13,7 @@ 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__)
@@ -280,14 +280,9 @@ def regmean_merging(
     return merged_params
 
 
+@auto_register_config
 class RegMeanAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
     _include_module_type = [nn.Linear]
-    _config_mapping = {
-        "num_regmean_examples": "num_regmean_examples",
-        "exclude_param_names_regex": "exclude_param_names_regex",
-        "reduce_non_diagonal_ratio": "reduce_non_diagonal_ratio",
-        "weight_transpose": "weight_transpose",
-    }
 
     def __init__(
         self,
@@ -298,10 +293,6 @@ class RegMeanAlgorithm(BaseAlgorithm, SimpleProfilerMixin):
         weight_transpose: bool,
         **kwargs,
     ):
-        self.num_regmean_examples = num_regmean_examples
-        self.exclude_param_names_regex = exclude_param_names_regex
-        self.reduce_non_diagonal_ratio = reduce_non_diagonal_ratio
-        self.weight_transpose = weight_transpose
         super().__init__(**kwargs)
 
     def run(self, modelpool: BaseModelPool, **kwargs):

fusion_bench/method/regmean_plusplus/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa F401
+from .clip_regmean_plusplus import RegMeanAlgorithmForCLIPPlusPlus
+from .regmean_plusplus import RegMeanAlgorithmPlusPlus

fusion_bench/method/regmean_plusplus/clip_regmean_plusplus.py

@@ -0,0 +1,199 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, cast  # noqa: F401
+
+import torch
+import torch.utils.data
+from omegaconf import DictConfig
+from torch import Tensor, nn
+from torch.nn.modules import Module
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+from fusion_bench.dataset.clip_dataset import CLIPDataset
+from fusion_bench.mixins import CLIPClassificationMixin
+
+from .regmean_plusplus import RegMeanAlgorithmPlusPlus
+
+log = logging.getLogger(__name__)
+
+
+class RegMeanAlgorithmForCLIPPlusPlus(
+    RegMeanAlgorithmPlusPlus,
+    CLIPClassificationMixin,
+):
+    _config_mapping = {
+        "_dataloader_kwargs": "dataloader_kwargs",
+    }
+
+    def __init__(self, *, dataloader_kwargs: DictConfig, **kwargs):
+        super().__init__(**kwargs)
+        self.dataloader_kwargs = dataloader_kwargs
+
+    def on_regmean_start(self):
+        self.setup_zero_shot_classification_head()
+
+    def compute_logits(self, module, batch, task: str) -> Tensor:
+        images, _ = batch
+        text_embeds = self.zeroshot_weights[task]
+
+        image_embeds = module(images)[1]
+        image_embeds = self.visual_projection(image_embeds)
+
+        # normalize embeddings
+        image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+
+        # cosine similarity
+        logits_per_text = (
+            torch.matmul(text_embeds, image_embeds.t()) * self.logit_scale_exp
+        )
+        logits_per_image = logits_per_text.t()
+
+        return logits_per_image
+
+    def get_regmean_weights(
+        self,
+        model_name: str,
+        layer: Module,
+        batches_input: List[Tensor],
+        linear_modules_to_merge: Dict[str, Module],
+    ):
+        layer = self.fabric.setup(layer)
+
+        def compute_regmean_weights(module_name: str):
+            """
+            compute the regmean weights, a hook function to deal with each module's input
+            :param module_name: str, module name
+            :return:
+            """
+
+            def hook(module: nn.Module, input: tuple, output: torch.Tensor):
+                # Tensor, shape (batch_size, sequence_length, hidden_dim)
+                x = cast(Tensor, input[0]).detach()
+                batch_num_actual_examples = x.shape[0]
+                # Tensor, shape (batch_size * sequence_length, hidden_dim)
+                x = x.reshape(-1, x.shape[-1])
+                # Tensor, shape (hidden_dim, hidden_dim)
+                xtx = torch.matmul(x.transpose(0, 1), x)
+                # store the averaged weights in regmean_weights
+                if module_name not in regmean_weights.keys():
+                    regmean_weights[module_name] = xtx / x.shape[0]
+                    num_computed_examples[module_name] = x.shape[0]
+                    num_actual_examples[module_name] = batch_num_actual_examples
+                else:
+                    regmean_weights[module_name] = (
+                        regmean_weights[module_name]
+                        * num_computed_examples[module_name]
+                        + xtx
+                    ) / (num_computed_examples[module_name] + x.shape[0])
+                    num_computed_examples[module_name] += x.shape[0]
+                    num_actual_examples[module_name] += batch_num_actual_examples
+
+            return hook
+
+        handles = []
+        # dictionary, regmean matrices for each linear module inputs
+        regmean_weights = {}
+        # dictionary, number of examples (multiplied the sequence length) used for computing regmean matrices
+        num_computed_examples = {}
+        # dictionary, number of actual examples used for computing regmean matrices
+        num_actual_examples = {}
+
+        for module_name, linear_module_to_merge in linear_modules_to_merge.items():
+            # register a hook in the forward process
+            handle = linear_module_to_merge.register_forward_hook(
+                compute_regmean_weights(module_name=module_name)
+            )
+            handles.append(handle)
+        _ = self.layer_batches_forward(layer, batches_input)
+
+        # remove the added hook
+        for handle in handles:
+            handle.remove()
+
+        for module_name in regmean_weights.keys():
+            regmean_weights[module_name] = regmean_weights[module_name].detach().cpu()
+
+        return regmean_weights
+
+    def merge_embedding_layer(self, models_to_merge_dict: Dict[str, nn.Module]):
+        models_to_merge_param_dict = defaultdict(list)
+
+        # get the parameters of the embedding layer from each model
+        for model_to_merge in models_to_merge_dict.values():
+            model_to_merge_state_dict = model_to_merge.state_dict()
+
+            param_dict = {}
+            for name, param in model_to_merge_state_dict.items():
+                if name.startswith("vision_model.embeddings") or name.startswith(
+                    "vision_model.pre_layrnorm"
+                ):
+                    param_dict[name] = param
+
+            for param_name in param_dict.keys():
+                models_to_merge_param_dict[param_name].append(param_dict[param_name])
+
+        # merge the parameters of the embedding layer
+        merged_params_dict = {}
+        for param_name, param_list in models_to_merge_param_dict.items():
+            merged_params_dict[param_name] = torch.stack(param_list).mean(dim=0)
+
+        return merged_params_dict
+
+    def get_input_for_first_layer(self, model: nn.Module, train_dataset):
+        # setup dataloader
+        train_dataset = CLIPDataset(train_dataset, self.clip_processor)
+        train_dataloader = DataLoader(
+            train_dataset, shuffle=True, **self.dataloader_kwargs
+        )
+        train_dataloader = self.fabric.setup_dataloaders(train_dataloader)
+        model = self.fabric.setup(model)
+
+        def compute_input(model, batch):
+            images, _ = batch
+
+            images = images.to(model.device)
+            image_embeds = model.vision_model.embeddings(images)
+            image_embeds = model.vision_model.pre_layrnorm(image_embeds)
+            image_embeds = image_embeds.detach().cpu()
+
+            return image_embeds
+
+        num_computed_examples = 0
+        num_regmean_examples = self.num_regmean_examples
+
+        batches_input = []
+        for batch in train_dataloader:
+            if num_computed_examples >= num_regmean_examples:
+                break
+            batches_input.append(compute_input(model, batch))
+            num_computed_examples += batch[0].size(0)
+
+        return batches_input
+
+    def get_layers(self, model: nn.Module):
+        return model.vision_model.encoder.layers
+
+    def update_merged_params_dict(
+        self, merged_params_dict, new_merged_params, layer_idx
+    ):
+        for key, value in new_merged_params.items():
+            key = f"vision_model.encoder.layers.{layer_idx}.{key}"
+            merged_params_dict[key] = value
+
+        return merged_params_dict
+
+    def layer_batches_forward(
+        self, layer: nn.Module, batches_input: List[Tensor]
+    ) -> Tensor:
+        batches_output = []
+        for batch in batches_input:
+            device = next(layer.parameters()).device
+            batch = batch.to(device)
+            logits = (
+                layer(batch, attention_mask=None, causal_attention_mask=None)[0]
+                .detach()
+                .cpu()
+            )
+            batches_output.append(logits)
+        return batches_output

fusion_bench/method/regmean_plusplus/regmean_plusplus.py

@@ -0,0 +1,383 @@
+import logging
+import re
+from collections import defaultdict
+from typing import Dict, List, cast
+
+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.modelpool import BaseModelPool
+
+log = logging.getLogger(__name__)
+
+
+def get_param_names_to_merge(
+    input_param_names: List[str], exclude_param_names_regex: list
+):
+    """
+    get the names of parameters that need to be merged
+    :param input_param_names: list, names of input parameters
+    :param exclude_param_names_regex: list, regular expression of names of parameters that need to be excluded
+    :return:
+    """
+    param_names_to_merge = []
+    for param_name in input_param_names:
+        exclude = any(
+            [
+                re.match(exclude_pattern, param_name)
+                for exclude_pattern in exclude_param_names_regex
+            ]
+        )
+        if not exclude:
+            param_names_to_merge.append(param_name)
+    return param_names_to_merge
+
+
+def get_modules_to_merge(model: nn.Module, include_module_types: list):
+    """
+    get the model modules that need to be merged, whose type is in include_module_types
+    :param model: nn.Module, input model
+    :param include_module_types: list, module types that want to include
+    :return:
+    """
+    modules_to_merge: Dict[str, nn.Module] = {}
+    for module_name, module in model.named_modules():
+        is_valid_type = not include_module_types or any(
+            [
+                isinstance(module, include_module_type)
+                for include_module_type in include_module_types
+            ]
+        )
+        if is_valid_type:
+            modules_to_merge[module_name] = module
+    return modules_to_merge
+
+
+def reduce_non_diagonal_elements(
+    regmean_weights: torch.Tensor, reduce_non_diagonal_ratio: float
+):
+    """
+    reduce the non-diagonal elements in regmean_weights
+    :param regmean_weights: Tensor, shape (hidden_dim, hidden_dim), input regmean weights
+    :param reduce_non_diagonal_ratio: float, reduce non-diagonal elements in regmean weights by multiplying this scalar
+    :return:
+    """
+    # diagonal matrix with (1 - reduce_non_diagonal_ratio) as elements
+    diag_weights = torch.diag(
+        torch.ones(regmean_weights.shape[0]) - reduce_non_diagonal_ratio
+    ).to(regmean_weights.device)
+    # matrix with reduce_non_diagonal_ratio as elements
+    non_diag_weights = torch.zeros_like(diag_weights).fill_(reduce_non_diagonal_ratio)
+    # diagonal elements are unchanged, while non-diagonal elements are multiplied by reduce_non_diagonal_ratio
+    return regmean_weights * (diag_weights + non_diag_weights)
+
+
+def regmean_params_merge(
+    param_weight_list: List[Tensor],
+    param_regmean_list: List[Tensor],
+    reduce_non_diagonal_ratio: float = 1.0,
+    weight_transpose: bool = True,
+    module_name: str = "",
+    device="cpu",
+):
+    # two lists with length num_models_to_merge
+    param_multiplied_results, module_regmean_weights_list = [], []
+    for model_idx, module_regmean_weights in enumerate(param_regmean_list):
+        # reduce non-diagonal elements
+        module_regmean_weights = reduce_non_diagonal_elements(
+            regmean_weights=module_regmean_weights,
+            reduce_non_diagonal_ratio=reduce_non_diagonal_ratio,
+        )
+        module_regmean_weights_list.append(module_regmean_weights)
+
+        model_to_merge_param = param_weight_list[model_idx]
+        # since the weight shape of Linear module is (output_size, input_size), we need to transpose it
+        param_multiplied_results.append(
+            torch.matmul(
+                module_regmean_weights,
+                (
+                    model_to_merge_param.transpose(0, 1)
+                    if weight_transpose
+                    else model_to_merge_param
+                ),
+            )
+        )
+
+    # sum up module_regmean_weights and param_multiplied_results over all individual models
+    sum_module_regmean_weights = sum(module_regmean_weights_list)
+    sum_param_multiplied_results = sum(param_multiplied_results)
+
+    # get the inverse matrix
+    inv_sum_module_regmean_weights = torch.inverse(sum_module_regmean_weights)
+    # merge parameters with regmean
+    merged_param = torch.matmul(
+        inv_sum_module_regmean_weights, sum_param_multiplied_results
+    )
+    # transpose to the original shape of "weight" in Linear module
+    merged_param = merged_param.transpose(0, 1) if weight_transpose else merged_param
+
+    return merged_param
+
+
+def merging_with_regmean_weights(
+    models_to_merge_param_dict: dict,
+    models_to_merge_regmean_weights_list: list,
+    reduce_non_diagonal_ratio: float = 1.0,
+    weight_transpose: bool = True,
+):
+    """
+    merge parameters of different models with computed regmean weights
+    :param models_to_merge_param_dict: dict, dictionary of list, where key is the parameter name,
+    value is a list of the corresponding parameters of all the models that need to be merged
+    :param models_to_merge_regmean_weights_list: list, list of dictionaries with length len(models_to_merge),
+    each dictionary records the regmean weights (matrix) of parameters for each model that needs to be merged, key is module name
+    :param reduce_non_diagonal_ratio: float, reduce non-diagonal elements in regmean weights by multiplying this scalar
+    :return:
+    """
+    # dict, dictionary of model parameters
+    merged_params = {}
+
+    for param_name, param_value_list in models_to_merge_param_dict.items():
+        merged_by_regmean = False
+        # only perform regmean merging on the "weight" parameter of Linear module
+        if param_name.endswith(".weight"):
+            module_name = param_name[: -len(".weight")]
+            if module_name in models_to_merge_regmean_weights_list[0].keys():
+                # two lists with length num_models_to_merge
+                module_regmean_weights_list = []
+                for model_idx, model_to_merge_regmean_weights in enumerate(
+                    models_to_merge_regmean_weights_list
+                ):
+                    device = param_value_list[model_idx].device
+
+                    # Tensor, shape (hidden_dim, hidden_dim)
+                    module_regmean_weights = model_to_merge_regmean_weights[
+                        module_name
+                    ].to(device)
+                    module_regmean_weights_list.append(module_regmean_weights)
+
+                merged_params[param_name] = regmean_params_merge(
+                    param_weight_list=param_value_list,
+                    param_regmean_list=module_regmean_weights_list,
+                    reduce_non_diagonal_ratio=reduce_non_diagonal_ratio,
+                    weight_transpose=weight_transpose,
+                    module_name=module_name,
+                    device=device,
+                )
+
+                merged_by_regmean = True
+        # use average merging for parameters whose names are not end with ".weight" or not in Linear module
+        if not merged_by_regmean:
+            merged_params[param_name] = torch.stack(param_value_list, dim=0).mean(dim=0)
+
+    return merged_params
+
+
+class RegMeanAlgorithmPlusPlus(BaseAlgorithm, SimpleProfilerMixin):
+    _include_module_type = [nn.Linear]
+    _config_mapping = {
+        "num_regmean_examples": "num_regmean_examples",
+        "exclude_param_names_regex": "exclude_param_names_regex",
+        "reduce_non_diagonal_ratio": "reduce_non_diagonal_ratio",
+        "weight_transpose": "weight_transpose",
+    }
+
+    def __init__(
+        self,
+        *,
+        num_regmean_examples: int,
+        exclude_param_names_regex: list,
+        reduce_non_diagonal_ratio: float,
+        weight_transpose: bool,
+        **kwargs,
+    ):
+        self.num_regmean_examples = num_regmean_examples
+        self.exclude_param_names_regex = exclude_param_names_regex
+        self.reduce_non_diagonal_ratio = reduce_non_diagonal_ratio
+        self.weight_transpose = weight_transpose
+        super().__init__(**kwargs)
+
+    def run(self, modelpool: BaseModelPool, **kwargs):
+        if not isinstance(modelpool, BaseModelPool):
+            modelpool = BaseModelPool(modelpool)
+        self.modelpool = modelpool
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        models_to_merge_dict = {
+            name: model.to(device) for name, model in modelpool.named_models()
+        }
+        self.on_regmean_start()
+
+        # initialize the merged models as the pretrained model
+        merged_model = modelpool.load_pretrained_model().to(device)
+        merged_params_dict = {}
+
+        # 1. merge embedding layer
+        merged_embedding_dict = self.merge_embedding_layer(
+            models_to_merge_dict=models_to_merge_dict
+        )
+        merged_model.load_state_dict(merged_embedding_dict, strict=False)
+
+        with torch.no_grad():
+            # 1.1. compute input for the first layer
+            with (
+                self.profile("merging models"),
+                self.profile("computing first layer input"),
+            ):
+                batches_input_dict = defaultdict(list)
+                for name in tqdm(
+                    models_to_merge_dict.keys(), desc="computing input for first layer"
+                ):
+                    dataset = modelpool.load_train_dataset(name)
+
+                    batches_input_dict[name] = self.get_input_for_first_layer(
+                        merged_model, dataset
+                    )
+
+            # 2. iteratively merge layer by layer with regmean algorithm
+            backbone_layers = self.get_layers(merged_model)
+            num_layers = len(backbone_layers)
+
+            models_to_merge_layers_dict = defaultdict(list)
+            for name, model in models_to_merge_dict.items():
+                models_to_merge_layers_dict[name] = self.get_layers(model)
+
+            param_names_to_merge = None
+            for layer_idx, backbone_layer in tqdm(
+                enumerate(backbone_layers), desc="merging layers", total=num_layers
+            ):
+                # dictionary of list, where key is the parameter name,
+                # value is a list of the corresponding parameters of all the models that need to be merged
+                models_to_merge_param_dict = defaultdict(list)
+
+                # list of dictionaries with length len(models_to_merge),
+                # each dictionary records the regmean weights (matrix) of parameters for each model that needs to be merged
+                models_to_merge_regmean_weights_list = []
+
+                for name, layers_to_merge in models_to_merge_layers_dict.items():
+                    layer_to_merge = layers_to_merge[layer_idx]
+                    param_dict = layer_to_merge.state_dict()
+
+                    # exclude parameter whose name matches element in exclude_param_names_regex
+                    if param_names_to_merge is None:
+                        param_names_to_merge = get_param_names_to_merge(
+                            input_param_names=list(param_dict.keys()),
+                            exclude_param_names_regex=self.config.get(
+                                "exclude_param_names_regex", []
+                            ),
+                        )
+
+                    for param_name in param_names_to_merge:
+                        models_to_merge_param_dict[param_name].append(
+                            param_dict[param_name]
+                        )
+
+                    linear_modules_to_merge = get_modules_to_merge(
+                        model=layer_to_merge,
+                        include_module_types=self._include_module_type,
+                    )
+                    assert (
+                        len(linear_modules_to_merge) > 0
+                    ), "No linear modules to merge"
+
+                    # 2.1. compute regmean weights for each model
+                    with (
+                        self.profile("merging models"),
+                        self.profile("computing regmean weights"),
+                    ):
+                        regmean_weights = self.get_regmean_weights(
+                            name,
+                            layer_to_merge,
+                            batches_input=batches_input_dict[name],
+                            linear_modules_to_merge=linear_modules_to_merge,
+                        )
+
+                        module_subset = get_param_names_to_merge(
+                            input_param_names=list(param_dict.keys()),
+                            exclude_param_names_regex=self.exclude_param_names_regex,
+                        )
+                        module_subset = [
+                            name.replace(".weight", "").replace(".bias", "")
+                            for name in module_subset
+                        ]
+                        module_subset = list(set(module_subset))
+                        regmean_weights = {
+                            module_name: regmean_weights[module_name]
+                            for module_name in module_subset
+                            if module_name in regmean_weights
+                        }
+
+                        models_to_merge_regmean_weights_list.append(regmean_weights)
+
+                # 2.2. merge parameters with regmean weights
+                with self.profile("merging models"):
+                    # merging with regmean weights
+                    merged_layer_params = merging_with_regmean_weights(
+                        models_to_merge_param_dict=models_to_merge_param_dict,
+                        models_to_merge_regmean_weights_list=models_to_merge_regmean_weights_list,
+                        reduce_non_diagonal_ratio=self.reduce_non_diagonal_ratio,
+                        weight_transpose=self.config.get("weight_transpose", True),
+                    )
+
+                    merged_params_dict = self.update_merged_params_dict(
+                        merged_params_dict=merged_params_dict,
+                        new_merged_params=merged_layer_params,
+                        layer_idx=layer_idx,
+                    )
+
+                # 2.3. compute input for the next layer
+                with (
+                    self.profile("merging models"),
+                    self.profile("forwarding next layer"),
+                ):
+                    if layer_idx < num_layers - 1:
+                        backbone_layer.load_state_dict(
+                            merged_layer_params, strict=False
+                        )
+                        batches_output_dict = defaultdict(list)
+                        for name in models_to_merge_dict.keys():
+                            batches_output_dict[name] = self.layer_batches_forward(
+                                backbone_layer, batches_input_dict[name]
+                            )
+                        batches_input_dict = batches_output_dict
+
+            # 3. load state dict to the merged model
+            merged_model.load_state_dict(merged_params_dict, strict=False)
+
+        self.print_profile_summary()
+        return merged_model
+
+    def merge_embedding_layer(self, models_to_merge_dict: Dict[str, nn.Module]):
+        """
+        Merge the embedding layer of the model with the merged model.
+        This method should be implemented in subclasses if needed.
+        """
+        raise NotImplementedError()
+
+    def get_input_for_first_layer(self, model: nn.Module, train_dataset):
+        raise NotImplementedError
+
+    def get_layers(self, model: nn.Module):
+        raise NotImplementedError
+
+    def update_merged_params_dict(
+        self, merged_params_dict, new_merged_params, layer_idx
+    ):
+        raise NotImplementedError
+
+    def layer_batches_forward(self, layer: nn.Module, batches_input: List[Tensor]):
+        raise NotImplementedError
+
+    def on_regmean_start(self):
+        pass
+
+    def get_regmean_weights(
+        self,
+        model_name: str,
+        layer: nn.Module,
+        batches_input: List[Tensor],
+        linear_modules_to_merge: Dict[str, nn.Module],
+    ):
+        raise NotImplementedError