fusion-bench 0.2.9__py3-none-any.whl

fusion_bench/utils/cache_utils.py

@@ -0,0 +1,58 @@
+import logging
+import os
+import pickle
+from functools import wraps
+from pathlib import Path
+from typing import Any, Callable, Union
+
+__all__ = ["cache_to_disk"]
+
+
+log = logging.getLogger(__name__)
+
+
+def cache_to_disk(file_path: Union[str, Path]) -> Callable:
+    """
+    A decorator to cache the result of a function to a file. If the file exists,
+    the result is loaded from the file. Otherwise, the function is executed and
+    the result is saved to the file.
+
+    ## Example usage
+
+    ```python
+    @cache_to_disk("path_to_file.pkl")
+    def some_function(*args: Any, **kwargs: Any) -> Any:
+        # Function implementation
+        return "some result"
+    ```
+
+    Args:
+        file_path (str): The path to the file where the result should be cached.
+
+    Returns:
+        Callable: The decorated function.
+    """
+    if isinstance(file_path, str):
+        file_path = Path(file_path)
+    assert isinstance(file_path, Path)
+
+    def decorator(func: Callable) -> Callable:
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if os.path.exists(file_path):
+                log.info(
+                    f"Loading cached result of {func.__name__} from {file_path}",
+                    stacklevel=2,
+                )
+                with open(file_path, "rb") as f:
+                    return pickle.load(f)
+            else:
+                result = func(*args, **kwargs)
+                file_path.parent.mkdir(parents=True, exist_ok=True)
+                with open(file_path, "wb") as f:
+                    pickle.dump(result, f)
+                return result
+
+        return wrapper
+
+    return decorator

fusion_bench/utils/data.py

