fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/method/pwe_moe/phn/solvers.py

@@ -0,0 +1,195 @@
+from abc import abstractmethod
+from typing import Tuple
+
+import cvxopt
+import cvxpy as cp
+import numpy as np
+import torch
+from torch import Tensor
+
+"""Implementation of Pareto HyperNetworks with:
+1. Linear scalarization
+3. EPO
+
+
+EPO code from: https://github.com/dbmptr/EPOSearch
+"""
+
+
+class Solver:
+    def __init__(self, n_tasks):
+        super().__init__()
+        self.n_tasks = n_tasks
+
+    @abstractmethod
+    def get_weighted_loss(self, losses, ray, parameters=None, **kwargs):
+        pass
+
+    def __call__(self, losses, ray, parameters, **kwargs):
+        return self.get_weighted_loss(losses, ray, parameters, **kwargs)
+
+
+class LinearScalarizationSolver(Solver):
+    """For LS we use the preference ray to weigh the losses"""
+
+    def __init__(self, n_tasks):
+        super().__init__(n_tasks)
+
+    def get_weighted_loss(self, losses, ray, parameters=None, **kwargs):
+        return (losses * ray).sum()
+
+
+class EPOSolver(Solver):
+    """Wrapper over EPO"""
+
+    def __init__(self, n_tasks, n_params):
+        super().__init__(n_tasks)
+        self.solver = EPO(n_tasks=n_tasks, n_params=n_params)
+
+    def get_weighted_loss(
+        self, losses, ray, parameters: Tuple[Tensor] = None, **kwargs
+    ):
+        assert parameters is not None
+        return self.solver.get_weighted_loss(losses, ray, parameters)
+
+
+class EPO:
+    def __init__(self, n_tasks, n_params):
+        self.n_tasks = n_tasks
+        self.n_params = n_params
+
+    def __call__(self, losses, ray, parameters):
+        return self.get_weighted_loss(losses, ray, parameters)
+
+    @staticmethod
+    def _flattening(grad: Tuple[Tensor]) -> Tensor:
+        return torch.cat(tuple(g.flatten() for g in grad), dim=0)
+
+    def get_weighted_loss(
+        self,
+        losses: Tuple[Tensor],
+        ray: Tensor,
+        parameters: Tuple[Tensor],
+    ):
+        lp = ExactParetoLP(m=self.n_tasks, n=self.n_params, r=ray.cpu().numpy())
+
+        grads = []
+        for i, loss in enumerate(losses):
+            g = torch.autograd.grad(
+                loss, parameters, retain_graph=True, allow_unused=True
+            )
+            g = tuple(filter(lambda x: x is not None, g))
+            flat_grad = self._flattening(g)
+            grads.append(flat_grad.data)
+
+        G = torch.stack(grads)
+        GG_T = G @ G.T
+        GG_T = GG_T.detach().cpu().numpy()
+
+        numpy_losses = losses.detach().cpu().numpy()
+
+        try:
+            alpha = lp.get_alpha(numpy_losses, G=GG_T, C=True)
+        except Exception as excep:
+            print(excep)
+            alpha = None
+
+        if alpha is None:  # A patch for the issue in cvxpy
+            alpha = (ray / ray.sum()).cpu().numpy()
+
+        alpha *= self.n_tasks
+        alpha = torch.from_numpy(alpha).to(losses.device)
+
+        weighted_loss = torch.sum(losses * alpha)
+        return weighted_loss
+
+
+class ExactParetoLP(object):
+    """modifications of the code in https://github.com/dbmptr/EPOSearch"""
+
+    def __init__(self, m, n, r, eps=1e-4):
+        cvxopt.glpk.options["msg_lev"] = "GLP_MSG_OFF"
+        self.m = m
+        self.n = n
+        self.r = r
+        self.eps = eps
+        self.last_move = None
+        self.a = cp.Parameter(m)  # Adjustments
+        self.C = cp.Parameter((m, m))  # C: Gradient inner products, G^T G
+        self.Ca = cp.Parameter(m)  # d_bal^TG
+        self.rhs = cp.Parameter(m)  # RHS of constraints for balancing
+
+        self.alpha = cp.Variable(m)  # Variable to optimize
+
+        obj_bal = cp.Maximize(self.alpha @ self.Ca)  # objective for balance
+        constraints_bal = [
+            self.alpha >= 0,
+            cp.sum(self.alpha) == 1,  # Simplex
+            self.C @ self.alpha >= self.rhs,
+        ]
+        self.prob_bal = cp.Problem(obj_bal, constraints_bal)  # LP balance
+
+        obj_dom = cp.Maximize(cp.sum(self.alpha @ self.C))  # obj for descent
+        constraints_res = [
+            self.alpha >= 0,
+            cp.sum(self.alpha) == 1,  # Restrict
+            self.alpha @ self.Ca >= -cp.neg(cp.max(self.Ca)),
+            self.C @ self.alpha >= 0,
+        ]
+        constraints_rel = [
+            self.alpha >= 0,
+            cp.sum(self.alpha) == 1,  # Relaxed
+            self.C @ self.alpha >= 0,
+        ]
+        self.prob_dom = cp.Problem(obj_dom, constraints_res)  # LP dominance
+        self.prob_rel = cp.Problem(obj_dom, constraints_rel)  # LP dominance
+
+        self.gamma = 0  # Stores the latest Optimum value of the LP problem
+        self.mu_rl = 0  # Stores the latest non-uniformity
+
+    def get_alpha(self, loss_terms, G: Tensor, r=None, C: bool = False, relax=False):
+        r = self.r if r is None else r
+        assert len(loss_terms) == len(G) == len(r) == self.m, "length != m"
+        rl, self.mu_rl, self.a.value = adjustments(loss_terms, r)
+        self.C.value = G if C else G @ G.T
+        self.Ca.value = self.C.value @ self.a.value
+
+        if self.mu_rl > self.eps:
+            J = self.Ca.value > 0
+            if len(np.where(J)[0]) > 0:
+                J_star_idx = np.where(rl == np.max(rl))[0]
+                self.rhs.value = self.Ca.value.copy()
+                self.rhs.value[J] = -np.inf  # Not efficient; but works.
+                self.rhs.value[J_star_idx] = 0
+            else:
+                self.rhs.value = np.zeros_like(self.Ca.value)
+            self.gamma = self.prob_bal.solve(solver=cp.GLPK, verbose=False)
+            self.last_move = "bal"
+        else:
+            if relax:
+                self.gamma = self.prob_rel.solve(solver=cp.GLPK, verbose=False)
+            else:
+                self.gamma = self.prob_dom.solve(solver=cp.GLPK, verbose=False)
+            self.last_move = "dom"
+
+        return self.alpha.value
+
+
+def mu(rl, normed=False):
+    if len(np.where(rl < 0)[0]):
+        raise ValueError(f"rl<0 \n rl={rl}")
+        # return None
+    m = len(rl)
+    l_hat = rl if normed else rl / rl.sum()
+    eps = np.finfo(rl.dtype).eps
+    l_hat = l_hat[l_hat > eps]
+    return np.sum(l_hat * np.log(l_hat * m))
+
+
+def adjustments(l, r=1):
+    m = len(l)
+    rl = r * l
+    l_hat = rl / rl.sum()
+    mu_rl = mu(l_hat, normed=True)
+    a = r * (np.log(l_hat * m) - mu_rl)
+    return rl, mu_rl, a

