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

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

fusion_bench/utils/fabric.py

@@ -3,9 +3,9 @@ from typing import Optional
 
 import lightning as L
 
-from fusion_bench.utils.pylogger import getRankZeroLogger
+from fusion_bench.utils.pylogger import get_rankzero_logger
 
-log = getRankZeroLogger(__name__)
+log = get_rankzero_logger(__name__)
 
 
 def seed_everything_by_time(fabric: Optional[L.Fabric] = None):

fusion_bench/utils/lazy_imports.py

@@ -72,3 +72,26 @@ class LazyImporter(ModuleType):
 
     def __reduce__(self):
         return (self.__class__, (self._name, self.__file__, self._import_structure))
+
+
+class LazyModule(ModuleType):
+    """Module wrapper for lazy import.
+    Adapted from Optuna: https://github.com/optuna/optuna/blob/1f92d496b0c4656645384e31539e4ee74992ff55/optuna/__init__.py
+
+    This class wraps specified module and lazily import it when they are actually accessed.
+
+    Args:
+        name: Name of module to apply lazy import.
+    """
+
+    def __init__(self, name: str) -> None:
+        super().__init__(name)
+        self._name = name
+
+    def _load(self) -> ModuleType:
+        module = importlib.import_module(self._name)
+        self.__dict__.update(module.__dict__)
+        return module
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._load(), item)

fusion_bench/utils/lazy_state_dict.py

@@ -2,7 +2,18 @@ import json
 import logging
 import os
 from copy import deepcopy
-from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 import torch
 from accelerate import init_empty_weights
@@ -11,10 +22,12 @@ from huggingface_hub import snapshot_download
 from safetensors import safe_open
 from safetensors.torch import load_file
 from torch import nn
+from torch.nn.modules.module import _IncompatibleKeys
 from transformers import AutoConfig
 
 from fusion_bench.utils.dtype import parse_dtype
 from fusion_bench.utils.packages import import_object
+from fusion_bench.utils.type import TorchModelType
 
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
@@ -49,7 +62,7 @@ def resolve_checkpoint_path(
         )
 
 
-class LazyStateDict:
+class LazyStateDict(Mapping[str, torch.Tensor], Generic[TorchModelType]):
     """
     Dictionary-like object that lazily loads a state dict from a checkpoint path.
     """
