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

fusion_bench/optim/exception.py

@@ -0,0 +1,47 @@
+class NoSparseGradientError(Exception):
+    """Raised when the gradient is sparse gradient.
+
+    :param optimizer_name: str. optimizer name.
+    :param note: str. special conditions to note (default '').
+    """
+
+    def __init__(self, optimizer_name: str, note: str = ""):
+        self.note: str = " " if not note else f" w/ {note} "
+        self.message: str = (
+            f"[-] {optimizer_name}{self.note}does not support sparse gradient."
+        )
+        super().__init__(self.message)
+
+
+class ZeroParameterSizeError(Exception):
+    """Raised when the parameter size is 0."""
+
+    def __init__(self):
+        self.message: str = "[-] parameter size is 0"
+        super().__init__(self.message)
+
+
+class NoClosureError(Exception):
+    """Raised when there's no closure function."""
+
+    def __init__(self, optimizer_name: str, note: str = ""):
+        self.message: str = f"[-] {optimizer_name} requires closure.{note}"
+        super().__init__(self.message)
+
+
+class NegativeLRError(Exception):
+    """Raised when learning rate is negative."""
+
+    def __init__(self, lr: float, lr_type: str = ""):
+        self.note: str = lr_type if lr_type else "learning rate"
+        self.message: str = f"[-] {self.note} must be positive. ({lr} > 0)"
+        super().__init__(self.message)
+
+
+class NegativeStepError(Exception):
+    """Raised when step is negative."""
+
+    def __init__(self, num_steps: int, step_type: str = ""):
+        self.note: str = step_type if step_type else "step"
+        self.message: str = f"[-] {self.note} must be positive. ({num_steps} > 0)"
+        super().__init__(self.message)

fusion_bench/optim/lr_scheduler/__init__.py

@@ -0,0 +1 @@
+from .linear_warmup import *

fusion_bench/optim/lr_scheduler/linear_warmup.py