fusion_bench/method/pwe_moe/utils.py

@@ -0,0 +1,43 @@
+import itertools
+
+import numpy as np
+
+
+def generate_simplex_grid(n: int, m: int):
+    """
+    Generate a uniform grid of points on the n-dimensional simplex.
+
+    Examples:
+
+        >>> generate_simplex_grid(3,2)
+        array([[0., 0., 1.],
+              [0., 1., 0.],
+              [1., 0., 0.]], dtype=float32)
+
+        >>> generate_simplex_grid(2,3)
+        array([[0. , 1. ],
+              [0.5, 0.5],
+              [1. , 0. ]], dtype=float32)
+
+    Args:
+        n (int): The dimension of the simplex.
+        m (int): The number of grid points along each dimension.
+
+    Returns:
+        list: A list of n-dimensional vectors representing the grid points.
+    """
+    m = m - 1
+    # Generate all combinations of indices summing up to m
+    indices = list(itertools.combinations_with_replacement(range(m + 1), n - 1))
+    # Initialize an empty list to store the grid points
+    grid_points = []
+
+    # Iterate over each combination of indices
+    for idx in indices:
+        # Append 0 and m to the indices
+        extended_idx = [0] + list(idx) + [m]
+        # Compute the vector components by taking the differences between consecutive indices and dividing by m
+        point = [(extended_idx[i + 1] - extended_idx[i]) / m for i in range(n)]
+        grid_points.append(point)
+
+    return np.array(grid_points, dtype=np.float32)

fusion_bench/method/rankone_moe/__init__.py

@@ -0,0 +1,3 @@
+# flake8: noqa F401
+from .clip_rankone_moe import CLIPRankOneMoEAlgorithm
+from .rankone_moe import RankOneMoEAlgorithm