@@ -66,8 +79,8 @@ class LazyStateDict:
     def __init__(
         self,
         checkpoint: str,
-        meta_module_class: Optional[Type[nn.Module]] = None,
-        meta_module: Optional[nn.Module] = None,
+        meta_module_class: Optional[Type[TorchModelType]] = None,
+        meta_module: Optional[TorchModelType] = None,
         cache_state_dict: bool = False,
         torch_dtype: Optional[torch.dtype] = None,
         device: str = "cpu",
@@ -88,15 +101,19 @@ class LazyStateDict:
             hf_proxies (Dict, optional): Proxies to use for downloading from Hugging Face Hub.
         """
         self.cache_state_dict = cache_state_dict
+
+        # Validate that both meta_module_class and meta_module are not provided
+        if meta_module_class is not None and meta_module is not None:
+            raise ValueError(
+                "Cannot provide both meta_module_class and meta_module, please provide only one."
+            )
+
         self.meta_module_class = meta_module_class
         if isinstance(self.meta_module_class, str):
             self.meta_module_class = import_object(self.meta_module_class)
         self.meta_module = meta_module
+
         if self.meta_module_class is not None:
-            if self.meta_module is not None:
-                raise ValueError(
-                    "Cannot provide both meta_module_class and meta_module, please provide only one."
-                )
             with init_empty_weights():
                 self.meta_module = self.meta_module_class.from_pretrained(
                     checkpoint,
@@ -168,12 +185,25 @@ class LazyStateDict:
     def config(self) -> "PretrainedConfig":
         return AutoConfig.from_pretrained(self._checkpoint)
 
+    @property
+    def dtype(self) -> torch.dtype:
+        """
+        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
+        """
+        if hasattr(self, "_cached_dtype"):
+            return self._cached_dtype
+
+        first_key = next(iter(self.keys()))
+        first_param = self[first_key]
+        self._cached_dtype = first_param.dtype
+        return self._cached_dtype
+
     def state_dict(self, keep_vars: bool = False) -> "LazyStateDict":
         """
         Args:
             keep_vars (bool): Ignored, as LazyStateDict does not support keep_vars. Just for compatibility.
         """
-        return self
+        return deepcopy(self)
 
     def _resolve_checkpoint_files(self, checkpoint: str):
         # reference: https://huggingface.co/docs/accelerate/v0.17.1/en/usage_guides/big_modeling
@@ -290,6 +320,18 @@ class LazyStateDict:
             )
             return tensor
 
+    def pop(self, key: str):
+        assert key in list(
+            self.keys()
+        ), "KeyError: Cannot pop a tensor for a key that does not exist in the LazyStateDict."
+        if self._state_dict_cache is not None and key in self._state_dict_cache:
+            if key in self._index:
+                self._index.pop(key)
+            return self._state_dict_cache.pop(key)
+        if key in self._index:
+            self._index.pop(key)
+        return None
+
     def __setitem__(self, key: str, value: torch.Tensor) -> None:
         """
         Set a tensor in the LazyStateDict. This will update the state dict cache if it is enabled.
@@ -300,9 +342,7 @@ class LazyStateDict:
         if self._state_dict_cache is not None:
             self._state_dict_cache[key] = value
         else:
-            log.warning(
-                "State dict cache is disabled, setting a tensor will not update the cache."
-            )
+            log.warning("State dict cache is disabled, initializing the cache.")
             self._state_dict_cache = {key: value}
 
     def __contains__(self, key: str) -> bool:
@@ -318,7 +358,7 @@ class LazyStateDict:
                     self._checkpoint_files[0], key, update_cache=False
                 )
                 return tensor is not None
-            except Exception:
+            except (KeyError, FileNotFoundError, RuntimeError, EOFError):
                 return False
         return False
 
@@ -388,8 +428,8 @@ class LazyStateDict:
             )
 
     def load_state_dict(
-        self, state_dict: Dict[str, torch.Tensor], strict: bool = True
-    ) -> None:
+        self, state_dict: Mapping[str, torch.Tensor], strict: bool = True
+    ) -> _IncompatibleKeys:
         """
         Load a state dict into this LazyStateDict.
         This method is only for compatibility with nn.Module and it overrides the cache of LazyStateDict.
@@ -398,13 +438,71 @@ class LazyStateDict:
             state_dict (Dict[str, torch.Tensor]): The state dict to load.
             strict (bool): Whether to enforce that all keys in the state dict are present in this LazyStateDict.
         """
+        if not isinstance(state_dict, Mapping):
+            raise TypeError(
+                f"Expected state_dict to be dict-like, got {type(state_dict)}."
+            )
+
+        missing_keys: list[str] = []
+        unexpected_keys: list[str] = []
+        error_msgs: list[str] = []
+
         log.warning(
             "Loading state dict into LazyStateDict is not recommended, as it may lead to unexpected behavior. "
             "Use with caution."
         )
+
+        # Check for unexpected keys in the provided state_dict
+        for key in state_dict:
+            if key not in self:
+                unexpected_keys.append(key)
+
+        # Check for missing keys that are expected in this LazyStateDict
+        for key in self.keys():
+            if key not in state_dict:
+                missing_keys.append(key)
+
+        # Handle strict mode
         if strict:
