fusion-bench 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

fusion_bench/method/isotropic_merging/iso_utils.py

@@ -0,0 +1,176 @@
+import math
+from typing import List
+
+import torch
+
+from fusion_bench.utils import timeit_context
+from fusion_bench.utils.type import StateDictType
+
+
+def iso_c(
+    task_vectors: List[StateDictType],
+    accelerator="cuda",
+    exclude_keys: List[str] = None,
+) -> StateDictType:
+    exclude_keys = [] if exclude_keys is None else exclude_keys
+
+    with torch.no_grad(), timeit_context("ISO-C Merging"):
+        new_vector = {}
+        for key in task_vectors[0]:
+            print(f"Merging {key}...")
+            original_device = task_vectors[0][key].device
+            tvs = [
+                task_vector[key].to(device=accelerator, non_blocking=True)
+                for task_vector in task_vectors
+            ]
+            num_tvs = len(tvs)
+            new_vector[key] = sum(tvs) / num_tvs
+            del tvs  # free memory
+
+            if len(task_vectors[0][key].shape) == 2 and key not in exclude_keys:
+                # if the key is a 2D matrix, we need to merge the task vectors in the common space
+                new_vector[key] *= num_tvs
+                U, S, V = torch.linalg.svd(new_vector[key], full_matrices=False)
+                S_mean = torch.ones_like(S) * S.mean()
+
+                new_vector[key] = torch.linalg.multi_dot(
+                    (
+                        U,
+                        torch.diag(S_mean),
+                        V,
+                    )
+                )
+            new_vector[key] = new_vector[key].to(
+                device=original_device, non_blocking=True
+            )
+    return new_vector
+
+
+@torch.no_grad()
+def iso_cts(
+    task_vectors: List[StateDictType],
+    common_space_fraction: float,
+    accelerator: str = "cuda",
+    exclude_keys: List[str] = None,
+):
+    exclude_keys = [] if exclude_keys is None else exclude_keys
+    new_vector = {}
+
+    print("ISO-CTS Merging")
+    for key in task_vectors[0]:
+        shape_ = task_vectors[0][key].shape
+        original_device = task_vectors[0][key].device
+        is_2d_matrix = (len(shape_) == 2) and (key not in exclude_keys)
+        if not is_2d_matrix:
+            print(f"Combining by avg {key}...")
+            for i, task_vector in enumerate(task_vectors):
+                vec = task_vector[key].to(device=accelerator, non_blocking=True)
+                if i == 0:
+                    new_vector[key] = vec.clone()
+                else:
+                    new_vector[key] += (vec - new_vector[key]) / (i + 1)
+
+            # move the new vector to the original device
+            new_vector[key] = new_vector[key].to(
+                device=original_device, non_blocking=True
+            )
+            continue
+
+        print(f"Computing common space using sum for {key}...")
+        combined_w = sum(
+            [
+                task_vector[key].to(device=accelerator, non_blocking=True)
+                for task_vector in task_vectors
+            ]
+        )
+
+        ### Calculate the common space size (making sure that task specific space is equally divisible) ###
+        common_space_index_s = int(min(shape_) * common_space_fraction)
+        _task_specific_total_space_index_s = round(
+            (min(shape_) - common_space_index_s) / len(task_vectors)
+        ) * len(task_vectors)
+        common_space_index_s = min(shape_) - _task_specific_total_space_index_s
+
+        u, s, v = torch.linalg.svd(combined_w, full_matrices=False)
+        common_space_u = u[:, :common_space_index_s]
+        common_space_s = s[:common_space_index_s]
+        common_space_v = v[:common_space_index_s, :]
+        ###################################################################
+
+        ### Calculate task specific space ###
+        n_dims_per_task = int((min(shape_) - common_space_index_s) / len(task_vectors))
+        for i, task_vector in enumerate(task_vectors):
+            w = task_vector[key].to(device=accelerator)
+
+            # calculate the projection onto task specific space to remove the common space
+            w_ts = w - common_space_u @ common_space_u.T @ w
+            u_ts, s_ts, v_ts = torch.linalg.svd(w_ts, full_matrices=False)
+
+            if i == 0:
+                combined_space_u = torch.zeros_like(u_ts, device=accelerator)
+                combined_space_s = torch.zeros_like(s_ts, device=accelerator)
+                combined_space_v = torch.zeros_like(v_ts, device=accelerator)
+
+            combined_space_u[:, i * n_dims_per_task : (i + 1) * n_dims_per_task] = u_ts[
+                :, :n_dims_per_task
+            ]
+            combined_space_s[i * n_dims_per_task : (i + 1) * n_dims_per_task] = s_ts[
+                :n_dims_per_task
+            ]
+            combined_space_v[i * n_dims_per_task : (i + 1) * n_dims_per_task, :] = v_ts[
+                :n_dims_per_task, :
+            ]
+        ###################################################################
+
+        combined_space_u[
+            :,
+            len(task_vectors) * n_dims_per_task : len(task_vectors) * n_dims_per_task
+            + common_space_index_s,
+        ] = common_space_u
+        combined_space_s[
+            len(task_vectors) * n_dims_per_task : len(task_vectors) * n_dims_per_task
+            + common_space_index_s
+        ] = common_space_s
+        combined_space_v[
+            len(task_vectors) * n_dims_per_task : len(task_vectors) * n_dims_per_task
+            + common_space_index_s,
+            :,
+        ] = common_space_v
+
+        ### Orthogonalize combined_space_u and combined_space_v ###
+        u_combined_space_u, s_combined_space_u, v_combined_space_u = torch.linalg.svd(
+            combined_space_u, full_matrices=False
+        )
+        u_combined_space_v, s_combined_space_v, v_combined_space_v = torch.linalg.svd(
+            combined_space_v, full_matrices=False
+        )
+        combined_space_u = u_combined_space_u @ v_combined_space_u
+        combined_space_v = u_combined_space_v @ v_combined_space_v
+        ###################################################################
+
+        combined_space_s = torch.ones_like(combined_space_s) * combined_space_s.mean()
+
+        new_vector[key] = torch.linalg.multi_dot(
+            (
+                combined_space_u,
+                torch.diag(combined_space_s),
+                combined_space_v,
+            )
+        )
+        new_vector[key] = new_vector[key].to(device=original_device, non_blocking=True)
+
+    return new_vector
+
+
+def check_parameterNamesMatch(checkpoints):
+    parameter_names = set(checkpoints[0].keys())
+
+    if len(checkpoints) >= 2:
+        # raise ValueError("Number of models is less than 2.")
+        for checkpoint in checkpoints[1:]:
+            current_parameterNames = set(checkpoint.keys())
+            if current_parameterNames != parameter_names:
+                raise ValueError(
+                    "Differing parameter names in models. "
+                    f"The different parameters are {parameter_names.symmetric_difference(current_parameterNames)}"
+                )