@@ -0,0 +1,222 @@
+"""
+Modified from pytorch_optimizer: https://github.com/kozistr/pytorch_optimizer/blob/main/pytorch_optimizer/lr_scheduler/linear_warmup.py
+"""
+
+import math
+from abc import ABC, abstractmethod
+from typing import List
+
+import numpy as np
+import torch
+
+from fusion_bench.optim.exception import NegativeLRError, NegativeStepError
+
+__all__ = [
+    "BaseLinearWarmupScheduler",
+    "LinearWarmupScheduler",
+    "CosineDecayWithWarmup",
+    "PolySchedulerWithWarmup",
+]
+
+
+class BaseLinearWarmupScheduler(ABC):
+    r"""BaseLinearWarmupScheduler class.
+
+    The LR Scheduler class based on this class has linear warmup strategy.
+
+    Args:
+        optimizer (torch.optim.Optimizer): Optimizer. It will set learning rate to all trainable parameters in optimizer.
+        T_max (int): Total steps to train.
+        max_lr (float): Maximum learning rate.
+        min_lr (float): Minimum learning rate.
+        init_lr (float): Initial learning rate.
+        warmup_steps (int): Steps to warm-up.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        T_max: int,
+        max_lr: float,
+        min_lr: float = 0.0,
+        init_lr: float = 0.0,
+        warmup_steps: int = 0,
+    ):
+        """
+        Initialize the BaseLinearWarmupScheduler.
+
+        Args:
+            optimizer (torch.optim.Optimizer): Optimizer to apply the learning rate schedule.
+            T_max (int): Total number of training steps.
+            max_lr (float): Maximum learning rate.
+            min_lr (float): Minimum learning rate.
+            init_lr (float): Initial learning rate.
+            warmup_steps (int): Number of steps for the warm-up phase.
+        """
+        self.optimizer = optimizer
+        self.total_steps = T_max
+        self.max_lr = max_lr
+        self.min_lr = min_lr
+        self.init_lr = init_lr
+        self.warmup_steps = warmup_steps
+
+        self.step_t: int = 0
+        self.base_lrs: List[float] = []
+
+        # record current value in self._last_lr to match API from torch.optim.lr_scheduler
+        self.last_lr: List[float] = [init_lr]
+
+        self.validate_parameters()
+
+        self._init_lr()
+
+    def validate_parameters(self):
+        """
+        Validate the parameters to ensure they are non-negative.
+
+        Raises:
+            NegativeLRError: If any of the learning rates are negative.
+            NegativeStepError: If any of the step values are negative.
+        """
+        if self.min_lr < 0:
+            raise NegativeLRError(self.min_lr, "min_lr")
+
+        if self.max_lr < 0:
+            raise NegativeLRError(self.max_lr, "max_lr")
+
+        if self.init_lr < 0:
+            raise NegativeLRError(self.init_lr, "init_lr")
+
+        if self.total_steps < 0:
+            raise NegativeStepError(self.total_steps, "T_max")
+
+        if self.warmup_steps < 0:
+            raise NegativeStepError(self.warmup_steps, "warmup_steps")
+
+    def _init_lr(self):
+        """
+        Initialize the learning rate for each parameter group in the optimizer.
+        """
+        self.base_lrs = []
+        for param_group in self.optimizer.param_groups:
+            param_group["lr"] = self.min_lr
+            self.base_lrs.append(self.min_lr)
+
+    def step(self):
+        """
+        Update the learning rate for the current step.
+
+        Returns:
+            float: The updated learning rate.
+        """
+        if self.step_t < self.warmup_steps:
+            value = (
+                self.init_lr
+                + (self.max_lr - self.init_lr) * self.step_t / self.warmup_steps
+            )
+        elif self.step_t == self.warmup_steps:
+            value = self.max_lr
+        else:
+            value = self._step()
+
+        self.step_t += 1
+
+        if self.optimizer is not None:
+            for param_group in self.optimizer.param_groups:
+                param_group["lr"] = value
+
+        self.last_lr = [value]
+
+        return value
+
+    @abstractmethod
+    def _step(self) -> float:  # pragma: no cover
+        """
+        Abstract method to calculate the learning rate for the current step.
+
+        Returns:
+            float: The calculated learning rate.
+        """
+        raise NotImplementedError
+
+    def get_lr(self) -> float:
+        """
+        Get the current learning rate.
+
+        Returns:
+            float: The current learning rate.
+        """
+        return self.last_lr[0]
+
+
+class LinearWarmupScheduler(BaseLinearWarmupScheduler):
+    r"""Linear LR Scheduler w/ linear warmup."""
+
+    def _step(self) -> float:
+        """
+        Calculate the learning rate for the current step using a linear decay.
+
+        Returns:
+            float: The calculated learning rate.
+        """
+        return self.max_lr + (self.min_lr - self.max_lr) * (
+            self.step_t - self.warmup_steps
+        ) / (self.total_steps - self.warmup_steps)
+
+
+class CosineDecayWithWarmup(BaseLinearWarmupScheduler):
+    r"""Cosine LR Scheduler w/ linear warmup."""
+
+    def _step(self) -> float:
+        """
+        Calculate the learning rate for the current step using a cosine decay.
+
+        Returns:
+            float: The calculated learning rate.
+        """
+        phase: float = (
+            (self.step_t - self.warmup_steps)
+            / (self.total_steps - self.warmup_steps)
+            * math.pi
+        )
+        return self.min_lr + (self.max_lr - self.min_lr) * (np.cos(phase) + 1.0) / 2.0
+
+
+class PolySchedulerWithWarmup(BaseLinearWarmupScheduler):
+    r"""Poly LR Scheduler.
+
+    Args:
+        poly_order (float): LR scheduler decreases with steps.
+    """
+
+    def __init__(self, optimizer, poly_order: float = 0.5, **kwargs):
+        """
+        Initialize the PolySchedulerWithWarmup.
+
+        Args:
+            optimizer (torch.optim.Optimizer): Optimizer to apply the learning rate schedule.
+            poly_order (float): Order of the polynomial for the learning rate decay.
+            kwargs: Additional arguments for the base class.
+
+        Raises:
+            ValueError: If poly_order is not positive.
+        """
+        self.poly_order = poly_order
+
+        if poly_order <= 0:
+            raise ValueError(f"[-] poly_order must be positive. {poly_order}")
+
+        super().__init__(optimizer, **kwargs)
+
+    def _step(self) -> float:
+        """
+        Calculate the learning rate for the current step using a polynomial decay.
+
+        Returns:
+            float: The calculated learning rate.
+        """
+        return (
+            self.min_lr
+            + (self.max_lr - self.min_lr)
+            * (self.step_t - self.warmup_steps) ** self.poly_order
+        )

fusion_bench/optim/lr_scheduler/utils/__init__.py

@@ -0,0 +1 @@
+from .visualization import *

fusion_bench/optim/lr_scheduler/utils/visualization.py

@@ -0,0 +1,119 @@
+"""
+This module provides utilities for visualizing learning rate schedulers.
+
+Functions:
+    simulate_scheduler(lr_scheduler, steps): Simulates the learning rate scheduler for a given number of steps.
+    plot_lr_schedulers(lr_schedulers, steps, titles): Plots the learning rates of one or more schedulers over a number of steps.
+"""
+
+from typing import TYPE_CHECKING, List, Union
+
+import matplotlib.pyplot as plt
+import torch
+
+if TYPE_CHECKING:
+    from torch.optim.lr_scheduler import LRScheduler
+
+__all__ = ["simulate_scheduler", "plot_lr_schedulers"]
+
+
+def simulate_scheduler(lr_scheduler, steps: int):
+    """
+    Simulates the learning rate scheduler for a given number of steps.
+
+    Args:
+        lr_scheduler (torch.optim.lr_scheduler.LRScheduler): The learning rate scheduler object.
+        steps (int): The number of steps to simulate.
+
+    Returns:
+        List[float]: A list of learning rates for each step.
+    """
+    lrs = []
+    for _ in range(steps):
+        lr = lr_scheduler.step()
+        lrs.append(lr)
+    return lrs
+
+
+def plot_lr_schedulers(
+    lr_schedulers: Union["LRScheduler", List["LRScheduler"]],
+    steps: int,
+    titles: Union[str, List[str]],
+    show_plot: bool = True,
+):
+    """
+    Plots the learning rates of one or more schedulers over a number of steps.
+
+    Args:
+        lr_schedulers (Union[LRScheduler, List[LRScheduler]]): One or more learning rate scheduler objects.
+        steps (int): The number of steps to simulate.
+        titles (Union[str, List[str]]): Titles for the plots.
+
+    Returns:
+        fig, axes: The matplotlib figure and axes objects.
+    """
+    # Handle single scheduler
+    if isinstance(lr_schedulers, torch.optim.lr_scheduler.LRScheduler):
+        lr_schedulers = [lr_schedulers]
+    if isinstance(titles, str):
+        titles = [titles]
+
+    fig, axs = plt.subplots(len(lr_schedulers), 1, figsize=(5, 3 * len(lr_schedulers)))
+    if len(lr_schedulers) == 1:
+        axs = [axs]
+
+    for i, (scheduler, title) in enumerate(zip(lr_schedulers, titles)):
+        lrs = simulate_scheduler(scheduler, steps)
+        axs[i].plot(lrs, label=title)
+        axs[i].set_title(title)
+        axs[i].set_xlabel("Steps")
+        axs[i].set_ylabel("Learning Rate")
+        axs[i].legend()
+        axs[i].grid(True)
+
+    plt.tight_layout()
+    if show_plot:
+        plt.show()
+    return fig, axs
+
+
+# Example usage
+if __name__ == "__main__":
+    from fusion_bench.optim.lr_scheduler.linear_warmup import (
+        CosineDecayWithWarmup,
+        LinearWarmupScheduler,
+        PolySchedulerWithWarmup,
+    )
+
+    # Dummy optimizer
+    optimizer = torch.optim.SGD(
+        [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))], lr=0.1
+    )
+
+    # Define the schedulers
+    linear_scheduler = LinearWarmupScheduler(
+        optimizer, t_max=100, max_lr=0.1, min_lr=0.01, init_lr=0.0, warmup_steps=10
+    )
+    cosine_scheduler = CosineDecayWithWarmup(
+        optimizer, t_max=100, max_lr=0.1, min_lr=0.01, init_lr=0.0, warmup_steps=10
+    )
+    poly_scheduler = PolySchedulerWithWarmup(
+        optimizer,
+        t_max=100,
+        max_lr=0.1,
+        min_lr=0.01,
+        init_lr=0.0,
+        warmup_steps=40,
+        poly_order=2.0,
+    )
+
+    # Plot the learning rates
+    plot_lr_schedulers(
+        [linear_scheduler, cosine_scheduler, poly_scheduler],
+        steps=100,
+        titles=[
+            "Linear Warmup",
+            "Cosine Decay with Warmup",
+            "Poly Scheduler with Warmup",
+        ],
+    )

fusion_bench/optim/mezo.py

@@ -5,8 +5,6 @@ import numpy as np
 import torch
 from torch.optim.optimizer import Optimizer
 
-from fusion_bench.utils import timeit_context
-
 log = logging.getLogger(__name__)
 
 

fusion_bench/programs/fabric_fusion_program.py

@@ -185,10 +185,13 @@ class FabricModelFusionProgram(
             report = taskpool.evaluate(merged_model)
             return report
         elif isinstance(merged_model, Dict):
-            model = merged_model.pop("model")
-            report: dict = taskpool.evaluate(model)
-            report.update(merged_model)
-            print(report)
+            report = {}
+            for key, item in merged_model.items():
+                if isinstance(item, nn.Module):
+                    report[key] = taskpool.evaluate(item)
+                else:
+                    # metadata
+                    report[key] = item
             return report
         elif isinstance(merged_model, Iterable):
             return [
@@ -236,7 +239,11 @@ class FabricModelFusionProgram(
             self.save_merged_model(merged_model)
             if self.taskpool is not None:
                 report = self.evaluate_merged_model(self.taskpool, merged_model)
-                print_json(report, print_type=False)
+                try:
+                    print_json(report, print_type=False)
+                except Exception as e:
+                    log.warning(f"Failed to pretty print the report: {e}")
+                    print(report)
                 if self.report_save_path is not None:
                     # save report (Dict) to a file
                     # if the directory of `save_report` does not exists, create it

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -3,7 +3,17 @@ import json
 import logging
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast  # noqa: F401
+from typing import (  # noqa: F401
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
 
 import torch
 from omegaconf import DictConfig
@@ -25,6 +35,10 @@ from fusion_bench.tasks.clip_classification import get_classnames_and_templates
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.parameters import count_parameters
 
+if TYPE_CHECKING:
+    from fusion_bench.models.surgery.surgerymodelwrapper import SurgeryModelWrapper
+
+# disable tokenizers parallelism by default to avoid deadlocks
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 log = logging.getLogger(__name__)
@@ -198,14 +212,16 @@ class CLIPVisionModelTaskPool(
         classifier: HFCLIPClassifier,
         test_loader: DataLoader,
         num_classes: int,
+        task_name: str = None,
     ):
         """