fusion_bench/method/rankone_moe/clip_rankone_moe.py

@@ -0,0 +1,160 @@
+import functools
+import logging
+import os
+from copy import deepcopy
+
+import torch
+from torch import Tensor
+from torch.utils.data import DataLoader
+from transformers.models.clip.modeling_clip import CLIPEncoder
+
+from fusion_bench.dataset import CLIPDataset
+from fusion_bench.method.task_arithmetic.task_arithmetic import task_arithmetic_merge
+from fusion_bench.mixins import CLIPClassificationMixin
+from fusion_bench.modelpool import CLIPVisionModelPool
+from fusion_bench.models.rankone_moe import RankOneMoE
+from fusion_bench.utils.data import InfiniteDataLoader
+
+from .rankone_moe import RankOneMoEAlgorithm
+
+log = logging.getLogger(__name__)
+
+
+class CLIPRankOneMoEAlgorithm(
+    RankOneMoEAlgorithm,
+    CLIPClassificationMixin,
+):
+    """
+    CLIPRankOneMoEAlgorithm is a class that implements the RankOneMoEAlgorithm (https://github.com/EnnengYang/RankOne-MoE)
+    for CLIP models. It extends the RankOneMoEAlgorithm and CLIPClassificationMixin classes.
+
+    Attributes:
+        modelpool (CLIPVisionModelPool): The model pool containing the CLIP models.
+    """
+
+    modelpool: CLIPVisionModelPool = None
+
+    def load_checkpoint(self, model, checkpoint):
+        """
+        Load the checkpoint file.
+
+        Args:
+            model: The model to load the checkpoint into.
+            checkpoint: The path to the checkpoint file.
+        """
+        state = {"model": model}
+        self._fabric.load(checkpoint, state)
+
+    def save_checkpoint(self, model, checkpoint):
+        """
+        Save the checkpoint file.
+
+        Args:
+            model: The model to save the checkpoint from.
+            checkpoint: The path to the checkpoint file.
+        """
+        self._fabric.save(checkpoint, {"model": model})
+
+    def construct_moe_model(self) -> RankOneMoE:
+        """
+        Construct the RankOne-MoE model using the models in the model pool.
+
+        Returns:
+            RankOne-MoE: The constructed MoE model.
+        """
+        base_model = self.modelpool.load_model("_pretrained_")
+        expert_models = [
+            self.modelpool.load_model(m) for m in self.modelpool.model_names
+        ]
+
+        # Merge the models using task arithmetic
+        moe_model = task_arithmetic_merge(
+            # This function modifies the model in place, so we need to pass a deepcopy
+            deepcopy(base_model),
+            expert_models,
+            scaling_factor=self.config.init_lambda,
+        ).requires_grad_(False)
+
+        # Up-scale MLP modules
+        base_encoder: CLIPEncoder = base_model.vision_model.encoder
+        moe_encoder: CLIPEncoder = moe_model.vision_model.encoder
+        expert_encoders = [m.vision_model.encoder for m in expert_models]
+
+        num_layers = len(base_encoder.layers)
+        for layer_idx in range(num_layers):
+            base_mlp = base_encoder.layers[layer_idx].mlp
+            expert_mlps = [e.layers[layer_idx].mlp for e in expert_encoders]
+
+            moe_encoder.layers[layer_idx].mlp = RankOneMoE(
+                hidden_size=base_encoder.config.hidden_size,
+                base_model=base_mlp,
+                expert_models=expert_mlps,
+                init_lambda=self.config.init_lambda,
+                batch_first=True,  # For open_clip models this is False
+                router_hidden_layers=self.config.router_hidden_layers,
+                batch_reduce=self.config.batch_reduce,
+                svd_accelerator=self.config.svd_accelerator,
+                rank_k=self.config.rank_k,
+                select_k=self.config.select_k,
+            )
+
+        return moe_model
+
+    @functools.cache
+    def get_shuffled_test_loader_iter(self, tta_dataset: str):
+        """
+        Get an iterator for the shuffled test data loader.
+
+        Args:
+            tta_dataset (str): The name of the test-time adaptation dataset.
+
+        Returns:
+            Iterator: An iterator for the shuffled test data loader.
+        """
+        dataset = self.modelpool.load_test_dataset(tta_dataset)
+        dataset = CLIPDataset(dataset, processor=self.clip_processor)
+        log.info("get_shuffled_test_loader_iter")
+        loader = DataLoader(
+            dataset,
+            batch_size=self.config.batch_size,
+            shuffle=True,
+            num_workers=self.config.num_workers,
+            pin_memory=True,
+        )
+        loader = self.fabric.setup_dataloaders(loader)
+        return iter(InfiniteDataLoader(loader))
+
+    def on_test_time_adaptation_start(self):
+        """
+        Load the CLIP processor and construct the zero-shot classification head for each task.
+        """
+        self.setup_zero_shot_classification_head()
+
+    def compute_logits(self, module, batch, task) -> Tensor:
+        """
+        Compute the logits for the given batch and task.
+
+        Args:
+            module: The model module.
+            batch: The input batch.
+            task: The task name.
+
+        Returns:
+            Tensor: The computed logits.
+        """
+        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

