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

fusion_bench/method/pruning/llama_random_prune.py

@@ -1,7 +1,7 @@
-from typing import Literal, Optional, Union  # noqa: F401
+from typing import Dict, Literal, Optional, Union  # noqa: F401
 
 import torch
-from torch import Dict, nn
+from torch import nn
 from tqdm.auto import tqdm
 from transformers import LlamaForCausalLM, LlamaModel
 

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,7 +2,17 @@ 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
@@ -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"
 
 
@@ -175,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
@@ -194,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

fusion_bench/mixins/fabric_training.py

@@ -0,0 +1,320 @@
+import itertools
+import logging
+import os
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Literal, Union
+
+import torch
+from torch import Tensor, nn
+from tqdm.auto import tqdm
+
+from .lightning_fabric import LightningFabricMixin
+
+if TYPE_CHECKING:
+    from lightning.fabric.wrappers import (
+        _FabricDataLoader,
+        _FabricModule,
+        _FabricOptimizer,
+    )
+
+log = logging.getLogger(__name__)
+
+
+class FabricTrainingMixin(LightningFabricMixin):
+    """
+    This is a general purpose mixin for training a model with PyTorch Lightning.
+    """
+
+    _latest_saved_checkpoint_global_step: int = -1
+    """The global step index of the latest saved checkpoint."""
+    _expected_total_steps: int = None
+    """The expected total number of steps of the entire training."""
+    is_training: bool
+    """Whether the training is in progress. If set to False, the training will stop."""
+    epoch_idx: int
+    """The epoch index, which is the number of epochs completed."""
+    global_step_idx: int
+    """The global step index, which is the number of parameter update steps."""
+    max_epochs: int
+    """Max number of epochs of the entire training."""
+    max_steps: int
+    """Max number of parameter update steps of the entire training."""
+    max_steps_per_epoch: int
+    """Max number of parameter update steps per epoch."""
+    gradient_clip_algorithm: Literal["value", "norm"]
+    """The algorithm to clip gradients. Available options: 'value', 'norm'."""
+    gradient_clip_val: float
+    """The value to clip gradients. If None, no clipping is applied."""
+    accumulate_grad_batches: int
+    """The number of gradient accumulation steps. The effective global batch size is `the batch size per device` x `the number of devices` x `the number of gradient accumulation steps`."""
+    lr_scheduler_interval: Literal["step", "epoch"]
+    """The interval to run the learning rate scheduler. Available options: 'step', 'epoch'."""
+    lr_scheduler_frequency: int
+    """The frequency to run the learning rate scheduler."""
+    checkpoint_save_interval: Literal["step", "epoch"]
+    """The interval to save the model checkpoint. Available options: 'step', 'epoch'."""
+    checkpoint_save_frequency: int
+    """The frequency to save the model checkpoint."""
+
+    def clip_gradients_if_needed(self, model, optimizer):
+        """
+        Clips gradients if the gradient clipping value is set.
+
+        Args:
+            model (nn.Module): The model whose gradients need to be clipped.
+            optimizer (torch.optim.Optimizer): The optimizer used for training.
+        """
+        fabric = self.fabric
+
+        if self.gradient_clip_val is not None:
+            if self.gradient_clip_algorithm == "value":
+                fabric.clip_gradients(model, optimizer, clip_val=self.gradient_clip_val)
+            elif self.gradient_clip_algorithm == "norm":
+                fabric.clip_gradients(model, optimizer, max_norm=self.gradient_clip_val)
+            else:
+                raise ValueError(
+                    f"Unknown gradient clip algorithm: {self.gradient_clip_algorithm}. Available options: 'value', 'norm'"
+                )
+
+    def compute_expected_total_steps(
+        self, train_dataloader: torch.utils.data.DataLoader
+    ):
+        """
+        Computes the expected total number of steps for the entire training.
+
+        Args:
+            train_dataloader (torch.utils.data.DataLoader): The dataloader for the training data.
+        """
+        # compute expected total steps
+        self._expected_total_steps = []
+        if self.max_steps > 0:
+            self._expected_total_steps.append(self.max_steps)
+        if self.max_steps_per_epoch > 0 and self.max_epochs > 0:
+            self._expected_total_steps.append(
+                self.max_steps_per_epoch * self.max_epochs
+            )
+        if self.max_epochs > 0:
+            self._expected_total_steps.append(
+                len(train_dataloader) * self.max_epochs // self.accumulate_grad_batches
+            )
+        self._expected_total_steps = min(self._expected_total_steps)
+        log.info(f"Expected total steps: {self._expected_total_steps}")
+
+    @property
+    def expected_total_steps(self):
+        """
+        The expected total number of steps of the entire training. You need to run `compute_expected_total_steps` method to compute this value before accessing it.
+
+        Raises:
+            ValueError: If the expected total steps have not been computed.
+        """
+        if self._expected_total_steps is None:
+            raise ValueError(
+                "The expected total steps have not been computed. Run `compute_expected_total_steps` method."
+            )
+        else:
+            return self._expected_total_steps
+
+    def conditional_checkpoint_save(
+        self,
+        stage: Literal["end_of_step", "end_of_epoch", "end_of_training"],
+        *args,
+        **kwargs,
+    ):
+        """
+        Conditionally saves a checkpoint based on the current training stage.
+
+        Args:
+            stage (Literal["end_of_step", "end_of_epoch", "end_of_training"]): The current stage of training.
+        """
+        if stage == "end_of_step":
+            if (
+                self.checkpoint_save_interval == "step"
+                and (self.global_step_idx + 1) % self.checkpoint_save_frequency == 0
+            ):
+                save_path = os.path.join(
+                    self.log_dir, "checkpoints", f"step={self.global_step_idx}.ckpt"
+                )
+                self.save_checkpoint(save_path, *args, **kwargs)
+        elif stage == "end_of_epoch":
+            if (
+                self.checkpoint_save_interval == "epoch"
+                and (self.epoch_idx + 1) % self.checkpoint_save_frequency == 0
+            ):
+                save_path = os.path.join(
+                    self.log_dir, "checkpoints", f"epoch={self.epoch_idx}.ckpt"
+                )
+                self.save_checkpoint(save_path, *args, **kwargs)
+        elif stage == "end_of_training":
+            # if the checkpoint has not been saved yet, save it
+            if self.global_step_idx > self._latest_saved_checkpoint_global_step:
+                save_path = os.path.join(
+                    self.log_dir,
+                    "checkpoints",
+                    f"epoch={self.epoch_idx}_step={self.global_step_idx}.ckpt",
+                )
+                self.save_checkpoint(save_path, *args, **kwargs)
+                try:
+                    os.symlink(
+                        src=save_path,
+                        dst=os.path.join(
+                            self.log_dir, "checkpoints", "latest_model.ckpt"
+                        ),
+                        target_is_directory=os.path.isdir(save_path),
+                    )
+                except Exception as e:
+                    log.error(f"Failed to create symlink: {e}")
+        else:
+            raise ValueError(
+                f"Unknown stage: {stage}. Available options: 'end_of_step', 'end_of_epoch', 'end_of_training'"
+            )
+
+    @abstractmethod
+    def save_checkpoint(self, path, **kwargs):
+        """
+        Saves a checkpoint of the model.
+
+        Args:
+            path (str): The path where the checkpoint will be saved.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError("save_checkpoint method is not implemented")
+
+    def train(
+        self,
+        model: Union[nn.Module, "_FabricModule"],
+        optimizer: Union[torch.optim.Optimizer, "_FabricOptimizer"],
+        lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
+    ):
+        """
+        Trains the model.
+
+        The global batch size is `the batch size per device` x `the number of devices` x `the number of gradient accumulation steps`.
+
+        Args:
+            model (Union[nn.Module, "_FabricModule"]): The model to be trained.
+            optimizer (Union[torch.optim.Optimizer, "_FabricOptimizer"]): The optimizer used for training.
+            lr_scheduler (torch.optim.lr_scheduler.LRScheduler): The learning rate scheduler.
+        """
+        fabric = self.fabric
+        self.is_training = True
+        # number of parameter update iterations, not the number of batches
+        self.global_step_idx = 0
+        model.train()
+        optimizer.zero_grad()
+        for epoch_idx in tqdm(
+            range(self.max_epochs) if self.max_epochs > 0 else itertools.count(0),
+            "Training Epoch",
+            dynamic_ncols=True,
+            leave=False,
+            disable=not fabric.is_global_zero,
+        ):
+            self.epoch_idx = epoch_idx
+            self.train_epoch(model, optimizer, lr_scheduler)
+            # run lr_scheduler at the end of the epoch if interval is set to "epoch"
+            if (
+                self.lr_scheduler_interval == "epoch"
+                and (epoch_idx + 1) % self.lr_scheduler_frequency == 0
+            ):
+                lr_scheduler.step()
+
+            # save the model at the end of the epoch if interval is set to "epoch" and frequency is met
+            self.conditional_checkpoint_save(stage="end_of_epoch")
+
+            if not self.is_training:
+                break
+
+        optimizer.zero_grad()
+        # save the model at the end of training
+        self.conditional_checkpoint_save(stage="end_of_training")
+        return model
+
+    @abstractmethod
+    def train_epoch(
+        self,
+        model: Union[nn.Module, "_FabricModule"],
+        optimizer: Union[torch.optim.Optimizer, "_FabricOptimizer"],
+        lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
+    ):
+        """
+        Trains the model for one epoch.
+
+        Args:
+            model (Union[nn.Module, "_FabricModule"]): The model to be trained.
+            optimizer (Union[torch.optim.Optimizer, "_FabricOptimizer"]): The optimizer used for training.
+            lr_scheduler (torch.optim.lr_scheduler.LRScheduler): The learning rate scheduler.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError(
+            "Copy this as a template and implement your own train_epoch method"
+        )
+        fabric = self.fabric
+
+        accumulated_loss = 0
+        for step_idx, batch in enumerate(
+            pbar := tqdm(
+                self.train_dataloader,
+                desc="Training Batches",
+                dynamic_ncols=True,
+                leave=False,
+                disable=not fabric.is_global_zero,
+            )
+        ):
+            is_accumulating = (step_idx + 1) % self.accumulate_grad_batches != 0
+
+            # disable gradient synchronization if accumulating gradients across steps for improved performance
+            with fabric.no_backward_sync(self.model, enabled=is_accumulating):
+                # use_cache=True is not compatible with gradient checkpointing, so we disable it here
+                output = self.compute_loss(batch)
+                loss = output["loss"] / self.accumulate_grad_batches
+
+                fabric.backward(loss)
+                accumulated_loss += loss.item()
+
+            # 1. update the model parameters if not accumulating gradients
+            # 2. step the lr_scheduler if interval is set to "step" and frequency is met
+            # 3. save the model if interval is set to "step" and frequency is met
+            # 4. log metrics
+            # 5. increase the global step index and reset the accumulated metrics
+            if not is_accumulating:
+                self.clip_gradients_if_needed(model, optimizer)
+
+                # run lr_scheduler at the end of the step if interval is set to "step"
+                if (
+                    self.lr_scheduler_interval == "step"
+                    and (self.global_step_idx + 1) % self.lr_scheduler_frequency == 0
+                ):
+                    lr_scheduler.step()
+
+                # update the model parameters and zero the gradients
+                optimizer.step()
+                optimizer.zero_grad()
+
+                metrics = {
+                    "train/loss": accumulated_loss,
+                    "train/lr": optimizer.param_groups[0]["lr"],
+                }
+
+                fabric.log_dict(metrics, step=self.global_step_idx)
+                pbar.set_postfix(metrics)
+
+                # save the model at the end of the step if interval is set to "step" and frequency is met
+                self.conditional_checkpoint_save(stage="end_of_step")
+
+                # break if max_steps_per_epoch is set, and exit epoch
+                if (
+                    self.max_steps_per_epoch > 0
+                    and step_idx + 1 >= self.max_steps_per_epoch
+                ):
+                    break
+                # break if max_steps is set, and exit training
+                if self.max_steps > 0 and self.global_step_idx >= self.max_steps - 1:
+                    self.is_training = False
+                    break
+
+                self.global_step_idx += 1
+                accumulated_loss = 0