@@ -0,0 +1,165 @@
+import pickle
+from pathlib import Path
+from typing import Literal, Optional, Union
+
+import numpy as np
+import torch
+import torch.utils.data
+from torch.utils.data import DataLoader, Dataset
+
+
+class InfiniteDataLoader:
+    def __init__(self, data_loader: DataLoader):
+        self.data_loader = data_loader
+        self.data_iter = iter(data_loader)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            data = next(self.data_iter)
+        except StopIteration:
+            self.data_iter = iter(self.data_loader)  # Reset the data loader
+            data = next(self.data_iter)
+        return data
+
+
+def load_tensor_from_file(file_path: Union[str, Path], 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.
+
+    Args:
+        file_path (str): The path to the file to load.
+        device: The device to move the tensor to. By default the tensor is loaded on the CPU.
+
+    Returns:
+        torch.Tensor: The tensor loaded from the file.
+    """
+    if file_path.endswith(".np"):
+        tensor = torch.from_numpy(np.load(file_path)).detach_()
+    if file_path.endswith((".pt", ".pth")):
+        tensor = torch.load(file_path, map_location="cpu").detach_()
+    else:
+        try:
+            tensor = pickle.load(open(file_path, "rb"))
+        except Exception:
+            raise ValueError(f"Unsupported file format: {file_path}")
+
+    # Move tensor to device
+    assert isinstance(tensor, torch.Tensor), f"Expected tensor, got {type(tensor)}"
+    if device is not None:
+        tensor = tensor.to(device=device)
+    return tensor
+
+
+def train_validation_split(
+    dataset: Dataset,
+    validation_fraction: Optional[float] = 0.1,
+    validation_size: Optional[int] = None,
+    random_seed: Optional[int] = None,
+    return_split: Literal["all", "train", "val"] = "both",
+):
+    """
+    Split a dataset into a training and validation set.
+
+    Args:
+        dataset (Dataset): The dataset to split.
+        validation_fraction (Optional[float]): The fraction of the dataset to use for validation.
+        validation_size (Optional[int]): The number of samples to use for validation. `validation_fraction` must be set to `None` if this is provided.
+        random_seed (Optional[int]): The random seed to use for reproducibility.
+        return_split (Literal["all", "train", "val"]): The split to return.
+
+    Returns:
+        Tuple[Dataset, Dataset]: The training and validation datasets.
+    """
+    # Check the input arguments
+    assert (
+        validation_fraction is None or validation_size is None
+    ), "Only one of validation_fraction and validation_size can be provided"
+    assert (
+        validation_fraction is not None or validation_size is not None
+    ), "Either validation_fraction or validation_size must be provided"
+
+    # Compute the number of samples for training and validation
+    num_samples = len(dataset)
+    if validation_size is not None:
+        assert (
+            0 < validation_fraction < 1
+        ), "Validation fraction must be between 0 and 1"
+        num_validation_samples = int(num_samples * validation_fraction)
+        num_training_samples = num_samples - num_validation_samples
+    else:
+        assert (
+            validation_size < num_samples
+        ), "Validation size must be less than num_samples"
+        num_validation_samples = validation_size
+        num_training_samples = num_samples - num_validation_samples
+
+    # Split the dataset
+    generator = (
+        torch.Generator().manual_seed(random_seed) if random_seed is not None else None
+    )
+    training_dataset, validation_dataset = torch.utils.data.random_split(
+        dataset, [num_training_samples, num_validation_samples], generator=generator
+    )
+
+    # return the split as requested
+    if return_split == "all":
+        return training_dataset, validation_dataset
+    elif return_split == "train":
+        return training_dataset
+    elif return_split == "val":
+        return validation_dataset
+    else:
+        raise ValueError(f"Invalid return_split: {return_split}")
+
+
+def train_validation_test_split(
+    dataset: Dataset,
+    validation_fraction: float,
+    test_fraction: float,
+    random_seed: Optional[int] = None,
+    return_spilt: Literal["all", "train", "val", "test"] = "all",
+):
+    """
+    Split a dataset into a training, validation and test set.
+
+    Args:
+        dataset (Dataset): The dataset to split.
+        validation_fraction (float): The fraction of the dataset to use for validation.
+        test_fraction (float): The fraction of the dataset to use for test.
+        random_seed (Optional[int]): The random seed to use for reproducibility.
+        return_spilt (Literal["all", "train", "val", "test"]): The split to return.
+
+    Returns:
+        Tuple[Dataset, Dataset, Dataset]: The training, validation and test datasets.
+    """
+    num_samples = len(dataset)
+    assert 0 < validation_fraction < 1, "Validation fraction must be between 0 and 1"
+    assert 0 < test_fraction < 1, "Test fraction must be between 0 and 1"
+    generaotr = (
+        torch.Generator().manual_seed(random_seed) if random_seed is not None else None
+    )
+
+    num_validation_samples = int(num_samples * validation_fraction)
+    num_test_samples = int(num_samples * test_fraction)
+    num_training_samples = num_samples - num_validation_samples - num_test_samples
+    training_dataset, validation_dataset, test_dataset = torch.utils.data.random_split(
+        dataset,
+        [num_training_samples, num_validation_samples, num_test_samples],
+        generator=generaotr,
+    )
+
+    # return the split as requested
+    if return_spilt == "all":
+        return training_dataset, validation_dataset, test_dataset
+    elif return_spilt == "train":
+        return training_dataset
+    elif return_spilt == "val":
+        return validation_dataset
+    elif return_spilt == "test":
+        return test_dataset
+    else:
+        raise ValueError(f"Invalid return_split: {return_spilt}")

fusion_bench/utils/devices.py

@@ -0,0 +1,231 @@
+import gc
+import os
+from typing import List, Optional, Union
+
+import torch
+from transformers.utils import (
+    is_torch_bf16_gpu_available,
+    is_torch_cuda_available,
+    is_torch_mps_available,
+    is_torch_npu_available,
+    is_torch_xpu_available,
+)
+
+__all__ = [
+    "cuda_empty_cache",
+    "to_device",
+    "num_devices",
+    "get_device",
+    "get_current_device",
+    "get_device_memory_info",
+    "get_device_capabilities",
+]
+
+
+def cuda_empty_cache():
+    gc.collect()
+    torch.cuda.empty_cache()
+
+
+def to_device(obj, device: Optional[torch.device], **kwargs):
+    """
+    Move a given object to the specified device.
+
+    This function recursively moves tensors, modules, lists, tuples, and dictionaries to the specified device.
+    For unsupported types, the object is returned as is.
+
+    Args:
+        obj: The object to be moved to the device. This can be a torch.Tensor, torch.nn.Module, list, tuple, or dict.
+        device (torch.device): The target device to move the object to. This can be `None`.
+        **kwargs: Additional keyword arguments to be passed to the `to` method of torch.Tensor or torch.nn.Module. For example, `non_blocking=True`, `dtype=torch.float16`.
+
+    Returns:
+        The object moved to the specified device. The type of the returned object matches the type of the input object.
+
+    Examples:
+        >>> tensor = torch.tensor([1, 2, 3])
+        >>> to_device(tensor, torch.device('cuda'))
+        tensor([1, 2, 3], device='cuda:0')
+
+        >>> model = torch.nn.Linear(2, 2)
+        >>> to_device(model, torch.device('cuda'))
+        Linear(..., device='cuda:0')
+
+        >>> data = [torch.tensor([1, 2]), torch.tensor([3, 4])]
+        >>> to_device(data, torch.device('cuda'))
+        [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')]
+    """
+    if isinstance(obj, (torch.Tensor, torch.nn.Module)):
+        return obj.to(device, **kwargs)
+    elif isinstance(obj, list):
+        return [to_device(o, device) for o in obj]
+    elif isinstance(obj, tuple):
+        return tuple(to_device(o, device) for o in obj)
+    elif isinstance(obj, dict):
+        for key in obj:
+            obj[key] = to_device(obj[key], device)
+        return obj
+    else:
+        # the default behavior is to return the object as is
+        return obj
+
+
+def num_devices(devices: Union[int, List[int], str]) -> int:
+    """
+    Return the number of devices.
+
+    Args:
+        devices: `devices` can be a single int to specify the number of devices, or a list of device ids, e.g. [0, 1, 2, 3]， or a str of device ids, e.g. "0,1,2,3" and "[0, 1, 2]".
+
+    Returns:
+        The number of devices.
+    """
+    if isinstance(devices, int):
+        return devices
+    elif isinstance(devices, str):
+        return len(devices.split(","))
+    elif isinstance(devices, list):
+        return len(devices)
+    else:
+        raise TypeError(
+            f"devices must be a single int or a list of ints, but got {type(devices)}"
+        )
+
+
+def get_device(obj) -> torch.device:
+    """
+    Get the device of a given object.
+
+    Args:
+        obj: The object whose device is to be determined.
+
+    Returns:
+        torch.device: The device of the given object.
+
+    Raises:
+        ValueError: If the object type is not supported.
+    """
+    if isinstance(obj, torch.Tensor):
+        return obj.device
+    elif isinstance(obj, torch.nn.Module):
+        if hasattr(obj, "device"):
+            return obj.device
+        else:
+            return next(iter(obj.parameters())).device
+    elif isinstance(obj, torch.device):
+        return obj
+    else:
+        raise ValueError(f"Unsupported object type: {type(obj)}")
+
+
+def get_current_device() -> torch.device:
+    R"""
+    Gets the current available device for PyTorch operations.
+    This is used for distributed training.
+
+    This function checks the availability of various types of devices in the following order:
+    1. XPU (Intel's AI accelerator)
+    2. NPU (Neural Processing Unit)
+    3. MPS (Metal Performance Shaders, for Apple devices)
+    4. CUDA (NVIDIA's GPU)
+    5. CPU (Central Processing Unit, used as a fallback)
+
+    The function returns the first available device found in the above order. If none of the specialized devices
+    are available, it defaults to the CPU.
+
+    Returns:
+        torch.device: The current available device for PyTorch operations.
+
+    Environment Variables:
+        LOCAL_RANK: This environment variable is used to specify the device index for multi-device setups.
+                    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
+    """
+
+    if is_torch_xpu_available():
+        device = "xpu:{}".format(os.environ.get("LOCAL_RANK", "0"))
+    elif is_torch_npu_available():
+        device = "npu:{}".format(os.environ.get("LOCAL_RANK", "0"))
+    elif is_torch_mps_available():
+        device = "mps:{}".format(os.environ.get("LOCAL_RANK", "0"))
+    elif is_torch_cuda_available():
+        device = "cuda:{}".format(os.environ.get("LOCAL_RANK", "0"))
+    else:
+        device = "cpu"
+
+    return torch.device(device)
+
+
+def get_device_memory_info(device: torch.device, reset_stats: bool = True) -> dict:
+    """
+    Get memory information for a given device.
+
+    Args:
+        device (torch.device): The device for which to get memory information.
+
+    Returns:
+        dict: A dictionary containing memory information for the given device.
+    """
+    if device.type == "cuda":
+        total_memory = torch.cuda.get_device_properties(device).total_memory
+        reserved_memory = torch.cuda.memory_reserved(device)
+        allocated_memory = torch.cuda.memory_allocated(device)
+        peak_memory_active = torch.cuda.memory_stats(device).get(
+            "active_bytes.all.peak", 0
+        )
+        peak_mem_alloc = torch.cuda.max_memory_allocated(device)
+        peak_mem_reserved = torch.cuda.max_memory_reserved(device)
+
+        if reset_stats:
+            torch.cuda.reset_peak_memory_stats(device)
+
+        return {
+            "total_memory": total_memory,
+            "reserved_memory": reserved_memory,
+            "allocated_memory": allocated_memory,
+            "peak_memory_active": peak_memory_active,
+            "peak_memory_allocated": peak_mem_alloc,
+            "peak_memory_reserved": peak_mem_reserved,
+        }
+    else:
+        raise ValueError(
+            f"Memory information not available for device type: {device.type}"
+        )
+
+
+def get_device_capabilities(device: torch.device) -> dict:
+    """
+    Get capabilities information for a given device.
+
+    Args:
+        device (torch.device): The device for which to get capabilities information.
+
+    Returns:
+        dict: A dictionary containing capabilities information for the given device.
+    """
+    if device.type == "cuda":
+        return {
+            "name": torch.cuda.get_device_name(device),
+            "capability": torch.cuda.get_device_capability(device),
+            "total_memory": torch.cuda.get_device_properties(device).total_memory,
+            "multi_processor_count": torch.cuda.get_device_properties(
+                device
+            ).multi_processor_count,
+        }
+    else:
+        raise ValueError(
+            f"Capabilities information not available for device type: {device.type}"
+        )
+
+
+def cleanup_cuda():
+    """
+    Call gc collect, empty CUDA cache, and reset peak memory stats.
+    """
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()

fusion_bench/utils/dict.py

@@ -0,0 +1,43 @@
+from copy import deepcopy
+from typing import Iterable, List, Tuple, Union
+
+
+def dict_get(d: dict, keys: Iterable[str], default=None):
+    return [d.get(k, default) for k in keys]
+
+
+def dict_map(f, d: dict, *, max_level: int = -1, skip_levels=0, inplace=False):
+    """Apply function f to each element in dictionary d and return a new dictionary.
+
+    Args:
+        f (callable): function to apply
+        d (dict): input dictionary
+        max_level (int, optional): maximum depth to apply function, -1 means unlimited. Defaults to -1.
+        skip_levels (int, optional): number of levels to skip. Defaults to 0.
+        inplace (bool, optional): whether to modify input dictionary in place. Defaults to False.
+
+    Returns:
+        dict: transformed dictionary
+    """
+    if not isinstance(d, dict):
+        raise TypeError("dict_map: d must be a dict")
+
+    if inplace:
+        ans = d
+    else:
+        ans = deepcopy(d)
+
+    def dict_map_impl(from_dict, to_dict, level):
+        if level == max_level:
+            return
+        for k in from_dict.keys():
+            if isinstance(from_dict[k], dict):
+                dict_map_impl(from_dict[k], to_dict[k], level + 1)
+            else:
+                if level < skip_levels:
+                    continue
+                else:
+                    to_dict[k] = f(from_dict[k])
+
+    dict_map_impl(d, ans, 0)
+    return ans

fusion_bench/utils/dtype.py

@@ -0,0 +1,146 @@
+import contextlib
+from typing import Dict, Generator, Iterable, Optional, Tuple
+
+import torch
+from transformers.utils import (
+    is_torch_bf16_gpu_available,
+    is_torch_cuda_available,
+    is_torch_mps_available,
+    is_torch_npu_available,
+    is_torch_xpu_available,
+)
+
+PRECISION_STR_TO_DTYPE: Dict[str, torch.dtype] = {
+    "fp16": torch.float16,
+    "float16": torch.float16,
+    "bf16": torch.bfloat16,
+    "bfloat16": torch.bfloat16,
+    "float": torch.float32,
+    "fp32": torch.float32,
+    "float32": torch.float32,
+    "double": torch.float64,
+    "fp64": torch.float64,
+    "float64": torch.float64,
+}
+
+
+def parse_dtype(dtype: Optional[str]):
+    """
+    Parses a string representation of a data type and returns the corresponding torch.dtype.
+
+    Args:
+        dtype (Optional[str]): The string representation of the data type.
+                               Can be one of "float32", "float", "float64", "double",
+                               "float16", "half", "bfloat16", or "bf16".
+                               If None, returns None.
+
+    Returns:
+        torch.dtype: The corresponding torch.dtype if the input is a valid string representation.
+                     If the input is already a torch.dtype, it is returned as is.
+                     If the input is None, returns None.
+
+    Raises:
+        ValueError: If the input string does not correspond to a supported data type.
+    """
+    if isinstance(dtype, torch.dtype):
+        return dtype
+
+    if dtype is None:
+        return None
+
+    dtype = dtype.strip('"')
+    if dtype not in PRECISION_STR_TO_DTYPE:
+        raise ValueError(f"Unsupported dtype: {type(dtype)}")
+
+    dtype = PRECISION_STR_TO_DTYPE[dtype]
+    return dtype
+
+
+def get_dtype(obj) -> torch.dtype:
+    """
+    Get the data type (dtype) of a given object.
+
+    Returns:
+        torch.dtype: The data type of the given object.
+
+    Raises:
+        ValueError: If the object type is not supported.
+    """
+    if isinstance(obj, torch.Tensor):
+        return obj.dtype
+    elif isinstance(obj, torch.nn.Module):
+        if hasattr(obj, "dtype"):
+            return obj.dtype
+        else:
+            return next(iter(obj.parameters())).dtype
+    elif isinstance(obj, (torch.device, str)):
+        return parse_dtype(obj)
+    else:
+        raise ValueError(f"Unsupported object type: {type(obj)}")
+
+
+@contextlib.contextmanager
+def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
+    """
+    Context manager to set torch's default dtype.
+
+    Args:
+        dtype (torch.dtype): The desired default dtype inside the context manager.
+
+    Returns:
+        ContextManager: context manager for setting default dtype.
+
+    Example:
+        >>> with set_default_dtype(torch.bfloat16):
+        >>>     x = torch.tensor([1, 2, 3])
+        >>>     x.dtype
+        torch.bfloat16
+
+
+    """
+    old_dtype = torch.get_default_dtype()
+    torch.set_default_dtype(dtype)
+    try:
+        yield
+    finally:
+        torch.set_default_dtype(old_dtype)
+
+
+def infer_optim_dtype(model_dtype: "torch.dtype") -> "torch.dtype":
+    r"""
+    Infers the optimal dtype according to the model_dtype and device compatibility.
+    """
+    _is_fp16_available = is_torch_npu_available() or is_torch_cuda_available()
+    try:
+        _is_bf16_available = is_torch_bf16_gpu_available() or (
+            is_torch_npu_available() and torch.npu.is_bf16_supported()
+        )
+    except Exception:
+        _is_bf16_available = False
+
+    if _is_bf16_available and model_dtype == torch.bfloat16:
+        return torch.bfloat16
+    elif _is_fp16_available:
+        return torch.float16
+    else:
+        return torch.float32
+
+
+def validate_expected_param_dtype(
+    named_params: Iterable[Tuple[str, torch.nn.Parameter]], dtype: torch.dtype
+) -> None:
+    """
+    Validates that all input parameters have the expected dtype.
+
+    Args:
+        named_params (Iterable[Tuple[str, torch.nn.Parameter]]): Iterable of named parameters.
+        dtype (torch.dtype): Expected dtype.
+
+    Raises:
+        ValueError: If any parameter has a different dtype than `dtype`.
+    """
+    for name, param in named_params:
+        if param.dtype != dtype:
+            raise ValueError(
+                f"Parameter {name} has dtype {param.dtype}, but expected {dtype}"
+            )