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

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"
 
 
@@ -120,13 +132,13 @@ class CLIPClassificationMixin(LightningFabricMixin):
 
         # get cache directory
         if self.modelpool.has_pretrained:
-            model_name = self.modelpool.get_model_config(
-                "_pretrained_"
-            ).pretrained_model_name_or_path
+            model_name = self.modelpool.get_model_config("_pretrained_")
+            if not isinstance(model_name, str):
+                model_name = model_name.pretrained_model_name_or_path
         else:
-            model_name = self.modelpool.get_model_config(
-                self.modelpool.model_names[0]
-            ).pretrained_model_name_or_path
+            model_name = self.modelpool.get_model_config(self.modelpool.model_names[0])
+            if not isinstance(model_name, str):
+                model_name = model_name.pretrained_model_name_or_path
         cache_dir = os.path.join(
             self.zeroshot_weights_cache_dir,
             os.path.normpath(model_name.split("/")[-1]),
@@ -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,9 +1,11 @@
+import functools
 import logging
 import os
 from typing import TYPE_CHECKING, Any, List, Optional, TypeVar
 
 import lightning as L
 import torch
+from lightning.fabric.connector import _is_using_cli
 from lightning.fabric.loggers import TensorBoardLogger
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
@@ -13,6 +15,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 +35,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.
@@ -70,7 +80,8 @@ class LightningFabricMixin:
                 self._fabric_instance = L.Fabric()
             else:
                 self._fabric_instance = instantiate(config.fabric)
-            self._fabric_instance.launch()
+            if not _is_using_cli():  # if not using cli, launch the fabric
+                self._fabric_instance.launch()
             # Set the log directory in config if it is not already set
             if (
                 self.log_dir is not None

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

@@ -147,7 +147,6 @@ class BaseModelPool(BaseYAMLSerializableModel):
             DictConfig: The configuration for the specified model.
         """
         model_config = self._models[model_name]
-        assert isinstance(model_config, DictConfig), "Model config must be a DictConfig"
         if return_copy:
             model_config = deepcopy(model_config)
         return model_config

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

fusion_bench/modelpool/causal_lm/causal_lm.py

@@ -3,6 +3,7 @@ import os
 from copy import deepcopy
 from typing import Any, Optional, TypeAlias, Union, cast  # noqa: F401
 
+import peft
 from omegaconf import DictConfig, flag_override
 from torch import nn
 from torch.nn.modules import Module
@@ -23,28 +24,6 @@ log = logging.getLogger(__name__)
 CausalLM: TypeAlias = Union[LlamaForCausalLM, MistralForCausalLM, Any]
 
 
-def config_priority_get(priority_config, general_config, key, default):
-    """
-    Retrieve a configuration value with priority.
-
-    This function retrieves the value associated with `key` from `priority_config` if it exists.
-    If the key is not found in `priority_config`, it retrieves the value from `general_config`.
-    If the key is not found in either configuration, it returns the provided `default` value.
-
-    Args:
-        priority_config (dict): The configuration dictionary with higher priority.
-        general_config (dict): The general configuration dictionary.
-        key (str): The key to look up in the configuration dictionaries.
-        default: The default value to return if the key is not found in either configuration.
-
-    Returns:
-        The value associated with `key` from `priority_config` or `general_config`, or the `default` value if the key is not found.
-    """
-    if key in priority_config:
-        return priority_config[key]
-    return general_config.get(key, default)
-
-
 class CausalLMPool(BaseModelPool):
     _config_mapping = BaseModelPool._config_mapping | {
         "_tokenizer": "tokenizer",
@@ -138,3 +117,23 @@ class CausalLMBackbonePool(CausalLMPool):
             model_name_or_config, *args, **kwargs
         )
         return model.model.layers
+
+
+def load_peft_causal_lm(
+    base_model_path: str,
+    peft_model_path: str,
+    torch_dtype: str = "bfloat16",
+    is_trainable: bool = True,
+    merge_and_unload: bool = False,
+):
+    base_model = LlamaForCausalLM.from_pretrained(
+        base_model_path, torch_dtype=torch_dtype
+    )
+    model = peft.PeftModel.from_pretrained(
+        base_model,
+        peft_model_path,
+        is_trainable=is_trainable,
+    )
+    if merge_and_unload:
+        model = model.merge_and_unload()
+    return model