-        Evaluate the classifier on the test dataset.
+        Evaluate the classifier on the test dataset (single-task evaluation).
 
         Args:
             classifier (HFCLIPClassifier): The classifier to evaluate.
             test_loader (DataLoader): The data loader for the test dataset.
             num_classes (int): The number of classes in the classification task.
+            task_name (str): The name of the task.
 
         Returns:
             Dict[str, float]: A dictionary containing the accuracy and loss of the classifier on the test dataset.
@@ -228,7 +244,12 @@ class CLIPVisionModelTaskPool(
             )
         ):
             inputs, targets = batch
-            outputs = classifier(inputs, return_image_embeds=True, return_dict=True)
+            outputs = classifier(
+                inputs,
+                return_image_embeds=True,
+                return_dict=True,
+                task_name=task_name,
+            )
             logits: Tensor = outputs["logits"]
 
             loss = F.cross_entropy(logits, targets)
@@ -246,12 +267,18 @@ class CLIPVisionModelTaskPool(
         results = {"accuracy": acc, "loss": loss}
         return results
 
-    def evaluate(self, model: Union[CLIPVisionModel, CLIPVisionTransformer], name=None):
+    def evaluate(
+        self,
+        model: Union[CLIPVisionModel, CLIPVisionTransformer],
+        name=None,
+        **kwargs,
+    ):
         """
         Evaluate the model on the image classification task.
 
         Args:
             model (Union[CLIPVisionModel, CLIPVisionTransformer]): The model to evaluate.
+            name (Optional[str]): The name of the model. This will be logged into the report if not None.
 
         Returns:
             Dict[str, Any]: A dictionary containing the evaluation results for each task.
@@ -261,8 +288,17 @@ class CLIPVisionModelTaskPool(
 
         report = {}
         # CLIPVisionModel works the same with CLIPVisonTransformer, so we can use it directly
-        self.clip_model.vision_model = model
-        classifier = HFCLIPClassifier(self.clip_model, processor=self.processor)
+        if hasattr(model, "is_surgery_model") and model.is_surgery_model:
+            log.info("running evaluation on a surgery model.")
+            model: "SurgeryModelWrapper" = model
+            self.clip_model.vision_model = model
+        else:
+            # replace the vision encoder with the model
+            self.clip_model.vision_model = model
+        classifier = HFCLIPClassifier(
+            self.clip_model,
+            processor=self.processor,
+        )
         classifier = cast(HFCLIPClassifier, self.fabric.to_device(classifier))
         # collect basic model information
         training_params, all_params = count_parameters(model)
@@ -285,6 +321,7 @@ class CLIPVisionModelTaskPool(
                 classifier,
                 test_dataloader,
                 num_classes=len(classnames),
+                task_name=task_name,
             )
             report[task_name] = result
             self.on_task_evaluation_end()

fusion_bench/taskpool/llama/reward_model.py

@@ -0,0 +1,157 @@
+"""
+The dataset contains the following fields:
+
+- chosen_input_ids: The input token ids for the winner.
+- chosen_attention_mask: The attention mask for the winner.
+- rejected_input_ids: The input token ids for the loser.
+- rejected_attention_mask: The attention mask for the loser.
+"""
+
+import functools
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
+
+import lightning as L
+import numpy as np
+import torch
+from omegaconf import DictConfig
+from torch.utils.data import Subset
+from tqdm.auto import tqdm
+
+from fusion_bench.dataset.llama.collate import bradley_terry_rm_collate
+from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.taskpool import BaseTaskPool
+from fusion_bench.utils import instantiate
+
+if TYPE_CHECKING:
+    from transformers import LlamaForSequenceClassification
+
+
+def evaluate_batch(model: "LlamaForSequenceClassification", batch):
+    batch_size = batch["input_ids"].size(0)
+    assert batch_size % 2 == 0, "Batch size must be even."
+
+    outputs = model(
+        input_ids=batch["input_ids"],
+        attention_mask=batch["attention_mask"],
+    )
+
+    rewards = outputs[0]
+    chosen_reward = rewards[: batch_size // 2]
+    rejected_rewards = rewards[batch_size // 2 :]
+
+    loss = -torch.log(torch.sigmoid(chosen_reward - rejected_rewards)).mean()
+    correct = (chosen_reward > rejected_rewards).sum().item()
+    total = batch_size // 2
+
+    return {
+        "loss": loss.item(),
+        "correct": correct,
+        "total": total,
+    }
+
+
+def evaluate_dataloader(model: "LlamaForSequenceClassification", dataloader):
+    """
+    Compute the accuracy of the reward model on the given dataloader.
+
+    Args:
+        model: The reward model
+        dataloader: The dataloader for the dataset
+
+    Returns:
+        float: The accuracy of the reward model on the dataset
+    """
+    metrics = {
+        "loss": 0.0,
+        "correct": 0,
+        "total": 0,
+    }
+    with torch.no_grad():
+        for batch in (pbar := tqdm(dataloader)):
+            batch_result = evaluate_batch(model, batch)
+            new_total = metrics["total"] + batch_result["total"]
+            metrics["loss"] = (
+                metrics["loss"] * metrics["total"] / new_total
+                + batch_result["loss"] * batch_result["total"] / new_total
+            )
+            metrics["correct"] += batch_result["correct"]
+            metrics["total"] += batch_result["total"]
+            pbar.set_postfix(metrics)
+
+    metrics["accuracy"] = metrics["correct"] / metrics["total"]
+    return metrics
+
+
+class RewardModelEvaluationTaskPool(
+    BaseTaskPool,
+    LightningFabricMixin,
+):
+    def __init__(
+        self,
+        test_datasets: List[DictConfig],
+        dataloader_kwargs: DictConfig,
+        tokenizer: Optional[DictConfig],
+        max_num_samples: int = -1,
+        seed: int = 0,
+        **kwargs,
+    ):
+        self.seed = seed
+        L.seed_everything(seed)
+        self._test_datasets = test_datasets
+        self.dataloader_kwargs = dataloader_kwargs
+        self._tokenizer = tokenizer
+        self.max_num_samples = max_num_samples
+        super().__init__(**kwargs)
+
+    def setup(self):
+        if self._tokenizer is None:
+            # try to load the tokenizer from the model pool
+            tokenizer = self._program.modelpool.load_tokenizer()
+        else:
+            tokenizer = instantiate(self._tokenizer)
+        self.tokenizer = tokenizer
+
+        test_datasets = {
+            dataset_name: instantiate(self._test_datasets[dataset_name])
+            for dataset_name in self._test_datasets
+        }
+        if self.max_num_samples > 0:
+            test_datasets = {
+                dataset_name: Subset(
+                    test_dataset,
+                    np.random.permutation(len(test_dataset))[: self.max_num_samples],
+                )
+                for dataset_name, test_dataset in test_datasets.items()
+            }
+        test_dataloaders = {
+            dataset_name: torch.utils.data.DataLoader(
+                test_dataset,
+                collate_fn=functools.partial(
+                    bradley_terry_rm_collate,
+                    pad_token_id=tokenizer.pad_token_id,
+                ),
+                **self.dataloader_kwargs,
+            )
+            for dataset_name, test_dataset in test_datasets.items()
+        }
+
+        self.test_dataloaders = {
+            dataset_name: self.fabric.setup_dataloaders(test_dataloader)
+            for dataset_name, test_dataloader in test_dataloaders.items()
+        }
+
+    @torch.no_grad()
+    def evaluate(self, model: "LlamaForSequenceClassification"):
+        self.setup()
+
+        model = self.fabric.setup_module(model)
+        if model.config.pad_token_id is None:
+            model.config.pad_token_id = self.tokenizer.pad_token_id
+
+        model.eval()
+        report = {}
+        for dataset_name, test_dataloader in self.test_dataloaders.items():
+            report[dataset_name] = evaluate_dataloader(model, test_dataloader)
+
+        print(report)
+        return report