-            for key in state_dict:
-                if key not in self:
-                    raise KeyError(f"Key {key} not found in LazyStateDict.")
+            if len(unexpected_keys) > 0:
+                error_msgs.insert(
+                    0,
+                    "Unexpected key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in unexpected_keys)
+                    ),
+                )
+            if len(missing_keys) > 0:
+                error_msgs.insert(
+                    0,
+                    "Missing key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in missing_keys)
+                    ),
+                )
+
+        if len(error_msgs) > 0:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    self.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+
+        # Load the state dict values
         for key, value in state_dict.items():
-            self[key] = value
+            if key in self:  # Only set keys that exist in this LazyStateDict
+                self[key] = value
+
+        return _IncompatibleKeys(missing_keys, unexpected_keys)
+
+    def __getattr__(self, name: str):
+        if "meta_module" in self.__dict__:
+            meta_module = self.__dict__["meta_module"]
+            if meta_module is not None:
+                if "_parameters" in meta_module.__dict__:
+                    if name in meta_module.__dict__["_parameters"]:
+                        return self.get_parameter(name)
+                if "_modules" in meta_module.__dict__:
+                    if name in meta_module.__dict__["_modules"]:
+                        return self.get_submodule(name)
+        raise AttributeError(
+            f"'{type(self).__name__}' object has no attribute '{name}'"
+        )

fusion_bench/utils/modelscope.py

@@ -26,13 +26,13 @@ try:
     from huggingface_hub import snapshot_download as huggingface_snapshot_download
 except ImportError:
 
-    def _raise_hugggingface_not_installed_error(*args, **kwargs):
+    def _raise_huggingface_not_installed_error(*args, **kwargs):
         raise ImportError(
             "Hugging Face Hub is not installed. Please install it using `pip install huggingface_hub` to use Hugging Face models."
         )
 