fusion_bench/method/opcm/__init__.py

@@ -0,0 +1,4 @@
+from .opcm import OPCMForCLIP
+from .task_arithmetic import ContinualTaskArithmeticForCLIP
+from .ties_merging import ContinualTiesMergingForCLIP
+from .weight_average import ContinualWeightAverageForCLIP

fusion_bench/method/opcm/opcm.py

@@ -0,0 +1,277 @@
+import os
+import random
+import time
+from collections import defaultdict
+from copy import deepcopy
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, cast
+
+import lightning as L
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch import Tensor, nn
+from tqdm.auto import tqdm
+from transformers import CLIPVisionModel
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.taskpool import CLIPVisionModelTaskPool
+from fusion_bench.utils import instantiate
+from fusion_bench.utils.json import load_from_json, save_to_json
+from fusion_bench.utils.parameters import state_dict_to_vector
+from fusion_bench.utils.state_dict_arithmetic import state_dict_sub
+
+from .utils import frobenius_inner_product, get_task_vector_norm, is_leaf_module, svd
+
+if TYPE_CHECKING:
+    from torch.utils.tensorboard import SummaryWriter
+
+
+class OPCMForCLIP(
+    BaseAlgorithm,
+    LightningFabricMixin,
+):
+    def __init__(
+        self,
+        alpha: float,
+        shuffle_order: bool = True,
+        seed: Optional[int] = None,
+        save_on_every_step: bool = True,
+        evaluate_on_every_step: bool = False,
+        **kwargs,
+    ):
+        """
+        Continual Model Merging via SVD Projection.
+
+        Args:
+            alpha (float): the scaling factor for the SVD projection.
+            shuffle_order (bool): whether to shuffle the order of the models.
+            seed (Optional[int]): the seed to use.
+            save_on_every_step (bool): whether to save the merged model on every step.
+            evaluate_on_every_step (bool): whether to evaluate the merged model on every step.
+        """
+        self.alpha = alpha
+        self.shuffle_order = shuffle_order
+        self.seed = seed
+        self.save_on_every_step = save_on_every_step
+        self.evaluate_on_every_step = evaluate_on_every_step
+        super().__init__(**kwargs)
+
+    @torch.no_grad()
+    def run(self, modelpool: BaseModelPool):
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+        accelerator = self.fabric.device
+
+        pretrained_model = modelpool.load_pretrained_model()
+
+        model_names = modelpool.model_names
+        if self.shuffle_order:
+            random.shuffle(model_names)
+
+        self.taskpool = cast(CLIPVisionModelTaskPool, self._program.taskpool)
+        self._test_datasets = deepcopy(self.taskpool._test_datasets)
+        """Configuration for the test datasets"""
+
+        # log the model names
+        if self.log_dir is not None:
+            save_to_json(model_names, Path(self.log_dir) / "model_names.json")
+            tensorboard_summarywriter: "SummaryWriter" = self.tensorboard_summarywriter
+            tensorboard_summarywriter.add_text(
+                "global/model_names", str(model_names), global_step=0
+            )
+
+        # get the average model
+        merged_model = modelpool.load_model(model_names[0])
+
+        if self.evaluate_on_every_step:
+            self.taskpool._is_setup = False
+            self.taskpool._test_datasets = DictConfig(
+                {model_names[0]: self._test_datasets[model_names[0]]}
+            )
+            report = self.taskpool.evaluate(deepcopy(merged_model))
+            save_to_json(report, Path(self.log_dir) / "report_0.json")
+
+        self.avg_task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
+        self.all_task_vector_norm = [self.avg_task_vector_norm]
+        self.fabric.log("model/task_vector_norm", self.avg_task_vector_norm, step=0)
+        self.fabric.log("model/avg_task_vector_norm", self.avg_task_vector_norm, step=0)
+        self.fabric.log(
+            "model/merged_task_vector_norm", self.avg_task_vector_norm, step=0
+        )
+
+        self.previous_lambda_t = 1
+        self.lambda_t = None
+        self.fabric.log("model/lambda_t", self.previous_lambda_t, step=0)
+        self.fabric.log("empirical/lambda_t", 1, step=0)
+
+        if self.save_on_every_step:
+            self.save_merged_model(merged_model, 0)
+
+        for model_idx, model_name in tqdm(
+            enumerate(model_names[1:]), desc="Processing models"
+        ):
+            model_idx += 1
+            task_model = modelpool.load_model(model_name)
+
+            self.all_task_vector_norm.append(
+                get_task_vector_norm(task_model, pretrained_model)
+            )
+            self.avg_task_vector_norm = np.mean(self.all_task_vector_norm)
+            self.fabric.log(
+                "model/task_vector_norm", self.all_task_vector_norm[-1], step=model_idx
+            )
+            self.fabric.log(
+                "model/avg_task_vector_norm", self.avg_task_vector_norm, step=model_idx
+            )
+
+            self.lambda_t = 1  # temporary value
+
+            for module_name, module in tqdm(
+                list(merged_model.named_modules()),
+                desc=f"Processing {model_name}",
+                leave=False,
+            ):
+                if not is_leaf_module(module):
+                    continue
+
+                if isinstance(module, nn.Linear):
+                    module.weight.data = self.merge_linear_weights(
+                        module.weight,
+                        pretrained_model.get_submodule(module_name).weight,
+                        task_model.get_submodule(module_name).weight,
+                        param_name=".".join([module_name, "weight"]),
+                        alpha=self.alpha,
+                        accelerator=accelerator,
+                    )
+                    if module.bias is not None:
+                        module.bias.data = self.merge_other_parameters(
+                            module.bias,
+                            pretrained_model.get_submodule(module_name).bias,
+                            task_model.get_submodule(module_name).bias,
+                            param_name=".".join([module_name, "bias"]),
+                            accelerator=accelerator,
+                        )
+                else:
+                    for param_name, param in module.named_parameters():
+                        param.data = self.merge_other_parameters(
+                            merged_W=param,
+                            pretrained_W=pretrained_model.get_submodule(
+                                module_name
+                            ).get_parameter(param_name),
+                            task_W=task_model.get_submodule(module_name).get_parameter(
+                                param_name
+                            ),
+                            param_name=".".join([module_name, param_name]),
+                            accelerator=accelerator,
+                        )
+
+            task_vector_norm = get_task_vector_norm(merged_model, pretrained_model)
+            self.lambda_t *= task_vector_norm / self.avg_task_vector_norm
+            for param_name, param in merged_model.named_parameters():
+                param.data = pretrained_model.get_parameter(param_name) + (
+                    param - pretrained_model.get_parameter(param_name)
+                ) * (self.avg_task_vector_norm / task_vector_norm)
+            self.fabric.log("model/lambda_t", self.lambda_t, step=model_idx)
+            self.fabric.log(
+                "empirical/lambda_t", np.sqrt(model_idx + 1), step=model_idx
+            )
+            self.previous_lambda_t = self.lambda_t
+            self.lambda_t = None
+
+            self.fabric.log(
+                "model/merged_task_vector_norm",
+                get_task_vector_norm(merged_model, pretrained_model),
+                step=model_idx,
+            )
+
+            if self.save_on_every_step:
+                self.save_merged_model(merged_model, model_idx)
+
+            if self.evaluate_on_every_step:
+                self.taskpool._is_setup = False
+                self.taskpool._test_datasets = DictConfig(
+                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
+                )
+                report = self.taskpool.evaluate(deepcopy(merged_model))
+                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
+
+        return merged_model
+
+    def save_merged_model(self, merged_model: CLIPVisionModel, step: int):
+        os.makedirs(Path(self.log_dir) / "checkpoints", exist_ok=True)
+        merged_model.save_pretrained(
+            Path(self.log_dir) / "checkpoints" / f"merged_model_{step}"
+        )
+
+    def merge_linear_weights(
+        self,
+        merged_W: Tensor,
+        pretrained_W: Tensor,
+        task_W: Tensor,
+        param_name: str,
+        alpha: float,
+        accelerator: str = "cpu",
+    ):
+        original_device = merged_W.device
+        merged_W = merged_W.to(accelerator)
+        pretrained_W = pretrained_W.to(accelerator)
+        task_W = task_W.to(accelerator)
+
+        previous_merged_tv = merged_W - pretrained_W
+        task_tv = task_W - pretrained_W
+
+        u, s, v = svd(previous_merged_tv)
+        rank = s.size(0)
+        split_rank = (s.cumsum(dim=0) / s.sum() > alpha).float().argmax().item()
+
+        projected_task_tv = u.T @ task_tv @ v
+        projected_task_tv.diag().fill_(0)
+
+        projected_task_tv[:split_rank, :split_rank] = 0
+
+        cleaned_task_tv = u @ projected_task_tv @ v.T
+
+        previous_lambda_t = self.previous_lambda_t
+        lambda_t = self.lambda_t
+        new_merged_W = (
+            pretrained_W
+            + (previous_lambda_t * previous_merged_tv + cleaned_task_tv) / lambda_t
+        )
+        return new_merged_W.to(original_device)
+
+    def merge_other_parameters(
+        self,
+        merged_W: Tensor,
+        pretrained_W: Tensor,
+        task_W: Tensor,
+        param_name: str,
+        accelerator: str = "cpu",
+    ):
+        original_device = merged_W.device
+        merged_W = merged_W.to(accelerator)
+        pretrained_W = pretrained_W.to(accelerator)
+        task_W = task_W.to(accelerator)
+
+        previous_merged_tv = merged_W - pretrained_W
+        task_tv = task_W - pretrained_W
+
+        previous_lambda_t = self.previous_lambda_t
+        lambda_t = self.lambda_t
+
+        new_merged_W = (
+            pretrained_W + (previous_lambda_t * previous_merged_tv + task_tv) / lambda_t
+        )
+        return new_merged_W.to(original_device)
+
+    def compute_lambda_t(
+        self, previous_merged_tv: Tensor, task_tv: Tensor, previous_lambda_t: float
+    ):
+        previous_merged_tv = torch.flatten(previous_merged_tv)
+        task_tv = torch.flatten(task_tv)
+
+        lambda_t = torch.linalg.vector_norm(
+            previous_lambda_t * previous_merged_tv + task_tv
+        ) / torch.linalg.vector_norm(previous_merged_tv)
+        return lambda_t.item()