fusion_bench/method/rankone_moe/rankone_moe.py

@@ -0,0 +1,249 @@
+import logging
+from abc import abstractmethod
+from typing import cast  # noqa: F401
+
+import lightning as L
+import lightning.fabric.wrappers
+import torch
+from lightning.pytorch.profilers import SimpleProfiler
+from omegaconf import DictConfig
+from torch import Tensor
+from torch.utils.data import DataLoader
+from tqdm.autonotebook import tqdm
+
+from fusion_bench.compat.method.base_algorithm import ModelFusionAlgorithm
+from fusion_bench.compat.modelpool import ModelPool
+from fusion_bench.models.rankone_moe import RankOneMoE
+from fusion_bench.utils import timeit_context
+from fusion_bench.utils.parameters import print_parameters
+
+log = logging.getLogger(__name__)
+
+
+def entropy_loss(logits: Tensor) -> Tensor:
+    """
+    Compute the entropy loss of a set of logits.
+
+    Args:
+        logits (Tensor): The logits to compute the entropy loss of.
+
+    Returns:
+        Tensor: The entropy loss of the logits.
+    """
+    probs = torch.softmax(logits, dim=-1)
+    return -torch.sum(probs * torch.log(probs + 1e-8), dim=-1).mean()
+
+
+class RankOneMoEAlgorithm(ModelFusionAlgorithm):
+    """
+    Algorithm for fusing models using RankOne-MoE (https://github.com/EnnengYang/RankOne-MoE).
+
+    This class provides methods for constructing the MoE model, performing test-time adaptation,
+    and running the fusion process.
+
+    Attributes:
+        _fabric (L.Fabric): The fabric for distributed training.
+        modelpool (ModelPool): The pool of models to be fused.
+        profiler (SimpleProfiler): The profiler for measuring performance.
+    """
+
+    _fabric: L.Fabric = None
+    modelpool: ModelPool = None
+
+    def __init__(self, algorithm_config: DictConfig):
+        """
+        Initialize the RankOneMoEAlgorithm with the given configuration.
+
+        Args:
+            algorithm_config (DictConfig): The configuration for the algorithm.
+        """
+        super().__init__(algorithm_config)
+
+        if self._fabric is None and torch.cuda.is_available():
+            self._fabric = L.Fabric(
+                devices=self.config.get("devices", 1),
+            )
+            self._fabric.launch()
+        else:
+            assert "No CUDA device available."
+        self.profiler = SimpleProfiler(
+            self.config.get("cache_dir", "outputs"), "we_moe_profiler.txt"
+        )
+
+    @abstractmethod
+    def load_checkpoint(self, model, checkpoint):
+        """
+        Load the checkpoint file.
+
+        Args:
+            model: The model to load the checkpoint into.
+            checkpoint: The checkpoint file to load.
+        """
+        pass
+
+    @abstractmethod
+    def save_checkpoint(self, model, checkpoint):
+        """
+        Save the checkpoint file.
+
+        Args:
+            model: The model to save the checkpoint from.
+            checkpoint: The checkpoint file to save.
+        """
+        pass
+
+    @abstractmethod
+    def construct_moe_model(self) -> RankOneMoE:
+        """
+        Construct the Mixture of Experts model using the models in the model pool.
+
+        Returns:
+            RankOne-MoE: The constructed MoE model.
+        """
+        pass
+
+    def on_test_time_adaptation_start(self):
+        """
+        Hook method called at the start of test-time adaptation.
+        """
+        pass
+
+    @abstractmethod
+    def get_shuffled_test_loader_iter(self, task: str) -> DataLoader:
+        """
+        Get an iterator for the shuffled test data loader for a specific task.
+
+        Args:
+            task (str): The task for which to get the test data loader.
+
+        Returns:
+            DataLoader: The shuffled test data loader iterator.
+        """
+        pass
+
+    @abstractmethod
+    def compute_logits(self, module, batch, task) -> Tensor:
+        """
+        Compute the logits for a given batch and task.
+
+        Args:
+            module: The model module to use for computing logits.
+            batch: The batch of data.
+            task: The task for which to compute logits.
+
+        Returns:
+            Tensor: The computed logits.
+        """
+        pass
+
+    def test_time_adaptation(self, module: RankOneMoE):
+        """
+        Perform test-time adaptation for the given module.
+
+        Args:
+            module (RankOne-MoE): The MoE module to adapt.
+
+        Returns:
+            RankOne-MoE: The adapted MoE module.
+        """
+        self.on_test_time_adaptation_start()
+
+        # configure optimizer
+        if self.config.optimizer == "adam":
+            optimizer = torch.optim.Adam(
+                [p for p in module.parameters() if p.requires_grad], lr=self.config.lr
+            )
+        else:
+            raise ValueError(f"Unsupported optimizer: {self.config.optimizer}")
+
+        if self._fabric is not None:
+            module, optimizer = self._fabric.setup(module, optimizer)
+
+        module.train()
+
+        if self.config.get("fast_dev_run", False):
+            log.info("Running fast_dev_run, only one step")
+            pbar = tqdm(
+                range(1),
+                "Test-time adaptation",
+                dynamic_ncols=True,
+            )
+        else:
+            pbar = tqdm(
+                range(self.config.max_steps),
+                "Test-time adaptation",
+                dynamic_ncols=True,
+            )
+        for step_idx in pbar:
+            if self.config.use_grad_accumulate:
+                for task in self.modelpool.model_names:
+                    with self.profiler.profile("data time"):
+                        batch = next(self.get_shuffled_test_loader_iter(task))
+                    with self.profiler.profile("forward pass"):
+                        logits = self.compute_logits(module, batch, task)
+                        assert (
+                            logits.dim() == 2
+                        ), f"Expected logits to be 2D, got {logits.dim()}"
+                        loss = entropy_loss(logits)
+                    # .backward() accumulates when .zero_grad() wasn't called
+                    # this can save memory
+                    with self.profiler.profile("backward pass"):
+                        self._fabric.backward(loss, retain_graph=True)
+            else:
+                loss = 0
+                for task in self.modelpool.model_names:
+                    with self.profiler.profile("data time"):
+                        batch = next(self.get_shuffled_test_loader_iter(task))
+                    with self.profiler.profile("forward pass"):
+                        logits = self.compute_logits(module, batch, task)
+                        assert (
+                            logits.dim() == 2
+                        ), f"Expected logits to be 2D, got {logits.dim()}"
+                        loss = loss + entropy_loss(logits)
+                with self.profiler.profile("backward pass"):
+                    self._fabric.backward(loss, retain_graph=True)
+
+            with self.profiler.profile("optimizer step"):
+                optimizer.step()
+                optimizer.zero_grad()
+
+            # print([m for m in module.parameters() if m.requires_grad][0])
+
+        return module
+
+    def run(self, modelpool: ModelPool):
+        """
+        Run the RankOneMoEAlgorithm to fuse models using RankOne-MoE.
+
+        Args:
+            modelpool (ModelPool): The pool of models to be fused.
+
+        Returns:
+            RankOne-MoE: The fused RankOne MoE model.
+        """
+        log.info("Fusing models using RankOne-MoE modules.")
+        self.modelpool = modelpool
+
+        with timeit_context("upscaling models to a RankOne-MoE model"):
+            moe_model = self.construct_moe_model()
+            print_parameters(moe_model)
+
+        if self.config.get("checkpoint", False):
+            log.info(
+                f"load checkpoint from {self.config.checkpoint}, test-time adaptation will be skipped."
+            )
+            self.load_checkpoint(moe_model, self.config.checkpoint)
+        else:
+            with self.profiler.profile("test-time adaptation"):
+                moe_model = self.test_time_adaptation(moe_model)
+            if self.config.get("save_checkpoint", False):
+                log.info(f"save checkpoint to {self.config.save_checkpoint}")
+                self.save_checkpoint(moe_model, self.config.save_checkpoint)
+
+            if lightning.fabric.wrappers.is_wrapped(moe_model):
+                moe_model = lightning.fabric.wrappers._unwrap_objects(moe_model)
+
+        # enable sample-wise adaptation
+        moe_model.batch_reduce = False
+        print(self.profiler.summary())
+        return moe_model

fusion_bench/method/regmean/__init__.py

@@ -0,0 +1,4 @@
+# flake8: noqa F401
+from .clip_regmean import RegMeanAlgorithmForCLIP
+from .gpt2_regmean import RegMeanAlgorithmForGPT2
+from .regmean import RegMeanAlgorithm