fusion_bench/mixins/lightning_fabric.py

@@ -1,3 +1,4 @@
+import functools
 import logging
 import os
 from typing import TYPE_CHECKING, Any, List, Optional, TypeVar
@@ -13,6 +14,7 @@ from fusion_bench.utils.instantiate import instantiate
 
 if TYPE_CHECKING:
     import lightning.fabric.loggers.tensorboard
+    from lightning.fabric.strategies import FSDPStrategy
 
 log = logging.getLogger(__name__)
 
@@ -32,6 +34,13 @@ def get_policy(*args: str) -> set:
     return {import_object(arg) for arg in args}
 
 
+def get_size_based_auto_wrap_policy(*args, **kwargs):
+    from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
+
+    policy = functools.partial(size_based_auto_wrap_policy, *args, **kwargs)
+    return policy
+
+
 class LightningFabricMixin:
     """
     A mixin class for integrating Lightning Fabric into a project.

fusion_bench/modelpool/__init__.py

@@ -16,6 +16,7 @@ _import_structure = {
         "HuggingFaceGPT2ClassificationPool",
         "GPT2ForSequenceClassificationPool",
     ],
+    "seq_classification_lm": ["SeqenceClassificationModelPool"],
 }
 
 
@@ -31,6 +32,7 @@ if TYPE_CHECKING:
     from .nyuv2_modelpool import NYUv2ModelPool
     from .PeftModelForSeq2SeqLM import PeftModelForSeq2SeqLMPool
     from .seq2seq_lm import Seq2SeqLMPool
+    from .seq_classification_lm import SeqenceClassificationModelPool
 
 else:
     sys.modules[__name__] = LazyImporter(

fusion_bench/modelpool/causal_lm/__init__.py

@@ -1,2 +1,2 @@
 # flake8: noqa F401
-from .causal_lm import CausalLMBackbonePool, CausalLMPool
+from .causal_lm import CausalLMBackbonePool, CausalLMPool, load_peft_causal_lm