fusion-bench 0.2.22__py3-none-any.whl → 0.2.23__py3-none-any.whl

fusion_bench/models/we_moe.py

@@ -1,6 +1,6 @@
 import functools
 import logging
-from typing import List
+from typing import Generic, List
 
 import torch
 import torch.func
@@ -9,7 +9,7 @@ from torch.func import functional_call
 from torch.nn import functional as F
 
 from fusion_bench.models.utils import del_attr, get_attr, set_attr
-from fusion_bench.utils.type import StateDictType
+from fusion_bench.utils.type import StateDictType, TorchModelType
 
 log = logging.getLogger(__name__)
 
@@ -76,15 +76,15 @@ def construct_weight_ensembling_gate(
     return gate
 
 
-class WeightEnsemblingMoE(nn.Module):
+class WeightEnsemblingMoE(nn.Module, Generic[TorchModelType]):
     # variable to store the merged state dict temporarily
     _merged_state_dict: StateDictType = None
 
     def __init__(
         self,
         hidden_size: int,
-        base_model: nn.Module,
-        expert_models: List[nn.Module],
+        base_model: TorchModelType,
+        expert_models: List[TorchModelType],
         init_lambda: float = 0.2,
         batch_first: bool = False,
         router_hidden_layers: int = 2,
@@ -101,8 +101,8 @@ class WeightEnsemblingMoE(nn.Module):
         Args:
 
             hidden_size (int): The size of the hidden layer in the models.
-            base_model (nn.Module): The base model that will be used as a reference for the expert models.
-            expert_models (List[nn.Module]): A list of expert models that will be combined.
+            base_model (TorchModelType): The base model that will be used as a reference for the expert models.
+            expert_models (List[TorchModelType]): A list of expert models that will be combined.
             init_lambda (float, optional): The initial lambda value for the weight ensembling gate. Defaults to 0.2.
             batch_first (bool, optional): If True, the input tensors are expected to have the batch size as the first dimension. Defaults to False.
             router_hidden_layers (int, optional): The number of hidden layers in the router. Defaults to 2.
@@ -145,7 +145,7 @@ class WeightEnsemblingMoE(nn.Module):
             self._merged_state_dict,
         )
 
-    def merge_weights(self, expert_weights):
+    def merge_weights(self, expert_weights) -> StateDictType:
         state_dict = self.base_model.state_dict(keep_vars=True)
         for weight, task_vector in zip(expert_weights, self.task_vectors):
             for name, param in task_vector.named_parameters():

fusion_bench/taskpool/base_pool.py

@@ -5,33 +5,115 @@ from fusion_bench.mixins import BaseYAMLSerializable
 
 
 class BaseTaskPool(BaseYAMLSerializable):
+    """Abstract base class for task pools in the FusionBench framework.
+
+    A task pool represents a collection of evaluation tasks that can be used to
+    assess model performance across multiple benchmarks or datasets. This base
+    class defines the common interface that all task pool implementations must
+    follow, ensuring consistency across different task types and evaluation
+    scenarios.
+
+    Task pools are designed to be configurable through YAML files and can be
+    used in various model fusion and evaluation workflows. They provide a
+    standardized way to evaluate models on multiple tasks and aggregate results.
+
+    The class inherits from BaseYAMLSerializable to support configuration
+    management and serialization capabilities.
+
+    Attributes:
+        _program: Optional program reference for execution context.
+        _config_key: Configuration key used for YAML configuration ("taskpool").
+
+    Abstract Methods:
+        evaluate: Must be implemented by subclasses to define task-specific
+            evaluation logic.
+
+    Example:
+        Implementing a custom task pool:
+
+        ```python
+        class MyTaskPool(BaseTaskPool):
+
+
+            def evaluate(self, model, **kwargs):
+                results = {}
+                for task_name in self.tasks:
+                    # Implement task-specific evaluation
+                    results[task_name] = self._evaluate_task(model, task_name)
+                return results
+        ```
+    """
+
     _program = None
     _config_key = "taskpool"
 
     @abstractmethod
     def evaluate(self, model: Any, *args: Any, **kwargs: Any) -> Dict[str, Any]:
-        """
-        Evaluate the model on all tasks in the task pool, and return a report.
+        """Evaluate a model on all tasks in the task pool and return aggregated results.
 
-        Take image classification as an example, the report will look like:
+        This abstract method defines the core evaluation interface that all task pool
+        implementations must provide. It should evaluate the given model on all tasks
+        managed by the pool and return a structured report of the results.
 
-        ```python
-        {
-            "mnist": {
-                "accuracy": 0.8,
-                "loss": 0.2,
-            },
-            <task_name>: {
-                <metric_name>: <metric_value>,
-                ...
-            },
-        }
-        ```
+        The evaluation process typically involves:
+        1. Iterating through all tasks in the pool
+        2. Running model inference on each task's dataset
+        3. Computing task-specific metrics
+        4. Aggregating results into a standardized report format
 
         Args:
-            model: The model to evaluate.
+            model: The model to evaluate. Can be any model type (PyTorch model,
+                Hugging Face model, etc.) that is compatible with the specific
+                task pool implementation.
+            *args: Additional positional arguments that may be needed for
+                task-specific evaluation procedures.
+            **kwargs: Additional keyword arguments for evaluation configuration,
+                such as batch_size, device, evaluation metrics, etc.
 
         Returns:
-            report (dict): A dictionary containing the results of the evaluation for each task.
+            Dict[str, Any]: A dictionary containing evaluation results for each task.
+                The structure follows the pattern:
+
+                ```python
+                {
+                    "task_name_1": {
+                        "metric_1": value,
+                        "metric_2": value,
+                        ...
+                    },
+                    "task_name_2": {
+                        "metric_1": value,
+                        "metric_2": value,
+                        ...
+                    },
+                    ...
+                }
+                ```
+
+        Example:
+            For an image classification task pool:
+
+            ```python
+            results = task_pool.evaluate(model)
+            # Returns:
+            # {
+            #     "mnist": {
+            #         "accuracy": 0.95,
+            #         "loss": 0.15,
+            #     },
+            #     "cifar10": {
+            #         "accuracy": 0.87,
+            #         "loss": 0.42,
+            #     }
+            # }
+            ```
+
+        Raises:
+            NotImplementedError: This method must be implemented by subclasses.
+
+        Note:
+            Implementations should ensure that the returned dictionary structure
+            is consistent and that metric names are standardized across similar
+            task types to enable meaningful comparison and aggregation.
         """
         pass

fusion_bench/taskpool/clip_vision/taskpool.py

@@ -309,7 +309,7 @@ class CLIPVisionModelTaskPool(
             self.setup()
 
         report = {}
-        # CLIPVisionModel works the same with CLIPVisonTransformer, so we can use it directly
+        # CLIPVisionModel works the same with CLIPVisionTransformer, so we can use it directly
         if hasattr(model, "is_surgery_model") and model.is_surgery_model:
             log.info("running evaluation on a surgery model.")
             model: "SurgeryModelWrapper" = model

fusion_bench/taskpool/dummy.py

@@ -1,5 +1,10 @@
 """
-This is the dummy task pool that is used for debugging purposes.
+Dummy task pool implementation for debugging and testing purposes.
+
+This module provides a minimal task pool implementation that can be used for
+debugging model fusion workflows, testing infrastructure, and validating model
+architectures without running expensive evaluation procedures. It's particularly
+useful during development and prototyping phases.
 """
 
 from typing import Optional
@@ -14,14 +19,41 @@ from fusion_bench.utils.parameters import count_parameters, print_parameters
 
 
 def get_model_summary(model: nn.Module) -> dict:
-    """
-    Generate a report for the given model.
+    """Generate a comprehensive summary report for a PyTorch model.
+
+    Analyzes the given model to extract key information about its architecture,
+    parameter count, and training characteristics. This function is useful for
+    model introspection and comparative analysis during model fusion workflows.
+
+    The summary includes both trainable and total parameter counts, which helps
+    in understanding model complexity and memory requirements. The trainable
+    percentage is particularly useful for identifying models with frozen layers
+    or parameter-efficient fine-tuning setups.
 
     Args:
-        model: The model to generate the report for.
+        model: The PyTorch model to analyze. Can be any nn.Module instance
+            including complex models, fusion models, or pre-trained models.
 
     Returns:
-        dict: The generated report.
+        dict: A structured report containing model information:
+            - model_info: Dictionary with parameter statistics
+                - trainable_params: Number of trainable parameters
+                - all_params: Total number of parameters (trainable + frozen)
+                - trainable_percentage: Ratio of trainable to total parameters
+
+    Example:
+        ```python
+        >>> model = MyModel()
+        >>> summary = get_model_summary(model)
+        >>> print(summary)
+        {
+            "model_info": {
+                "trainable_params": 1234567,
+                "all_params": 1234567,
+                "trainable_percentage": 1.0
+            }
+        }
+        ```
     """
     report = {}
     training_params, all_params = count_parameters(model)
@@ -34,21 +66,77 @@ def get_model_summary(model: nn.Module) -> dict:
 
 
 class DummyTaskPool(BaseTaskPool):
+    """A lightweight task pool implementation for debugging and development workflows.
+
+    This dummy task pool provides a minimal evaluation interface that focuses on
+    model introspection rather than task-specific performance evaluation. It's
+    designed for development scenarios where you need to test model fusion
+    pipelines, validate architectures, or debug workflows without the overhead
+    of running actual evaluation tasks.
+
+    The task pool is particularly useful when:
+        - You want to verify model fusion works correctly
+        - You need to check parameter counts after fusion
+        - You're developing new fusion algorithms
+        - You want to test infrastructure without expensive evaluations
+
+    Example:
+        ```python
+        >>> taskpool = DummyTaskPool(model_save_path="/tmp/fused_model")
+        >>> results = taskpool.evaluate(fused_model)
+        >>> print(f"Model has {results['model_info']['trainable_params']} parameters")
+        ```
     """
-    This is a dummy task pool used for debugging purposes. It inherits from the base TaskPool class.
-    """
 
-    def __init__(self, model_save_path: Optional[str] = None):
-        super().__init__()
+    def __init__(self, model_save_path: Optional[str] = None, **kwargs):
+        """Initialize the dummy task pool with optional model saving capability.
+
+        Args:
+            model_save_path: Optional path where the evaluated model should be saved.
+                If provided, the model will be serialized and saved to this location
+                after evaluation using the separate_save utility. If None, no model
+                saving will be performed.
+
+        Example:
+            ```python
+            >>> # Create taskpool without saving
+            >>> taskpool = DummyTaskPool()
+
+            >>> # Create taskpool with model saving
+            >>> taskpool = DummyTaskPool(model_save_path="/path/to/save/model.pth")
+            ```
+        """
+        super().__init__(**kwargs)
         self.model_save_path = model_save_path
 
     def evaluate(self, model):
-        """
-        Evaluate the given model.
-        This method does nothing but print the parameters of the model in a human-readable format.
+        """Perform lightweight evaluation and analysis of the given model.
+
+        This method provides a minimal evaluation that focuses on model introspection
+        rather than task-specific performance metrics. It performs parameter analysis,
+        optionally saves the model, and returns a summary report.
+
+        The evaluation process includes:
+        1. Printing human-readable parameter information (rank-zero only)
+        2. Optionally saving the model if a save path was configured
+        3. Generating and returning a model summary report
 
         Args:
-            model: The model to evaluate.
+            model: The model to evaluate. Can be any PyTorch nn.Module including
+                fusion models, pre-trained models, or custom architectures.
+
+        Returns:
+            dict: A model summary report containing parameter statistics and
+                architecture information. See get_model_summary() for detailed
+                format specification.
+
+        Example:
+            ```python
+            >>> taskpool = DummyTaskPool(model_save_path="/tmp/model.pth")
+            >>> model = torch.nn.Linear(10, 5)
+            >>> results = taskpool.evaluate(model)
+            >>> print(f"Trainable params: {results['model_info']['trainable_params']}")
+            ```
         """
         if rank_zero_only.rank == 0:
             print_parameters(model, is_human_readable=True)

fusion_bench/taskpool/lm_eval_harness/taskpool.py

@@ -16,6 +16,47 @@ log = logging.getLogger(__name__)
 
 
 class LMEvalHarnessTaskPool(BaseTaskPool, LightningFabricMixin):
+    """A task pool implementation that interfaces with the LM Evaluation Harness framework.
+
+    This class provides a wrapper around the LM Evaluation Harness (lm-eval) library,
+    enabling evaluation of language models on various standardized benchmarks and tasks.
+    It inherits from BaseTaskPool and LightningFabricMixin to provide distributed
+    computing capabilities through PyTorch Lightning Fabric.
+
+    The task pool supports evaluation on multiple tasks simultaneously and provides
+    flexible configuration options for batch processing, output formatting, and
+    logging. It automatically handles model setup and wrapping for distributed
+    evaluation when using Lightning Fabric.
+
+    Args:
+        tasks: A single task name or list of task names to evaluate on.
+            Examples: "hellaswag", ["arc_easy", "arc_challenge", "hellaswag"]
+        apply_chat_template: Whether to apply chat template formatting to inputs.
+            Useful for instruction-tuned or chat models.
+        include_path: Path to additional task definitions or custom tasks.
+        batch_size: Number of samples to process in each batch. Larger values
+            may improve throughput but require more memory.
+        metadata: Additional metadata to include in evaluation results.
+        verbosity: Logging verbosity level for the evaluation process.
+        output_path: Custom path for saving evaluation results. If None,
+            results are saved to the default log directory.
+        log_samples: Whether to log individual sample predictions and targets.
+            Useful for debugging but increases output size significantly.
+        _usage_: Internal usage tracking string.
+        _version_: Internal version tracking string.
+        **kwargs: Additional arguments passed to the LM Evaluation Harness.
+
+    Example:
+        ```python
+        >>> taskpool = LMEvalHarnessTaskPool(
+        ...     tasks=["arc_easy", "hellaswag"],
+        ...     batch_size=8,
+        ...     verbosity="INFO"
+        ... )
+        >>> results = taskpool.evaluate(model)
+        ```
+    """
+
     def __init__(
         self,
         tasks: Union[str, List[str]],
@@ -44,6 +85,45 @@ class LMEvalHarnessTaskPool(BaseTaskPool, LightningFabricMixin):
         self.log_samples = log_samples
 
     def evaluate(self, model, *command_line_args, **kwargs):
+        """Evaluate a language model on the configured tasks using LM Evaluation Harness.
+
+        This method wraps the model with the LM Evaluation Harness framework and
+        executes evaluation on all configured tasks. It automatically handles
+        command-line argument construction, model wrapping with Lightning Fabric
+        for distributed evaluation, and result logging.
+
+        The evaluation process includes:
+        1. Building command-line arguments from instance configuration
+        2. Setting up the LM Evaluation Harness argument parser
+        3. Wrapping the model with Lightning Fabric if not already wrapped
+        4. Creating an HFLM (Hugging Face Language Model) wrapper
+        5. Executing the evaluation through the LM-Eval CLI interface
+
+        Args:
+            model: The language model to evaluate. Can be a Hugging Face model,
+                PyTorch model, or any model compatible with the LM Evaluation Harness.
+                The model will be automatically wrapped with Lightning Fabric for
+                distributed evaluation if not already wrapped.
+            *command_line_args: Additional positional command-line arguments
+                (currently unused but preserved for interface compatibility).
+            **kwargs: Additional keyword arguments that will be converted to
+                command-line flags and passed to the LM Evaluation Harness.
+                Keys will be prefixed with '--' and values converted to strings.
+
+        Returns:
+            None: Results are written to the configured output path and logged.
+
+        Example:
+            ```python
+            >>> taskpool = LMEvalHarnessTaskPool(tasks=["arc_easy"])
+            >>> taskpool.evaluate(model, limit=100, device="cuda")
+            ```
+
+        Note:
+            The method leverages the LM Evaluation Harness's command-line interface
+            internally, which provides standardized evaluation procedures and
+            ensures compatibility with the broader evaluation ecosystem.
+        """
         command_line_args = []
         if self.include_path is not None:
             command_line_args.extend(["--include_path", self.include_path])

fusion_bench/taskpool/nyuv2_taskpool.py

@@ -15,9 +15,37 @@ log = logging.getLogger(__name__)
 
 
 class NYUv2TaskPool(TaskPool):
+    """Task pool for multi-task learning evaluation on the NYUv2 dataset.
+
+    This task pool provides evaluation capabilities for multi-task learning models
+    on the NYU Depth V2 (NYUv2) dataset, which is a popular benchmark for indoor
+    scene understanding. The dataset supports multiple computer vision tasks
+    including semantic segmentation, depth estimation, and surface normal prediction.
+
+    The task pool is designed to work with encoder-decoder architectures where
+    a shared encoder processes input images and task-specific decoders generate
+    predictions for different tasks. It integrates with PyTorch Lightning for
+    streamlined training and evaluation workflows.
+
+    Supported Tasks:
+        - Semantic segmentation
+        - Depth estimation
+        - Surface normal prediction
+    """
+
     _trainer: L.Trainer = None
 
     def __init__(self, taskpool_config: DictConfig):
+        """Initialize the NYUv2 task pool with configuration settings.
+
+        Args:
+            taskpool_config: Configuration object containing all necessary
+                parameters for the task pool, including:
+                - data_dir: Path to the directory containing NYUv2 dataset
+                - tasks: List of tasks to evaluate (e.g., ["semantic", "depth"])
+                - batch_size: Batch size for evaluation data loader
+                - num_workers: Number of worker processes for data loading
+        """
         self.config = taskpool_config
 
     def load_datasets(self):

fusion_bench/utils/__init__.py

@@ -20,3 +20,4 @@ from .packages import import_object
 from .parameters import *
 from .pylogger import get_rankzero_logger
 from .timer import timeit_context
+from .type import BoolStateDictType, StateDictType, TorchModelType

fusion_bench/utils/data.py

@@ -1,6 +1,6 @@
 import pickle
 from pathlib import Path
-from typing import Literal, Optional, Union
+from typing import Any, Literal, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -37,7 +37,9 @@ class InfiniteDataLoader:
         return data
 
 
-def load_tensor_from_file(file_path: Union[str, Path], device=None) -> torch.Tensor:
+def load_tensor_from_file(
+    file_path: Union[str, Path], device: Optional[Union[str, torch.device]] = None
+) -> torch.Tensor:
     """
     Loads a tensor from a file, which can be either a .pt, .pth or .np file.
     If the file is not one of these formats, it will try to load it as a pickle file.
@@ -72,7 +74,7 @@ def train_validation_split(
     validation_size: Optional[int] = None,
     random_seed: Optional[int] = None,
     return_split: Literal["all", "train", "val"] = "both",
-):
+) -> Union[Tuple[Dataset, Dataset], Dataset]:
     """
     Split a dataset into a training and validation set.
 
@@ -134,7 +136,7 @@ def train_validation_test_split(
     test_fraction: float,
     random_seed: Optional[int] = None,
     return_spilt: Literal["all", "train", "val", "test"] = "all",
-):
+) -> Union[Tuple[Dataset, Dataset, Dataset], Dataset]:
     """
     Split a dataset into a training, validation and test set.
 

fusion_bench/utils/devices.py

@@ -1,7 +1,7 @@
 import gc
 import logging
 import os
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import torch
 from transformers.utils import (
@@ -12,6 +12,8 @@ from transformers.utils import (
     is_torch_xpu_available,
 )
 
+from .type import T
+
 __all__ = [
     "clear_cuda_cache",
     "to_device",
@@ -37,7 +39,7 @@ def clear_cuda_cache():
         log.warning("CUDA is not available. No cache to clear.")
 
 
-def to_device(obj, device: Optional[torch.device], **kwargs):
+def to_device(obj: T, device: Optional[torch.device], **kwargs: Any) -> T:
     """
     Move a given object to the specified device.
 
@@ -102,7 +104,7 @@ def num_devices(devices: Union[int, List[int], str]) -> int:
         )
 
 
