fusion-bench 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

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/simple_average.py

@@ -11,8 +11,8 @@ from fusion_bench.modelpool import BaseModelPool
 from fusion_bench.utils.state_dict_arithmetic import (
     state_dict_add,
     state_dict_avg,
-    state_dict_mul,
     state_dict_div,
+    state_dict_mul,
 )
 from fusion_bench.utils.type import StateDictType
 

fusion_bench/method/surgery/__init__.py

@@ -0,0 +1,3 @@
+from .clip_layer_wise_adamerging_surgery import (
+    CLIPLayerWiseAdaMergingSurgeryAlgorithm,
+)

fusion_bench/method/surgery/clip_layer_wise_adamerging_surgery.py

@@ -0,0 +1,157 @@
+"""
+Implementation of the Layer-Wise AdaMerging+Surgery Algorithm.
+
+For more details, please refer to:
+
+- (ICLR 2024) Yang, et.al. AdaMerging: Adaptive Model Merging for Multi-Task Learning. http://arxiv.org/abs/2310.02575
+- (ICML 2024) Yang, et.al. Representation Surgery for Multi-Task Model Merging. https://arxiv.org/abs/2402.02705
+
+Basic Example:
+
+```shell
+fusion_bench \
+    method=surgery/adamerging_surgery \
+    modelpool=CLIPVisionModelPool/clip-vit-base-patch32_TA8 \
+    taskpool=CLIPVisionModelTaskPool/clip-vit-classification_TA8
+```
+"""
+
+import copy
+import functools
+import gc
+import logging
+from typing import TYPE_CHECKING, cast
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import CLIPVisionModel
+
+from fusion_bench.dataset.clip_dataset import CLIPDataset
+from fusion_bench.method.adamerging.layer_wise_adamerging import (
+    LayerWiseAdaMergingAlgorithm,
+)
+from fusion_bench.method.adamerging.utils import get_memory_usage
+from fusion_bench.mixins import CLIPClassificationMixin
+from fusion_bench.modelpool import CLIPVisionModelPool
+from fusion_bench.models.surgery.surgerymodelwrapper import SurgeryModelWrapper
+from fusion_bench.models.wrappers.layer_wise_fusion import LayerWiseMergedModel
+
+log = logging.getLogger(__name__)
+
+
+class CLIPLayerWiseAdaMergingSurgeryAlgorithm(
+    CLIPClassificationMixin,
+    LayerWiseAdaMergingAlgorithm,
+):
+
+    def on_test_time_adaptation_start(self):
+        """
+        Here we load the CLIP processor and construct the zero-shot classification head for each task.
+        """
+        self.setup_zero_shot_classification_head()
+
+    @functools.cache
+    def get_shuffled_test_loader_iter(self, task: str):
+        return super().get_shuffled_test_loader_iter(
+            task,
+            batch_size=self.config.batch_size,
+            num_workers=self.config.num_workers,
+        )
+
+    def run(self, modelpool: CLIPVisionModelPool, **kwargs):
+        """
+        Run the Layer-Wise AdaMerging+Surgery Algorithm.
+
+        This method constructs the wrapped model and performs test-time adaptation if necessary. Then, it will perform surgery.
+
+        Args:
+            modelpool (ModelPool): The model pool containing the pretrained and fine-tuned models.
+
+        Returns:
+            LayerWiseMergedModel: The merged model after test-time adaptation.
+        """
+        log.info("Fusing models using layer-wise adaptive merging.")
+        self.modelpool = modelpool
+        self.log_hyperparams(self.config)
+
+        # === Start of the AdaMerging Algorithm ===
+        with self.profile("construct the wrapped model"):
+            module = cast(
+                LayerWiseMergedModel[CLIPVisionModel],
+                self.construct_layer_wise_merged_model(modelpool),
+            )
+
+        if self.config.weights is not None:
+            # skip the test-time adaptation
+            merged_model = copy.deepcopy(module.merge_and_unload())
+        else:
+            with self.profile("test-time adaptation"):
+                module = self.test_time_adaptation(module)
+            if self.config.get("save_merging_weights", False):
+                self.save_merging_weights(
+                    self.config.save_merging_weights, module.merge_weight
+                )
+            merged_model = copy.deepcopy(module.merge_and_unload())
+
+        # free memory
+        del module
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        # === Start of the Surgery Algorithm ===
+        log.info("start performing Surgery")
+        alpha_model = SurgeryModelWrapper(
+            merged_model,
+            modelpool.model_names,
+            projection_dim=merged_model.config.projection_dim,
+        )
+        alpha_model = self.fabric.setup(alpha_model)
+        log.info(get_memory_usage("after freeing memory, the memory usage of GPU is:"))
+
+        optimizer = torch.optim.Adam(
+            alpha_model.collect_trainable_params(),
+            lr=1e-3,
+            betas=(0.9, 0.999),
+            weight_decay=0.0,
+        )
+
+        finetuned_models = {
+            model_name: modelpool.load_model(model_name)
+            for model_name in modelpool.model_names
+        }
+        for name, model in finetuned_models.items():
+            model.requires_grad_(False)
+            model = self.fabric.to_device(model)
+            model.eval()
+
+        for iteration in tqdm(
+            range(self.config.surgery_steps),
+            "surgery",
+            dynamic_ncols=True,
+        ):
+            for dataset_name in modelpool.model_names:
+                batch = next(self.get_shuffled_test_loader_iter(dataset_name))
+                finetuned_feature = self.compute_features(
+                    finetuned_models[dataset_name], batch[0]
+                )
+                features, _, _ = alpha_model.compute_surgery_features(
+                    lambda model: self.compute_features(model, batch[0]),
+                    dataset_name,
+                )
+
+                loss = F.l1_loss(features, finetuned_feature)
+
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+
+            if ((iteration + 1) % self.config.eval_iterations) == 0:
+                # print(list(alpha_model.collect_trainable_params()))
+                # Evaluate try to use the test module in fusion bench
+                log.info(f"iteration: {iteration+1}")
+                self._program.evaluate_merged_model(self._program.taskpool, alpha_model)
+
+        log.info("test the result of Adamerging")
+        return merged_model

