fusion-bench 0.2.24__py3-none-any.whl → 0.2.26__py3-none-any.whl

fusion_bench/method/dop/dop.py

@@ -0,0 +1,366 @@
+"""
+Continual Model Merging without Data: Dual Projections for Balancing Stability and Plasticity. NeurIPS, 2025.
+
+
+Example:
+
+fusion_bench \
+    method=dop/dop \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8_model_only \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+"""
+
+import logging
+import os
+import random
+from copy import deepcopy
+from pathlib import Path
+from typing import Dict, 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 torch.autograd import Variable
+from tqdm.auto import tqdm
+from transformers import CLIPVisionModel
+
+from fusion_bench import BaseAlgorithm, BaseModelPool, auto_register_config
+from fusion_bench.method.simple_average import simple_average
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.taskpool import CLIPVisionModelTaskPool
+from fusion_bench.utils import seed_everything_by_time
+from fusion_bench.utils.json import save_to_json
+
+from .min_norm_solvers import MinNormSolver, gradient_normalizers
+from .utils import is_leaf_module, svd
+
+log = logging.getLogger(__name__)
+
+
+@auto_register_config
+class ContinualDOPForCLIP(BaseAlgorithm, LightningFabricMixin):
+
+    def __init__(
+        self,
+        seed: Optional[int] = None,
+        shuffle_order: bool = False,
+        save_on_every_step: bool = True,
+        evaluate_on_every_step: bool = False,
+        lr: float = 1e-4,
+        num_steps: int = 200,
+        mgda: bool = True,
+        ema: bool = True,
+        ema_beta: float = 0.99,
+        alpha: float = None,
+        svd_epsilon: float = 1.0,
+        svd_proj_space: str = "uv",
+        **kwargs,
+    ):
+        self.lr = lr
+        self.num_steps = num_steps
+        self.mgda = mgda
+        self.ema = ema
+        self.ema_beta = ema_beta
+        self.alpha = alpha
+        self.svd_epsilon = svd_epsilon
+        self.svd_proj_space = svd_proj_space
+        self.seed = seed
+        self.shuffle_order = shuffle_order
+        self.save_on_every_step = save_on_every_step
+        self.evaluate_on_every_step = evaluate_on_every_step
+
+        assert (
+            self.svd_epsilon >= 0 and self.svd_epsilon <= 1
+        ), "The svd_epsilon should be in the range of [0, 1]"
+        assert (
+            self.alpha >= 0 and self.alpha <= 1
+        ), "The alpha should be in the range of [0, 1]"
+        super().__init__(**kwargs)
+
+    def print_params(self, pretrained_model):
+        total_params = 0
+        linear_params = 0
+        linear_weight_params = 0
+        for module_name, module in pretrained_model.named_modules():
+            if not is_leaf_module(module):
+                continue
+            if isinstance(module, nn.Linear):
+                linear_params += sum(p.numel() for n, p in module.named_parameters())
+                linear_weight_params += sum(
+                    p.numel() for n, p in module.named_parameters() if "weight" in n
+                )
+            total_params += sum(p.numel() for p in module.parameters())
+
+        linear_ratio = linear_params / total_params * 100
+        linear_weight_ratio = linear_weight_params / total_params * 100
+        print(f"Total Parameters: {total_params}")
+        print(f"Linear Parameters: {linear_params}")
+        print(f"Linear Weight Parameters: {linear_weight_params}")
+        print(f"Linear Ratio: {linear_ratio:.2f}%")
+        print(f"Linear Weight Ratio: {linear_weight_ratio:.2f}%")
+
+    def run(self, modelpool: BaseModelPool):
+        if self.seed is not None:
+            L.seed_everything(self.seed)
+        else:
+            seed_everything_by_time(self.fabric)
+
+        # get the model names, shuffle if needed
+        # the model names will be saved to the log directory as `model_names.json`
+        model_names = modelpool.model_names
+        if self.shuffle_order:
+            random.shuffle(model_names)
+        if self.log_dir is not None:
+            save_to_json(model_names, os.path.join(self.log_dir, "model_names.json"))
+
+        if self.evaluate_on_every_step:
+            """Configuration for the test datasets"""
+            self.taskpool = cast(CLIPVisionModelTaskPool, self._program.taskpool)
+            self._test_datasets = deepcopy(self.taskpool._test_datasets)
+
+        pretrained_model = modelpool.load_pretrained_model()
+
+        merged_model = None
+        for model_idx, model_name in enumerate(model_names):
+            print(
+                f"--------- Optimizing {model_idx + 1}/{len(model_names)}-th with {model_name} ---------"
+            )
+            if model_idx == 0:
+                merged_model = modelpool.load_model(model_names[0])
+            else:
+                merged_model = self._layer_wise_optimize(
+                    model_names=["merged", model_name],
+                    pretrained_model=deepcopy(pretrained_model),
+                    finetuned_models={
+                        "merged": merged_model,
+                        model_name: modelpool.load_model(model_name),
+                    },
+                    model_idx=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 _layer_wise_optimize(
+        self,
+        model_names: List[str],
+        pretrained_model: nn.Module,
+        finetuned_models: Dict[str, nn.Module],
+        model_idx: int,
+    ):
+        time_cost = []
+        for module_name, module in pretrained_model.named_modules():
+            if not is_leaf_module(module):
+                continue
+
+            if isinstance(module, nn.Linear):
+                if module.weight.requires_grad:
+                    import time
+
+                    start_time = time.time()
+                    merged_weight = self._optimize_weight(
+                        module.weight,
+                        {
+                            model_name: finetuned_models[model_name]
+                            .get_submodule(module_name)
+                            .weight
+                            for model_name in model_names
+                        },
+                        module_name,
+                        model_idx,
+                    )
+                    end_time = time.time()
+                    time_cost.append(end_time - start_time)
+                    module.weight.data = merged_weight.data
+                else:
+                    module.weight.data = simple_average(
+                        [
+                            finetuned_models[model_name]
+                            .get_submodule(module_name)
+                            .weight
+                            for model_name in model_names
+                        ]
+                    )
+                if module.bias is not None:
+                    module.bias.data = simple_average(
+                        [
+                            finetuned_models[model_name].get_submodule(module_name).bias
+                            for model_name in model_names
+                        ]
+                    )
+            else:
+                simple_average(
+                    [
+                        finetuned_models[model_name].get_submodule(module_name)
+                        for model_name in model_names
+                    ],
+                    base_module=module,
+                )
+
+        return pretrained_model
+
+    def _optimize_weight(
+        self,
+        pretrained_weight: Tensor,
+        finetuned_weights: Dict[str, Tensor],
+        module_name: str,
+        model_idx: int,
+    ):
+        assert (
+            self.fabric.world_size == 1
+        ), "This algorithm is not currently supported in distributed training"
+
+        pretrained_weight = self.fabric.to_device(pretrained_weight.detach())
+        finetuned_weights = {
+            model_name: self.fabric.to_device(finetuned_weight.detach())
+            for model_name, finetuned_weight in finetuned_weights.items()
+        }
+
+        merged_weight = self.fabric.to_device(
+            nn.Parameter(
+                simple_average(
+                    [
+                        finetuned_weight.detach()
+                        for finetuned_weight in finetuned_weights.values()
+                    ]
+                ),
+                requires_grad=True,
+            )
+        )
+
+        # Compute SVD of the difference between the finetuned and pretrained weights
+        proj_u_dict = {}
+        proj_v_dict = {}
+        proj_s_dict = {}
+        for i, finetuned_weight in enumerate(finetuned_weights.values()):
+            finetuned_tv = finetuned_weight - pretrained_weight
+            u, s, v = svd(finetuned_tv, full_matrices=True)
+            epsilon = 1.0 if self.svd_epsilon > 1.0 else self.svd_epsilon
+            cumsum_ratio = s.cumsum(dim=0) / s.sum()
+            split_rank = torch.searchsorted(cumsum_ratio, epsilon).item()
+            u_main = u[:, :split_rank]
+            v_main = v[:, :split_rank]
+            s_main = s[:split_rank]
+            proj_u_dict[i] = u_main
+            proj_v_dict[i] = v_main
+            proj_s_dict[i] = s_main
+
+        if self.mgda:
+            if self.ema:
+                ema_sol = [self.alpha, 1 - self.alpha]
+            # This is multiple-gradient descent algorithm (MGDA) optimization
+            optimizer = torch.optim.Adam([merged_weight], lr=self.lr)
+            all_losses = [[], []]
+            all_alphas = [[], []]
+            for step_idx in tqdm(
+                range(self.num_steps), desc=f"Optimizing {module_name} weight"
+            ):
+                # Scaling the loss functions based on the algorithm choice
+                loss_data = {}
+                grads = {}
+                for i, finetuned_weight in enumerate(finetuned_weights.values()):
+                    proj_u = proj_u_dict[i]
+                    proj_v = proj_v_dict[i]
+                    proj_s = proj_s_dict[i]
+                    delta_tv = merged_weight - finetuned_weight
+                    loss_i = self.cal_loss_i(delta_tv, proj_s, proj_u, proj_v)
+                    loss_data[i] = float(loss_i.data)
+
+                    all_losses[i].append(float(loss_i.data))
+
+                    optimizer.zero_grad()
+                    loss_i.backward()
+                    grads[i] = Variable(
+                        merged_weight.grad.data.clone(), requires_grad=False
+                    )
+
+                # Normalize all gradients
+                gn = gradient_normalizers(
+                    grads=grads, losses=loss_data, normalization_type="loss"
+                )
+                for i, _ in enumerate(finetuned_weights.values()):
+                    grads[i] = grads[i] / float(gn[i])
+
+                # Frank-Wolfe iteration to compute scales.
+                sol, min_norm = MinNormSolver.find_min_norm_element(
+                    [[grads[i]] for i in range(len(finetuned_weights.values()))]
+                )
+
+                if self.ema:
+                    ema_sol = [
+                        self.ema_beta * ema_sol[i] + (1 - self.ema_beta) * float(sol[i])
+                        for i in range(len(sol))
+                    ]
+                    sol = ema_sol
+                    all_alphas[0].append(ema_sol[0])
+                    all_alphas[1].append(ema_sol[1])
+
+                # Scaled back-propagation
+                loss = 0
+                for i, finetuned_weight in enumerate(finetuned_weights.values()):
+                    # Comptue gradients of each loss function wrt parameters
+                    proj_u = proj_u_dict[i]
+                    proj_v = proj_v_dict[i]
+                    proj_s = proj_s_dict[i]
+                    delta_tv = merged_weight - finetuned_weight
+                    loss_i = self.cal_loss_i(delta_tv, proj_s, proj_u, proj_v)
+                    loss += float(sol[i]) * loss_i
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        else:
+            # This is a naive weighted optimization
+            optimizer = torch.optim.Adam([merged_weight], lr=self.lr)
+            for step_idx in tqdm(
+                range(self.num_steps), desc=f"Optimizing {module_name} weight"
+            ):
+                loss = 0
+                for i, finetuned_weight in enumerate(finetuned_weights.values()):
+                    proj_u = proj_u_dict[i]
+                    proj_v = proj_v_dict[i]
+                    proj_s = proj_s_dict[i]
+                    delta_tv = merged_weight - finetuned_weight
+                    loss_i = self.cal_loss_i(delta_tv, proj_s, proj_u, proj_v)
+                    loss += self.alpha * loss_i if i == 0 else (1 - self.alpha) * loss_i
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+        return merged_weight.detach().cpu()
+
+    def cal_loss_i(self, delta_tv, proj_s, proj_u, proj_v):
+        proj_delta_1 = torch.diag(proj_s) @ proj_u.T @ delta_tv
+        proj_delta_2 = delta_tv @ proj_v @ torch.diag(proj_s)
+        loss_i_u = torch.linalg.matrix_norm(proj_delta_1, ord="fro") ** 2
+        loss_i_v = torch.linalg.matrix_norm(proj_delta_2, ord="fro") ** 2
+        if self.svd_proj_space == "uv":
+            loss_i = loss_i_u + loss_i_v
+        elif self.svd_proj_space == "u":
+            loss_i = loss_i_u
+        elif self.svd_proj_space == "v":
+            loss_i = loss_i_v
+        else:
+            raise ValueError("Invalid svd_proj_space")
+
+        return loss_i
+
+    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}"
+        )