-def get_device(obj) -> torch.device:
+def get_device(obj: Any) -> torch.device:
     """
     Get the device of a given object.
 
@@ -151,6 +153,7 @@ def get_current_device() -> torch.device:
                     If not set, it defaults to "0".
 
     Example:
+
         >>> device = get_current_device()
         >>> print(device)
         xpu:0  # or npu:0, mps:0, cuda:0, cpu depending on availability
@@ -241,7 +244,7 @@ def cleanup_cuda():
     torch.cuda.reset_peak_memory_stats()
 
 
-def print_memory_usage(print_fn=print):
+def print_memory_usage(print_fn=print) -> str:
     """
     Print the current GPU memory usage.
 

fusion_bench/utils/dtype.py

@@ -1,5 +1,5 @@
 import contextlib
-from typing import Dict, Generator, Iterable, Optional, Tuple
+from typing import Dict, Generator, Iterable, Optional, Tuple, Union
 
 import torch
 from transformers.utils import (
@@ -25,7 +25,7 @@ PRECISION_STR_TO_DTYPE: Dict[str, torch.dtype] = {
 }
 
 
-def parse_dtype(dtype: Optional[str]):
+def parse_dtype(dtype: Optional[str]) -> Optional[torch.dtype]:
     """
     Parses a string representation of a data type and returns the corresponding torch.dtype.
 
@@ -92,6 +92,7 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
         ContextManager: context manager for setting default dtype.
 
     Example:
+
         >>> with set_default_dtype(torch.bfloat16):
         >>>     x = torch.tensor([1, 2, 3])
         >>>     x.dtype