-    huggingface_snapshot_download = _raise_hugggingface_not_installed_error
-    hf_hub_download = _raise_hugggingface_not_installed_error
+    huggingface_snapshot_download = _raise_huggingface_not_installed_error
+    hf_hub_download = _raise_huggingface_not_installed_error
 
 __all__ = [
     "load_dataset",

fusion_bench/utils/packages.py

@@ -1,7 +1,7 @@
 import importlib.metadata
 import importlib.util
 from functools import lru_cache
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from packaging import version
 
@@ -69,7 +69,7 @@ def is_vllm_available():
     return _is_package_available("vllm")
 
 
-def import_object(abs_obj_name: str):
+def import_object(abs_obj_name: str) -> Any:
     """
     Imports a class from a module given the absolute class name.
 
@@ -84,7 +84,7 @@ def import_object(abs_obj_name: str):
     return getattr(module, obj_name)
 
 
-def compare_versions(v1, v2):
+def compare_versions(v1: str, v2: str) -> int:
     """Compare two version strings.
     Returns -1 if v1 < v2, 0 if v1 == v2, 1 if v1 > v2"""
 

fusion_bench/utils/parameters.py

@@ -129,7 +129,6 @@ def human_readable(num: int) -> str:
     Converts a number into a human-readable string with appropriate magnitude suffix.
 
     Examples:
-
         ```python
         print(human_readable(1500))
         # Output: '1.50K'
@@ -201,7 +200,6 @@ def count_parameters(module: nn.Module, non_zero_only: bool = False) -> tuple[in
         tuple: A tuple containing the number of trainable parameters and the total number of parameters.
 
     Examples:
-
         ```python
         # Count the parameters
         trainable_params, all_params = count_parameters(model)

fusion_bench/utils/path.py

@@ -1,6 +1,9 @@
+import logging
 import os
 from typing import List
 
+log = logging.getLogger(__name__)
+
 
 def path_is_dir_and_not_empty(path: str):
     if path is None:
@@ -20,3 +23,56 @@ def listdir_fullpath(dir: str) -> List[str]:
     assert os.path.isdir(dir), "Argument 'dir' must be a Directory"
     names = os.listdir(dir)
     return [os.path.join(dir, name) for name in names]
+
+
+def create_symlink(src_dir: str, dst_dir: str, link_name: str = None):
+    """
+    Creates a symbolic link from src_dir to dst_dir.
+
+    Args:
+        src_dir (str): The source directory to link to.
+        dst_dir (str): The destination directory where the symlink will be created.
+        link_name (str, optional): The name of the symlink. If None, uses the basename of src_dir.
+
+    Raises:
+        OSError: If the symbolic link creation fails.
+        ValueError: If src_dir does not exist or is not a directory.
+    """
+    if not os.path.exists(src_dir):
+        raise ValueError(f"Source directory does not exist: {src_dir}")
+
+    if not os.path.isdir(src_dir):
+        raise ValueError(f"Source path is not a directory: {src_dir}")
+
+    # Avoid creating symlink if source and destination are the same
+    if os.path.abspath(src_dir) == os.path.abspath(dst_dir):
+        log.warning(
+            "Source and destination directories are the same, skipping symlink creation"
+        )
+        return
+
+    # Create destination directory if it doesn't exist
+    os.makedirs(dst_dir, exist_ok=True)
+
+    # Determine link name
+    if link_name is None:
+        link_name = os.path.basename(src_dir)
+
+    link_path = os.path.join(dst_dir, link_name)
+
+    try:
+        # if the system is windows, use the `mklink` command in "CMD" to create the symlink
+        if os.name == "nt":
+            os.system(
+                f"mklink /J {os.path.abspath(link_path)} {os.path.abspath(src_dir)}"
+            )
+        else:
+            os.symlink(
+                src_dir,
+                link_path,
+                target_is_directory=True,
+            )
+        log.info(f"Created symbolic link: {link_path} -> {src_dir}")
+    except OSError as e:
+        log.warning(f"Failed to create symbolic link: {e}")
+        raise

fusion_bench/utils/pylogger.py

@@ -74,7 +74,7 @@ RankZeroLogger.manager = logging.Manager(RankZeroLogger.root)
 RankZeroLogger.manager.setLoggerClass(RankZeroLogger)
 
 
-def getRankZeroLogger(name=None):
+def get_rankzero_logger(name=None):
     """
     Return a logger with the specified name, creating it if necessary.
 

fusion_bench/utils/timer.py

@@ -6,38 +6,120 @@ log = logging.getLogger(__name__)
 
 class timeit_context:
     """
-    Usage:
+    A context manager for measuring and logging execution time of code blocks.
 
-    ```python
-    with timeit_context() as timer:
-        ... # code block to be measured
-    ```
+    This context manager provides precise timing measurements with automatic logging
+    of elapsed time. It supports nested timing contexts with proper indentation
+    for hierarchical timing analysis, making it ideal for profiling complex
+    operations with multiple sub-components.
+
+    Args:
+        msg (str, optional): Custom message to identify the timed code block.
+            If provided, logs "[BEGIN] {msg}" at start and includes context
+            in the final timing report. Defaults to None.
+        loglevel (int, optional): Python logging level for output messages.
+            Uses standard logging levels (DEBUG=10, INFO=20, WARNING=30, etc.).
+            Defaults to logging.INFO.
+
+    Example:
+        Basic usage:
+        ```python
+        with timeit_context("data loading"):
+            data = load_large_dataset()
+        # Logs: [BEGIN] data loading
+        # Logs: [END]   Elapsed time: 2.34s
+        ```
+
+        Nested timing:
+        ```python
+        with timeit_context("model training"):
+            with timeit_context("data preprocessing"):
+                preprocess_data()
+            with timeit_context("forward pass"):
+                model(data)
+        # Output shows nested structure:
+        # [BEGIN] model training
+        #   [BEGIN] data preprocessing
+        #   [END]   Elapsed time: 0.15s
+        #   [BEGIN] forward pass
+        #   [END]   Elapsed time: 0.89s
+        # [END]   Elapsed time: 1.04s
+        ```
+
+        Custom log level:
+        ```python
+        with timeit_context("debug operation", loglevel=logging.DEBUG):
+            debug_function()
+        ```
     """
 
     nest_level = -1
 
     def _log(self, msg):
+        """
+        Internal method for logging messages with appropriate stack level.
+
+        This helper method ensures that log messages appear to originate from
+        the caller's code rather than from internal timer methods, providing
+        more useful debugging information.
+
+        Args:
+            msg (str): The message to log at the configured log level.
+        """
         log.log(self.loglevel, msg, stacklevel=3)
 
     def __init__(self, msg: str = None, loglevel=logging.INFO) -> None:
+        """
+        Initialize a new timing context with optional message and log level.
+
+        Args:
+            msg (str, optional): Descriptive message for the timed operation.
+                If provided, will be included in the begin/end log messages
+                to help identify what is being timed. Defaults to None.
+            loglevel (int, optional): Python logging level for timer output.
+                Common values include:
+                - logging.DEBUG (10): Detailed debugging information
+                - logging.INFO (20): General information (default)
+                - logging.WARNING (30): Warning messages
+                - logging.ERROR (40): Error messages
+                Defaults to logging.INFO.
+        """
         self.loglevel = loglevel
         self.msg = msg
 
     def __enter__(self) -> None:
         """
-        Sets the start time and logs an optional message indicating the start of the code block execution.
+        Enter the timing context and start the timer.
 
-        Args:
-            msg: str, optional message to log
+        This method is automatically called when entering the 'with' statement.
+        It records the current timestamp, increments the nesting level for
+        proper log indentation, and optionally logs a begin message.
+
+        Returns:
+            None: This context manager doesn't return a value to the 'as' clause.
+                  All timing information is handled internally and logged automatically.
         """
         self.start_time = time.time()
         timeit_context.nest_level += 1
         if self.msg is not None:
             self._log("  " * timeit_context.nest_level + "[BEGIN] " + str(self.msg))
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         """
-        Calculates the elapsed time and logs it, along with an optional message indicating the end of the code block execution.
+        Exit the timing context and log the elapsed time.
+
+        This method is automatically called when exiting the 'with' statement,
+        whether through normal completion or exception. It calculates the total
+        elapsed time and logs the results with proper nesting indentation.
+
+        Args:
+            exc_type (type): Exception type if an exception occurred, None otherwise.
+            exc_val (Exception): Exception instance if an exception occurred, None otherwise.
+            exc_tb (traceback): Exception traceback if an exception occurred, None otherwise.
+
+        Returns:
+            None: Does not suppress exceptions (returns None/False implicitly).
+                  Any exceptions that occurred in the timed block will propagate normally.
         """
         end_time = time.time()
         elapsed_time = end_time - self.start_time

{fusion_bench-0.2.21.dist-info → fusion_bench-0.2.23.dist-info}/METADATA

@@ -1,30 +1,8 @@
 Metadata-Version: 2.4
 Name: fusion_bench
-Version: 0.2.21
+Version: 0.2.23
 Summary: A Comprehensive Benchmark of Deep Model Fusion
 Author-email: Anke Tang <tang.anke@foxmail.com>
-License: MIT License
-        
-        Copyright (c) 2024 Anke Tang
-        
-        Permission is hereby granted, free of charge, to any person obtaining a copy
-        of this software and associated documentation files (the "Software"), to deal
-        in the Software without restriction, including without limitation the rights
-        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-        copies of the Software, and to permit persons to whom the Software is
-        furnished to do so, subject to the following conditions:
-        
-        The above copyright notice and this permission notice shall be included in all
-        copies or substantial portions of the Software.
-        
-        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-        SOFTWARE.
-        
 Project-URL: Repository, https://github.com/tanganke/fusion_bench
 Project-URL: Homepage, https://github.com/tanganke/fusion_bench
 Project-URL: Issues, https://github.com/tanganke/fusion_bench/issues