fusion_bench/method/opcm/task_arithmetic.py

@@ -0,0 +1,115 @@
+import os
+import random
+import time
+from collections import defaultdict
+from copy import deepcopy
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, cast
+
+import lightning as L
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch import Tensor, nn
+from tqdm.auto import tqdm
+from transformers import CLIPVisionModel
+
+from fusion_bench import BaseAlgorithm, BaseModelPool
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.taskpool import CLIPVisionModelTaskPool
+from fusion_bench.utils.json import load_from_json, save_to_json
+from fusion_bench.utils.state_dict_arithmetic import state_dict_add, state_dict_sub
+
+if TYPE_CHECKING:
+    from torch.utils.tensorboard import SummaryWriter
+
+
+class ContinualTaskArithmeticForCLIP(BaseAlgorithm, LightningFabricMixin):
+    def __init__(
+        self,
+        scaling_factor: float,
+        shuffle_order: bool = True,
+        seed: Optional[int] = None,
+        save_on_every_step: bool = True,
+        evaluate_on_every_step: bool = False,
+        **kwargs,
+    ):
+        """
+        Continual Model Merging via Task Arithmetic.
+
+        Args:
+            scaling_factor (float): the scaling factor to use.
+            shuffle_order (bool): whether to shuffle the order of the models.
+            seed (Optional[int]): the seed to use.
+            save_on_every_step (bool): whether to save the merged model on every step.
+            evaluate_on_every_step (bool): whether to evaluate the merged model on every step.
+        """
+        self.scaling_factor = scaling_factor
+        self.shuffle_order = shuffle_order
+        self.seed = seed
+        self.save_on_every_step = save_on_every_step
+        self.evaluate_on_every_step = evaluate_on_every_step
+        super().__init__(**kwargs)
+
+    @torch.no_grad()
+    def run(self, modelpool: BaseModelPool):
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+
+        model_names = modelpool.model_names
+        if self.shuffle_order:
+            random.shuffle(model_names)
+
+        self.taskpool = cast(CLIPVisionModelTaskPool, self._program.taskpool)
+        self._test_datasets = deepcopy(self.taskpool._test_datasets)
+        """Configuration for the test datasets"""
+
+        # log the model names
+        if self.log_dir is not None:
+            save_to_json(model_names, Path(self.log_dir) / "model_names.json")
+            tensorboard_summarywriter: "SummaryWriter" = self.tensorboard_summarywriter
+            tensorboard_summarywriter.add_text(
+                "global/model_names", str(model_names), global_step=0
+            )
+
+        # get the average model
+        pretrained_model = modelpool.load_pretrained_model()
+        merged_model = deepcopy(pretrained_model)
+
+        for model_idx, model_name in tqdm(
+            enumerate(model_names), desc="Processing models"
+        ):
+            task_model = modelpool.load_model(model_name)
+
+            for param_name, param in task_model.named_parameters():
+                if not param.requires_grad:
+                    continue
+
+                task_param = param
+                merged_param = merged_model.get_parameter(param_name)
+                pretrained_param = pretrained_model.get_parameter(param_name)
+
+                new_param = merged_param + self.scaling_factor * (
+                    task_param - pretrained_param
+                )
+                merged_model.get_parameter(param_name).data = new_param
+
+            if self.save_on_every_step:
+                self.save_merged_model(merged_model, model_idx)
+
+            if self.evaluate_on_every_step:
+                self.taskpool._is_setup = False
+                self.taskpool._test_datasets = DictConfig(
+                    {n: self._test_datasets[n] for n in model_names[: model_idx + 1]}
+                )
+                report = self.taskpool.evaluate(deepcopy(merged_model))
+                save_to_json(report, Path(self.log_dir) / f"report_{model_idx}.json")
+
+        return merged_model
+
+    def save_merged_model(self, merged_model: CLIPVisionModel, step: int):
+        os.makedirs(Path(self.log_dir) / "checkpoints", exist_ok=True)
+        torch.save(
+            merged_model.state_dict(),
+            Path(self.log_dir) / "checkpoints" / f"model_{step}.pth",
+        )