fusion_bench/method/dop/min_norm_solvers.py

@@ -0,0 +1,227 @@
+# This code is from
+# Multi-Task Learning as Multi-Objective Optimization
+# Ozan Sener, Vladlen Koltun
+# Neural Information Processing Systems (NeurIPS) 2018
+# https://github.com/intel-isl/MultiObjectiveOptimization
+from typing import Union
+
+import numpy as np
+import torch
+
+
+def np_sum(x: Union[torch.Tensor, np.ndarray]) -> float:
+    if isinstance(x, torch.Tensor):
+        return x.sum().item()
+    return np.sum(x)
+
+
+def to_numpy(x: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu().numpy()
+    return x
+
+
+class MinNormSolver:
+    MAX_ITER = 250
+    STOP_CRIT = 1e-5
+
+    def _min_norm_element_from2(v1v1, v1v2, v2v2):
+        """
+        Analytical solution for min_{c} |cx_1 + (1-c)x_2|_2^2
+        d is the distance (objective) optimzed
+        v1v1 = <x1,x1>
+        v1v2 = <x1,x2>
+        v2v2 = <x2,x2>
+        """
+        if v1v2 >= v1v1:
+            # Case: Fig 1, third column
+            gamma = 0.999
+            cost = v1v1
+            return gamma, cost
+        if v1v2 >= v2v2:
+            # Case: Fig 1, first column
+            gamma = 0.001
+            cost = v2v2
+            return gamma, cost
+        # Case: Fig 1, second column
+        gamma = -1.0 * ((v1v2 - v2v2) / (v1v1 + v2v2 - 2 * v1v2))
+        cost = v2v2 + gamma * (v1v2 - v2v2)
+        return gamma, cost
+
+    def _min_norm_2d(vecs, dps):
+        R"""
+        Find the minimum norm solution as combination of two points
+        This is correct only in 2D
+        ie. min_c |\sum c_i x_i|_2^2 st. \sum c_i = 1 , 1 >= c_1 >= 0 for all i, c_i + c_j = 1.0 for some i, j
+        """
+        dmin = 1e8
+        for i in range(len(vecs)):
+            for j in range(i + 1, len(vecs)):
+                if (i, j) not in dps:
+                    dps[(i, j)] = 0.0
+                    for k in range(len(vecs[i])):
+                        dps[(i, j)] += (
+                            torch.mul(vecs[i][k], vecs[j][k]).sum().data.cpu()
+                        )
+                    dps[(j, i)] = dps[(i, j)]
+                if (i, i) not in dps:
+                    dps[(i, i)] = 0.0
+                    for k in range(len(vecs[i])):
+                        dps[(i, i)] += (
+                            torch.mul(vecs[i][k], vecs[i][k]).sum().data.cpu()
+                        )
+                if (j, j) not in dps:
+                    dps[(j, j)] = 0.0
+                    for k in range(len(vecs[i])):
+                        dps[(j, j)] += (
+                            torch.mul(vecs[j][k], vecs[j][k]).sum().data.cpu()
+                        )
+                c, d = MinNormSolver._min_norm_element_from2(
+                    dps[(i, i)], dps[(i, j)], dps[(j, j)]
+                )
+                if d < dmin:
+                    dmin = d
+                    sol = [(i, j), c, d]
+        return sol, dps
+
+    def _projection2simplex(y):
+        R"""
+        Given y, it solves argmin_z |y-z|_2 st \sum z = 1 , 1 >= z_i >= 0 for all i
+        """
+        m = len(y)
+        sorted_y = np.flip(np.sort(y), axis=0)
+        tmpsum = 0.0
+        tmax_f = (np.sum(y) - 1.0) / m
+        for i in range(m - 1):
+            tmpsum += sorted_y[i]
+            tmax = (tmpsum - 1) / (i + 1.0)
+            if tmax > sorted_y[i + 1]:
+                tmax_f = tmax
+                break
+        return np.maximum(y - tmax_f, np.zeros(y.shape))
+
+    def _next_point(cur_val, grad, n):
+        proj_grad = grad - (np.sum(grad) / n)
+        tm1 = -1.0 * cur_val[proj_grad < 0] / proj_grad[proj_grad < 0]
+        tm2 = (1.0 - cur_val[proj_grad > 0]) / (proj_grad[proj_grad > 0])
+
+        skippers = np_sum(tm1 < 1e-7) + np_sum(tm2 < 1e-7)
+        t = 1
+        if len(tm1[tm1 > 1e-7]) > 0:
+            t = np.min(to_numpy(tm1[tm1 > 1e-7]))
+        if len(tm2[tm2 > 1e-7]) > 0:
+            t = min(t, np.min(to_numpy(tm2[tm2 > 1e-7])))
+
+        next_point = proj_grad * t + to_numpy(cur_val)
+        next_point = MinNormSolver._projection2simplex(next_point)
+        return next_point
+
+    def find_min_norm_element(vecs):
+        R"""
+        Given a list of vectors (vecs), this method finds the minimum norm element in the convex hull
+        as min |u|_2 st. u = \sum c_i vecs[i] and \sum c_i = 1.
+        It is quite geometric, and the main idea is the fact that if d_{ij} = min |u|_2 st u = c x_i + (1-c) x_j; the solution lies in (0, d_{i,j})
+        Hence, we find the best 2-task solution, and then run the projected gradient descent until convergence
+        """
+        # Solution lying at the combination of two points
+        dps = {}
+        init_sol, dps = MinNormSolver._min_norm_2d(vecs, dps)
+
+        n = len(vecs)
+        sol_vec = np.zeros(n)
+        sol_vec[init_sol[0][0]] = init_sol[1]
+        sol_vec[init_sol[0][1]] = 1 - init_sol[1]
+
+        if n < 3:
+            # This is optimal for n=2, so return the solution
+            return sol_vec, init_sol[2]
+
+        iter_count = 0
+
+        grad_mat = np.zeros((n, n))
+        for i in range(n):
+            for j in range(n):
+                grad_mat[i, j] = dps[(i, j)]
+
+        while iter_count < MinNormSolver.MAX_ITER:
+            grad_dir = -1.0 * np.dot(grad_mat, sol_vec)
+            new_point = MinNormSolver._next_point(sol_vec, grad_dir, n)
+            # Re-compute the inner products for line search
+            v1v1 = 0.0
+            v1v2 = 0.0
+            v2v2 = 0.0
+            for i in range(n):
+                for j in range(n):
+                    v1v1 += sol_vec[i] * sol_vec[j] * dps[(i, j)]
+                    v1v2 += sol_vec[i] * new_point[j] * dps[(i, j)]
+                    v2v2 += new_point[i] * new_point[j] * dps[(i, j)]
+            nc, nd = MinNormSolver._min_norm_element_from2(v1v1, v1v2, v2v2)
+            new_sol_vec = nc * sol_vec + (1 - nc) * new_point
+            change = new_sol_vec - sol_vec
+            if np_sum(np.abs(change)) < MinNormSolver.STOP_CRIT:
+                return sol_vec, nd
+            sol_vec = new_sol_vec
+
+    def find_min_norm_element_FW(vecs):
+        R"""
+        Given a list of vectors (vecs), this method finds the minimum norm element in the convex hull
+        as min |u|_2 st. u = \sum c_i vecs[i] and \sum c_i = 1.
+        It is quite geometric, and the main idea is the fact that if d_{ij} = min |u|_2 st u = c x_i + (1-c) x_j; the solution lies in (0, d_{i,j})
+        Hence, we find the best 2-task solution, and then run the Frank Wolfe until convergence
+        """
+        # Solution lying at the combination of two points
+        dps = {}
+        init_sol, dps = MinNormSolver._min_norm_2d(vecs, dps)
+
+        n = len(vecs)
+        sol_vec = np.zeros(n)
+        sol_vec[init_sol[0][0]] = init_sol[1]
+        sol_vec[init_sol[0][1]] = 1 - init_sol[1]
+
+        if n < 3:
+            # This is optimal for n=2, so return the solution
+            return sol_vec, init_sol[2]
+
+        iter_count = 0
+
+        grad_mat = np.zeros((n, n))
+        for i in range(n):
+            for j in range(n):
+                grad_mat[i, j] = dps[(i, j)]
+
+        while iter_count < MinNormSolver.MAX_ITER:
+            t_iter = np.argmin(np.dot(grad_mat, sol_vec))
+
+            v1v1 = np.dot(sol_vec, np.dot(grad_mat, sol_vec))
+            v1v2 = np.dot(sol_vec, grad_mat[:, t_iter])
+            v2v2 = grad_mat[t_iter, t_iter]
+
+            nc, nd = MinNormSolver._min_norm_element_from2(v1v1, v1v2, v2v2)
+            new_sol_vec = nc * sol_vec
+            new_sol_vec[t_iter] += 1 - nc
+
+            change = new_sol_vec - sol_vec
+            if np_sum(np.abs(change)) < MinNormSolver.STOP_CRIT:
+                return sol_vec, nd
+            sol_vec = new_sol_vec
+
+
+def gradient_normalizers(grads, losses, normalization_type):
+    gn = {}
+    if normalization_type == "l2":
+        for t in grads:
+            gn[t] = np.sqrt(np.sum([gr.pow(2).sum().data.cpu() for gr in grads[t]]))
+    elif normalization_type == "loss":
+        for t in grads:
+            gn[t] = losses[t]
+    elif normalization_type == "loss+":
+        for t in grads:
+            gn[t] = losses[t] * np.sqrt(
+                np.sum([gr.pow(2).sum().data.cpu() for gr in grads[t]])
+            )
+    elif normalization_type == "none":
+        for t in grads:
+            gn[t] = 1.0
+    else:
+        print("ERROR: Invalid Normalization Type")
+    return gn

fusion_bench/method/dop/utils.py

@@ -0,0 +1,73 @@
+from typing import Tuple
+
+import torch
+from torch import Tensor, nn
+
+from fusion_bench.utils.parameters import state_dict_to_vector
+from fusion_bench.utils.state_dict_arithmetic import state_dict_sub
+
+
+def _svd(w: Tensor, full_matrices=True) -> Tuple[Tensor, Tensor, Tensor]:
+    """
+    Perform Singular Value Decomposition (SVD) on a tensor.
+
+    Args:
+        w (Tensor): The input tensor.
+        full_matrices (bool): Whether to compute the full-sized U and V matrices.
+
+    Returns:
+        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
+    """
+    u, s, vh = torch.linalg.svd(
+        w, full_matrices=full_matrices, driver="gesvd" if w.is_cuda else None
+    )
+    v = vh.T
+    return u, s, v
+
+
+def svd(
+    w: Tensor, full_matrices=True, accelerator=None
+) -> Tuple[Tensor, Tensor, Tensor]:
+    """
+    Perform SVD on a tensor, optionally using a specified accelerator.
+
+    Args:
+        w (Tensor): The input tensor.
+        full_matrices (bool): Whether to compute the full-sized U and V matrices.
+        accelerator (str): The device to perform the computation on.
+
+    Returns:
+        Tuple[Tensor, Tensor, Tensor]: The U, S, and V matrices from SVD.
+    """
+    if accelerator is None:
+        return _svd(w, full_matrices=full_matrices)
+    original_device = w.device
+    w = w.to(accelerator)
+    u, s, v = _svd(w)
+    return u.to(original_device), s.to(original_device), v.to(original_device)
+
+
+def frobenius_inner_product(w1: Tensor, w2: Tensor) -> Tensor:
+    return torch.trace(w1.T @ w2)
+
+
+def is_leaf_module(module: nn.Module) -> bool:
+    return len(list(module.children())) == 0
+
+
+def get_task_vector_norm(model: nn.Module, pretrained_model: nn.Module) -> Tensor:
+    """
+    Get the vector norm of the task model.
+
+    Args:
+        model (nn.Module): The task model.
+        pretrained_model (nn.Module): The pretrained model.
+
+    Returns:
+        Tensor: The vector norm of the task model.
+    """
+    return torch.linalg.norm(
+        state_dict_to_vector(
+            state_dict_sub(model.state_dict(), pretrained_model.state_dict())
+        )
+    )

fusion_bench/method/opcm/opcm.py

@@ -87,6 +87,7 @@ class OPCMForCLIP(
         # get the average model
         with self.profile("loading model"):
             merged_model = modelpool.load_model(model_names[0])
+            assert merged_model is not None, "Failed to load the first model"
 
         if self.evaluate_on_every_step:
             with self.profile("evaluating model"):

fusion_bench/method/pwe_moe/module.py

@@ -13,8 +13,6 @@ import torch.func
 from torch import Tensor, nn
 from torch.nn import functional as F
 
-from fusion_bench.utils import join_list
-
 log = logging.getLogger(__name__)
 
 

fusion_bench/method/tall_mask/task_arithmetic.py

@@ -15,7 +15,7 @@ from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_add,
     state_dict_binary_mask,
     state_dict_diff_abs,
-    state_dict_hadmard_product,
+    state_dict_hadamard_product,
     state_dict_mul,
     state_dict_sub,
     state_dict_sum,
@@ -111,7 +111,7 @@ class TallMaskTaskArithmeticAlgorithm(
 
         with self.profile("compress and retrieve"):
             for model_name in modelpool.model_names:
-                retrieved_task_vector = state_dict_hadmard_product(
+                retrieved_task_vector = state_dict_hadamard_product(
                     tall_masks[model_name], multi_task_vector
                 )
                 retrieved_state_dict = state_dict_add(