fusion_bench/mixins/__init__.py

@@ -10,10 +10,12 @@ _import_structure = {
     "serialization": ["YAMLSerializationMixin", "BaseYAMLSerializableModel"],
     "simple_profiler": ["SimpleProfilerMixin"],
     "clip_classification": ["CLIPClassificationMixin"],
+    "fabric_training": ["FabricTrainingMixin"],
 }
 
 if TYPE_CHECKING:
     from .clip_classification import CLIPClassificationMixin
+    from .fabric_training import FabricTrainingMixin
     from .lightning_fabric import LightningFabricMixin
     from .serialization import BaseYAMLSerializableModel, YAMLSerializationMixin
     from .simple_profiler import SimpleProfilerMixin

fusion_bench/mixins/clip_classification.py

@@ -2,14 +2,24 @@ import functools
 import logging
 import os
 from copy import deepcopy
-from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union, cast  # noqa: F401
+from typing import (  # noqa: F401
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import torch
+from omegaconf import DictConfig
 from torch import nn
 from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import CLIPModel, CLIPProcessor, CLIPVisionModel
-from omegaconf import DictConfig
 
 from fusion_bench.dataset.clip_dataset import CLIPDataset
 from fusion_bench.mixins import LightningFabricMixin
@@ -18,10 +28,12 @@ from fusion_bench.models.hf_clip import HFCLIPClassifier
 from fusion_bench.tasks.clip_classification import get_classnames_and_templates
 from fusion_bench.utils.data import InfiniteDataLoader
 
-log = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from transformers.models.clip.modeling_clip import CLIPVisionTransformer
 
-TensorOrModule = TypeVar("TensorOrModule", torch.Tensor, torch.nn.Module, Any)
+log = logging.getLogger(__name__)
 
+# disable tokenizers parallelism by default to avoid deadlocks
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 
@@ -43,12 +55,7 @@ class CLIPClassificationMixin(LightningFabricMixin):
     # a dict of zeroshot weights for each task, each key is the task name
     zeroshot_weights_cache_dir: str = "outputs/cache/clip_zeroshot_weights"
     zeroshot_weights: Dict[str, torch.Tensor] = {}
-
-    def __init__(self, algorithm_config: DictConfig) -> None:
-        super().__init__(algorithm_config)
-        self.whether_setup_zero_shot_classification_head = (
-            False  # We want to only do this once
-        )
+    whether_setup_zero_shot_classification_head = False
 
     @property
     def clip_processor(self):
@@ -180,13 +187,30 @@ class CLIPClassificationMixin(LightningFabricMixin):
 
     def compute_logits(
         self,
-        module: Union[nn.Module, CLIPVisionModel],
+        module: Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"],
         images: torch.Tensor,
         task: str,
+        image_embeds: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
+        """
+        Compute the logits of the images for a given task.
+
+        Args:
+            module (Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"]): The module to compute the logits.
+            images (torch.Tensor): The images to compute the logits.
+            task (str): The task to compute the logits.
+            image_embeds (Optional[torch.Tensor]): The precomputed image embeddings. If None, the image embeddings will be computed.
+
+        Returns:
+            torch.Tensor: The logits of the images.
+        """
         text_embeds = self.zeroshot_weights[task]
 
-        image_embeds = module(images)[1]
+        if image_embeds is None:
+            image_embeds = module(images)[1]
+        assert isinstance(
+            image_embeds, torch.Tensor
+        ), f"`image_embeds` must be a tensor, but got {type(image_embeds)}"
         image_embeds = self.visual_projection(image_embeds)
 
         # normalize embeddings
@@ -199,3 +223,27 @@ class CLIPClassificationMixin(LightningFabricMixin):
         logits_per_image = logits_per_text.t()
 
         return logits_per_image
+
+    def compute_features(
+        self,
+        module: Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"],
+        images: torch.Tensor,
+        normalize: bool = True,
+    ) -> torch.Tensor:
+        """
+        Extracts image features using CLIP's vision encoder and visual projection.
+
+        Args:
+            module (Union[nn.Module, CLIPVisionModel, "CLIPVisionTransformer"]): The CLIP vision encoder module.
+            images (torch.Tensor): Input image batch to process.
+            normalize (bool): Whether to normalize the image embeddings.
+
+        Returns:
+            torch.Tensor: Normalized image embeddings with dimension matching CLIP's projection space (`projection_dim` in model config).
+        """
+        image_embeds = module(images)[1]
+        image_embeds = self.visual_projection(image_embeds)
+
+        if normalize:
+            image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+        return image_embeds