Model Utilities

Type Definitions

fusion_bench.utils.type

StateDictType = Dict[str, Tensor] module-attribute

BoolStateDictType = Dict[str, torch.BoolTensor] module-attribute

TorchModelType = TypeVar('TorchModelType', bound=(nn.Module)) module-attribute

Parameter Count and Manipulation

fusion_bench.utils.parameters

check_parameters_all_equal(list_of_param_names)

Checks if all models have the same parameters.

This function takes a list of parameter names or state dictionaries from different models. It checks if all models have the same parameters by comparing the parameter names. If any model has different parameters, it raises a ValueError with the differing parameters.

Parameters:

• list_of_param_names (List[Union[StateDict, List[str]]]) –

  A list of parameter names or state dictionaries.

Raises:

• ValueError –

  If any model has different parameters.

Returns:

• None –

  None

Source code in fusion_bench/utils/parameters.py

def check_parameters_all_equal(
    list_of_param_names: List[Union[StateDictType, nn.Module, List[str]]],
) -> None:
    """
    Checks if all models have the same parameters.

    This function takes a list of parameter names or state dictionaries from different models.
    It checks if all models have the same parameters by comparing the parameter names.
    If any model has different parameters, it raises a ValueError with the differing parameters.

    Args:
        list_of_param_names (List[Union[StateDict, List[str]]]): A list of parameter names or state dictionaries.

    Raises:
        ValueError: If any model has different parameters.

    Returns:
        None
    """
    if isinstance(list_of_param_names[0], Mapping):
        list_of_param_names = [list(i.keys()) for i in list_of_param_names]
    elif isinstance(list_of_param_names[0], nn.Module):
        list_of_param_names = [list(i.state_dict().keys()) for i in list_of_param_names]
    else:
        parameter_names = set(list_of_param_names[0])

        if len(list_of_param_names) >= 2:
            # raise ValueError("Number of models is less than 2.")
            for names in list_of_param_names[1:]:
                current_parameterNames = set(names)
                if current_parameterNames != parameter_names:
                    raise ValueError(
                        "Differing parameter names in models. "
                        f"The different parameters are {parameter_names.symmetric_difference(current_parameterNames)}"
                    )

count_parameters(module, non_zero_only=False)

Counts the number of trainable and total parameters in a PyTorch model.

Parameters:

• model (Module) –

  The PyTorch model for which to count parameters.

• non_zero_only (bool, default: False ) –

  If True, only non-zero parameters are counted. If False, all parameters are counted. Defaults to False.

Returns:

• tuple ( tuple[int, int] ) –

  A tuple containing the number of trainable parameters and the total number of parameters.

Examples:

# Count the parameters
trainable_params, all_params = count_parameters(model)

Source code in fusion_bench/utils/parameters.py

@torch.no_grad()
def count_parameters(module: nn.Module, non_zero_only: bool = False) -> tuple[int, int]:
    """
    Counts the number of trainable and total parameters in a PyTorch model.

    Args:
        model (nn.Module): The PyTorch model for which to count parameters.
        non_zero_only (bool, optional): If True, only non-zero parameters are counted. If False, all parameters are counted. Defaults to False.

    Returns:
        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)
        ```
    """
    trainable_params = 0
    all_param = 0

    for name, param in module.named_parameters():
        # count the number of parameters
        num_params = _numel(param, non_zero_only)

        # accumulate the number of trainable and total parameters
        all_param += num_params
        if param.requires_grad:
            trainable_params += num_params

    return trainable_params, all_param

get_parameter_statistics(module_or_state_dict, model_wise=False)

Get statistics of the parameters in a PyTorch model or state dictionary.

Parameters:

• module_or_state_dict (Union[Module, StateDictType]) –

  The PyTorch model for which to get parameter statistics.

Returns:

• dict ( dict ) –

  A dictionary containing the mean, standard deviation, min, and max of the parameters.

Source code in fusion_bench/utils/parameters.py

@torch.no_grad()
def get_parameter_statistics(
    module_or_state_dict: Union[nn.Module, StateDictType],
    model_wise: bool = False,
) -> dict:
    """
    Get statistics of the parameters in a PyTorch model or state dictionary.

    Args:
        module_or_state_dict (Union[nn.Module, StateDictType]): The PyTorch model for which to get parameter statistics.

    Returns:
        dict: A dictionary containing the mean, standard deviation, min, and max of the parameters.
    """
    stats = {}
    if isinstance(module_or_state_dict, nn.Module):
        state_dict = module_or_state_dict.state_dict()
    else:
        state_dict = module_or_state_dict

    if model_wise:
        # if model-wise, return the statistics for the entire model
        state_dict = {"model": state_dict_to_vector(state_dict)}

    for name, param in state_dict.items():
        stats[name] = {
            "mean": param.data.mean().item(),
            "std": param.data.std().item(),
            "min": param.data.min().item(),
            "max": param.data.max().item(),
        }

    return stats

get_parameter_summary(module_or_state_dict, non_zero_only=False)

Get a summary of the parameters in a PyTorch model.

Source code in fusion_bench/utils/parameters.py

@torch.no_grad()
def get_parameter_summary(
    module_or_state_dict: Union[nn.Module, StateDictType], non_zero_only: bool = False
) -> dict:
    """
    Get a summary of the parameters in a PyTorch model.
    """
    if isinstance(module_or_state_dict, nn.Module):
        state_dict = module_or_state_dict.state_dict(keep_vars=True)
    else:
        state_dict = module_or_state_dict

    trainable_params = 0
    all_param = 0
    bytes = 0

    for name, param in state_dict.items():
        # count the number of parameters
        num_params = _numel(param, non_zero_only)
        bytes += _numel(param, non_zero_only=False) * param.element_size()

        # accumulate the number of trainable and total parameters
        all_param += num_params
        if param.requires_grad:
            trainable_params += num_params

    return {
        "trainable_params": trainable_params,
        "all_param": all_param,
        "bytes": bytes,
    }

human_readable(num)

Converts a number into a human-readable string with appropriate magnitude suffix.

Examples:

print(human_readable(1500))
# Output: '1.50K'
print(human_readable(1500000))
# Output: '1.50M'

Parameters:

• num (int) –

  The number to convert.

Returns:

• str ( str ) –

  The human-readable string representation of the number.

Source code in fusion_bench/utils/parameters.py

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'
        print(human_readable(1500000))
        # Output: '1.50M'
        ```

    Args:
        num (int): The number to convert.

    Returns:
        str: The human-readable string representation of the number.
    """
    if num < 1000 and isinstance(num, int):
        return str(num)
    magnitude = 0
    while abs(num) >= 1000:
        magnitude += 1
        num /= 1000.0
    return "%.2f%s" % (num, ["", "K", "M", "B", "T", "P"][magnitude])

print_parameters(module, is_human_readable=True, print_fn=print, non_zero_only=False)

Prints the number of trainable and total parameters in a PyTorch model.

Parameters:

• module (Module) –

  The PyTorch model for which to print parameters.

• human_readable (bool) –

  If True, the parameter counts are converted to a human-readable format (e.g., '1.5M' instead of '1500000'). Defaults to True.

• print_fn (Callable, default: print ) –

  Function used to print the message.

• non_zero_only (bool, default: False ) –

  If True, only non-zero elements are counted. If False, all elements are counted. Defaults to False.

Prints

The number of trainable parameters, the total number of parameters, and the percentage of trainable parameters in the model.

Source code in fusion_bench/utils/parameters.py

def print_parameters(
    module: nn.Module,
    is_human_readable: bool = True,
    print_fn=print,
    non_zero_only: bool = False,
):
    """
    Prints the number of trainable and total parameters in a PyTorch model.

    Args:
        module (nn.Module): The PyTorch model for which to print parameters.
        human_readable (bool, optional): If True, the parameter counts are converted to a human-readable format (e.g., '1.5M' instead of '1500000'). Defaults to True.
        print_fn (Callable): Function used to print the message.
        non_zero_only (bool, optional): If True, only non-zero elements are counted. If False, all elements are counted. Defaults to False.

    Prints:
        The number of trainable parameters, the total number of parameters, and the percentage of trainable parameters in the model.
    """
    trainable_params, all_param = count_parameters(module, non_zero_only=non_zero_only)
    trainable_ratio = 100 * trainable_params / all_param
    if is_human_readable:
        trainable_params = human_readable(trainable_params)
        all_param = human_readable(all_param)

    print_fn(
        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {trainable_ratio:.4f}"
    )

state_dict_to_vector(state_dict, remove_keys=None)

Convert a state dictionary to a vector.

Parameters:

• state_dict (Union[dict[str, Tensor], Module]) –

  The state dictionary to convert.

• remove_keys (list, default: None ) –

  List of keys to remove from the state dictionary. Defaults to [].

Returns:

• –

  torch.Tensor: The converted vector.

Source code in fusion_bench/utils/parameters.py

def state_dict_to_vector(
    state_dict: Union[StateDictType, nn.Module],
    remove_keys: Optional[List[str]] = None,
):
    """
    Convert a state dictionary to a vector.

    Args:
        state_dict (Union[dict[str, torch.Tensor], nn.Module]): The state dictionary to convert.
        remove_keys (list, optional): List of keys to remove from the state dictionary. Defaults to [].

    Returns:
        torch.Tensor: The converted vector.
    """
    remove_keys = remove_keys if remove_keys is not None else []

    if isinstance(state_dict, nn.Module):
        shared_state_dict = state_dict.state_dict()
    else:
        shared_state_dict = copy.copy(state_dict)

    # remove the keys to be removed
    for key in remove_keys:
        if key in shared_state_dict:
            del shared_state_dict[key]

    # sort the reference dict
    sorted_shared_state_dict = OrderedDict(sorted(shared_state_dict.items()))

    vector = nn.utils.parameters_to_vector(
        [value.reshape(-1) for key, value in sorted_shared_state_dict.items()]
    )
    return vector

trainable_state_dict(module, prefix='', keep_vars=False)

Returns the state dictionary of the module containing only the trainable parameters.

Parameters:

• module (Module) –

  The neural network module.

• prefix (str, default: '' ) –

  The prefix to add to the parameter names. Defaults to "".

• keep_vars (bool, default: False ) –

  If True, the parameters are not detached. Defaults to False.

Returns:

• StateDictType –

  Dict[str, Tensor]: A dictionary containing the names and values of the trainable parameters.

Source code in fusion_bench/utils/parameters.py

def trainable_state_dict(
    module: nn.Module,
    prefix: str = "",
    keep_vars: bool = False,
) -> StateDictType:
    """
    Returns the state dictionary of the module containing only the trainable parameters.

    Args:
        module (nn.Module): The neural network module.
        prefix (str, optional): The prefix to add to the parameter names. Defaults to "".
        keep_vars (bool, optional): If True, the parameters are not detached. Defaults to False.

    Returns:
        Dict[str, Tensor]: A dictionary containing the names and values of the trainable parameters.
    """
    state_dict = {
        prefix + name: param if keep_vars else param.detach()
        for name, param in module.named_parameters()
        if param.requires_grad
    }
    return state_dict

vector_to_state_dict(vector, state_dict, remove_keys=None)

Convert a vector to a state dictionary.

Parameters:

• vector (Tensor) –

  The vector to convert.

• state_dict (Union[dict[str, Tensor], Module]) –

  The reference state dictionary to define the order of the vector.

• remove_keys (list, default: None ) –

  List of keys to remove from the reference state dictionary. Defaults to [].

Returns:

• dict ( Dict[str, Tensor] ) –

  The converted state dictionary.

Source code in fusion_bench/utils/parameters.py

def vector_to_state_dict(
    vector: torch.Tensor,
    state_dict: Union[StateDictType, nn.Module],
    remove_keys: Optional[List[str]] = None,
) -> Dict[str, torch.Tensor]:
    """
    Convert a vector to a state dictionary.

    Args:
        vector (torch.Tensor): The vector to convert.
        state_dict (Union[dict[str, torch.Tensor], nn.Module]): The reference state dictionary to define the order of the vector.
        remove_keys (list, optional): List of keys to remove from the reference state dictionary. Defaults to [].

    Returns:
        dict: The converted state dictionary.
    """
    remove_keys = remove_keys if remove_keys is not None else []

    # create a reference dict to define the order of the vector
    if isinstance(state_dict, nn.Module):
        reference_dict = state_dict.state_dict()
    else:
        # shallow copy the state_dict
        reference_dict = copy.copy(state_dict)

    # remove the keys to be removed
    for key in remove_keys:
        if key in reference_dict:
            del reference_dict[key]

    # sort the reference dict
    sorted_reference_dict = OrderedDict(sorted(reference_dict.items()))

    # create a shared state dict using the reference dict
    nn.utils.vector_to_parameters(vector, sorted_reference_dict.values())

    # add back the encoder and decoder embedding weights.
    if "transformer.shared.weight" in sorted_reference_dict:
        for key in remove_keys:
            sorted_reference_dict[key] = sorted_reference_dict[
                "transformer.shared.weight"
            ]
    return sorted_reference_dict

State Dict Arithmetic

fusion_bench.utils.state_dict_arithmetic

ArithmeticStateDict

Bases: OrderedDict

An OrderedDict subclass that supports arithmetic operations on state dictionaries.

This class provides convenient operator overloading for common state dict operations like addition, subtraction, multiplication, and division, while maintaining all the functionality of OrderedDict.

Examples:

>>> sd1 = ArithmeticStateDict({'weight': torch.tensor([1.0, 2.0]), 'bias': torch.tensor([0.5])})
>>> sd2 = ArithmeticStateDict({'weight': torch.tensor([2.0, 3.0]), 'bias': torch.tensor([1.0])})
>>> result = sd1 + sd2  # Element-wise addition
>>> result = sd1 - sd2  # Element-wise subtraction
>>> result = sd1 * 2.0  # Scalar multiplication
>>> result = sd1 / 2.0  # Scalar division
>>> result = sd1 @ sd2  # Hadamard product

Source code in fusion_bench/utils/state_dict_arithmetic.py

class ArithmeticStateDict(OrderedDict):
    """
    An OrderedDict subclass that supports arithmetic operations on state dictionaries.

    This class provides convenient operator overloading for common state dict operations
    like addition, subtraction, multiplication, and division, while maintaining all
    the functionality of OrderedDict.

    Examples:
        >>> sd1 = ArithmeticStateDict({'weight': torch.tensor([1.0, 2.0]), 'bias': torch.tensor([0.5])})
        >>> sd2 = ArithmeticStateDict({'weight': torch.tensor([2.0, 3.0]), 'bias': torch.tensor([1.0])})
        >>> result = sd1 + sd2  # Element-wise addition
        >>> result = sd1 - sd2  # Element-wise subtraction
        >>> result = sd1 * 2.0  # Scalar multiplication
        >>> result = sd1 / 2.0  # Scalar division
        >>> result = sd1 @ sd2  # Hadamard product
    """

    def __init__(self, *args, **kwargs):
        """Initialize ArithmeticStateDict with the same interface as OrderedDict."""
        super().__init__(*args, **kwargs)

    def __add__(
        self, other: Union["ArithmeticStateDict", StateDictType, Number]
    ) -> "ArithmeticStateDict":
        """
        Element-wise addition with another state dict or scalar.

        Args:
            other: Another state dict to add or a scalar to add to all elements.

        Returns:
            A new ArithmeticStateDict with the element-wise sum.
        """
        if isinstance(other, (int, float, Number)):
            # Scalar addition
            result_dict = state_dict_add_scalar(self, other)
            return ArithmeticStateDict(result_dict)
        elif isinstance(other, (dict, OrderedDict)):
            # State dict addition
            result_dict = state_dict_add(self, other, strict=True)
            return ArithmeticStateDict(result_dict)
        else:
            raise TypeError(
                f"Cannot add ArithmeticStateDict with {type(other).__name__}"
            )

    def __radd__(
        self, other: Union["ArithmeticStateDict", StateDictType, Number]
    ) -> "ArithmeticStateDict":
        """
        Right addition (other + self).
        Handles the case where sum() starts with 0 and scalar addition.
        """
        if other == 0:  # sum() starts with 0 by default
            return self
        elif isinstance(other, (int, float, Number)):
            # Scalar addition is commutative
            return self.__add__(other)
        elif isinstance(other, (dict, OrderedDict)):
            return self.__add__(other)
        else:
            raise TypeError(
                f"Cannot add {type(other).__name__} with ArithmeticStateDict"
            )

    def __sub__(
        self, other: Union["ArithmeticStateDict", StateDictType, Number]
    ) -> "ArithmeticStateDict":
        """
        Element-wise subtraction with another state dict or scalar.

        Args:
            other: Another state dict to subtract or a scalar to subtract from all elements.

        Returns:
            A new ArithmeticStateDict with the element-wise difference.
        """
        if isinstance(other, (int, float, Number)):
            # Scalar subtraction: subtract scalar from all elements
            result_dict = state_dict_add_scalar(self, -other)
            return ArithmeticStateDict(result_dict)
        elif isinstance(other, (dict, OrderedDict)):
            # State dict subtraction
            result_dict = state_dict_sub(self, other, strict=True)
            return ArithmeticStateDict(result_dict)
        else:
            raise TypeError(
                f"Cannot subtract {type(other).__name__} from ArithmeticStateDict"
            )

    def __rsub__(
        self, other: Union["ArithmeticStateDict", StateDictType, Number]
    ) -> "ArithmeticStateDict":
        """Right subtraction (other - self)."""
        if isinstance(other, (int, float, Number)):
            # Scalar - ArithmeticStateDict: subtract each element from scalar
            result = ArithmeticStateDict()
            for key, tensor in self.items():
                result[key] = other - tensor
            return result
        elif isinstance(other, (dict, OrderedDict)):
            result_dict = state_dict_sub(other, self, strict=True)
            return ArithmeticStateDict(result_dict)
        else:
            raise TypeError(
                f"Cannot subtract ArithmeticStateDict from {type(other).__name__}"
            )

    def __mul__(
        self, scalar: Union[Number, "ArithmeticStateDict", StateDictType]
    ) -> "ArithmeticStateDict":
        """
        Scalar multiplication or Hadamard product.

        Args:
            scalar: A scalar value for element-wise multiplication, or another state dict
                   for Hadamard product.

        Returns:
            A new ArithmeticStateDict with the result.
        """
        if isinstance(scalar, (int, float, Number)):
            result_dict = state_dict_mul(self, scalar)
            return ArithmeticStateDict(result_dict)
        elif isinstance(scalar, (dict, OrderedDict)):
            # Hadamard product for dict-like objects
            result_dict = state_dict_hadamard_product(self, scalar)
            return ArithmeticStateDict(result_dict)
        else:
            raise TypeError(
                f"Cannot multiply ArithmeticStateDict with {type(scalar).__name__}"
            )

    def __rmul__(
        self, scalar: Union[Number, "ArithmeticStateDict", StateDictType]
    ) -> "ArithmeticStateDict":
        """Right multiplication (scalar * self)."""
        return self.__mul__(scalar)

    def __truediv__(self, scalar: Number) -> "ArithmeticStateDict":
        """
        Scalar division.

        Args:
            scalar: A scalar value to divide by.

        Returns:
            A new ArithmeticStateDict with each element divided by scalar.

        Raises:
            ZeroDivisionError: If scalar is zero.
            TypeError: If scalar is not a number.
        """
        if not isinstance(scalar, (int, float, Number)):
            raise TypeError(
                f"Cannot divide ArithmeticStateDict by {type(scalar).__name__}"
            )

        result_dict = state_dict_div(self, scalar)
        return ArithmeticStateDict(result_dict)

    def __pow__(self, exponent: Number) -> "ArithmeticStateDict":
        """
        Element-wise power operation.

        Args:
            exponent: The exponent to raise each element to.

        Returns:
            A new ArithmeticStateDict with each element raised to the power.
        """
        if not isinstance(exponent, (int, float, Number)):
            raise TypeError(
                f"Cannot raise ArithmeticStateDict to power of {type(exponent).__name__}"
            )

        result_dict = state_dict_power(self, exponent)
        return ArithmeticStateDict(result_dict)

    def __matmul__(
        self, other: Union["ArithmeticStateDict", StateDictType]
    ) -> "ArithmeticStateDict":
        """
        Hadamard product (element-wise multiplication) using @ operator.

        Args:
            other: Another state dict for element-wise multiplication.

        Returns:
            A new ArithmeticStateDict with the Hadamard product.
        """
        if not isinstance(other, (dict, OrderedDict)):
            raise TypeError(
                f"Cannot compute Hadamard product with {type(other).__name__}"
            )

        result_dict = state_dict_hadamard_product(self, other)
        return ArithmeticStateDict(result_dict)

    def __rmatmul__(
        self, other: Union["ArithmeticStateDict", StateDictType]
    ) -> "ArithmeticStateDict":
        """Right matrix multiplication (other @ self)."""
        return self.__matmul__(other)

    def __iadd__(
        self, other: Union["ArithmeticStateDict", StateDictType, Number]
    ) -> "ArithmeticStateDict":
        """In-place addition."""
        if isinstance(other, (int, float, Number)):
            # Scalar addition
            for key in self:
                self[key] = self[key] + other
        elif isinstance(other, (dict, OrderedDict)):
            # State dict addition
            for key in self:
                if key in other:
                    self[key] = self[key] + other[key]
        else:
            raise TypeError(f"Cannot add {type(other).__name__} to ArithmeticStateDict")
        return self

    def __isub__(
        self, other: Union["ArithmeticStateDict", StateDictType, Number]
    ) -> "ArithmeticStateDict":
        """In-place subtraction."""
        if isinstance(other, (int, float, Number)):
            # Scalar subtraction
            for key in self:
                self[key] = self[key] - other
        elif isinstance(other, (dict, OrderedDict)):
            # State dict subtraction
            for key in self:
                if key in other:
                    self[key] = self[key] - other[key]
        else:
            raise TypeError(
                f"Cannot subtract {type(other).__name__} from ArithmeticStateDict"
            )
        return self

    def __imul__(
        self, scalar: Union[Number, "ArithmeticStateDict", StateDictType]
    ) -> "ArithmeticStateDict":
        """In-place multiplication."""
        if isinstance(scalar, (int, float, Number)):
            for key in self:
                self[key] = self[key] * scalar
        elif isinstance(scalar, (dict, OrderedDict)):
            for key in self:
                if key in scalar:
                    self[key] = self[key] * scalar[key]
        else:
            raise TypeError(
                f"Cannot multiply ArithmeticStateDict with {type(scalar).__name__}"
            )
        return self

    def __itruediv__(self, scalar: Number) -> "ArithmeticStateDict":
        """In-place division."""
        if not isinstance(scalar, (int, float, Number)):
            raise TypeError(
                f"Cannot divide ArithmeticStateDict by {type(scalar).__name__}"
            )
        if scalar == 0:
            raise ZeroDivisionError("Cannot divide by zero")

        for key in self:
            self[key] = self[key] / scalar
        return self

    def __ipow__(self, exponent: Number) -> "ArithmeticStateDict":
        """In-place power operation."""
        if not isinstance(exponent, (int, float, Number)):
            raise TypeError(
                f"Cannot raise ArithmeticStateDict to power of {type(exponent).__name__}"
            )

        for key in self:
            self[key] = self[key] ** exponent
        return self

    def abs(self) -> "ArithmeticStateDict":
        """
        Element-wise absolute value.

        Returns:
            A new ArithmeticStateDict with absolute values.
        """
        result = ArithmeticStateDict()
        for key, tensor in self.items():
            result[key] = torch.abs(tensor)
        return result

    def sqrt(self) -> "ArithmeticStateDict":
        """
        Element-wise square root.

        Returns:
            A new ArithmeticStateDict with square roots.
        """
        result = ArithmeticStateDict()
        for key, tensor in self.items():
            result[key] = torch.sqrt(tensor)
        return result

    def sum(self) -> "ArithmeticStateDict":
        """
        Sum with other ArithmeticStateDicts using the + operator.

        Args:
            *others: Other ArithmeticStateDicts to sum with.

        Returns:
            A new ArithmeticStateDict with the sum.
        """
        # This is used for when sum() is called on a list of ArithmeticStateDicts
        return self

    def to_device(
        self,
        device: Union[torch.device, str],
        copy: bool = False,
        inplace: bool = False,
    ) -> "ArithmeticStateDict":
        """
        Move all tensors to the specified device.

        Args:
            device: Target device.
            copy: Whether to force a copy.
            inplace: Whether to modify in place.

        Returns:
            ArithmeticStateDict with tensors on the target device.
        """
        if inplace:
            for key, tensor in self.items():
                self[key] = tensor.to(device, non_blocking=True, copy=copy)
            return self
        else:
            result = ArithmeticStateDict()
            for key, tensor in self.items():
                result[key] = tensor.to(device, non_blocking=True, copy=copy)
            return result

    def clone(self) -> "ArithmeticStateDict":
        """
        Create a deep copy with cloned tensors.

        Returns:
            A new ArithmeticStateDict with cloned tensors.
        """
        result = ArithmeticStateDict()
        for key, tensor in self.items():
            result[key] = tensor.clone()
        return result

    def detach(self) -> "ArithmeticStateDict":
        """
        Detach all tensors from the computation graph.

        Returns:
            A new ArithmeticStateDict with detached tensors.
        """
        result = ArithmeticStateDict()
        for key, tensor in self.items():
            result[key] = tensor.detach()
        return result

    def num_params(self) -> int:
        """
        Calculate the total number of parameters.

        Returns:
            Total number of parameters in all tensors.
        """
        return sum(tensor.numel() for tensor in self.values())

    @classmethod
    def from_state_dict(cls, state_dict: StateDictType) -> "ArithmeticStateDict":
        """
        Create an ArithmeticStateDict from a regular state dict.

        Args:
            state_dict: A regular state dictionary.

        Returns:
            A new ArithmeticStateDict with the same data.
        """
        return cls(state_dict)

    @classmethod
    def weighted_sum(
        cls,
        state_dicts: List[Union["ArithmeticStateDict", StateDictType]],
        weights: List[float],
    ) -> "ArithmeticStateDict":
        """
        Compute a weighted sum of multiple state dicts.

        Args:
            state_dicts: List of state dicts to combine.
            weights: List of weights for the combination.

        Returns:
            A new ArithmeticStateDict with the weighted sum.
        """
        result_dict = state_dict_weighted_sum(state_dicts, weights)
        return cls(result_dict)

    @classmethod
    def average(
        cls, state_dicts: List[Union["ArithmeticStateDict", StateDictType]]
    ) -> "ArithmeticStateDict":
        """
        Compute the average of multiple state dicts.

        Args:
            state_dicts: List of state dicts to average.

        Returns:
            A new ArithmeticStateDict with the average.
        """
        result_dict = state_dict_avg(state_dicts)
        return cls(result_dict)

__add__(other)

Element-wise addition with another state dict or scalar.

Parameters:

• other (Union[ArithmeticStateDict, StateDictType, Number]) –

  Another state dict to add or a scalar to add to all elements.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the element-wise sum.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __add__(
    self, other: Union["ArithmeticStateDict", StateDictType, Number]
) -> "ArithmeticStateDict":
    """
    Element-wise addition with another state dict or scalar.

    Args:
        other: Another state dict to add or a scalar to add to all elements.

    Returns:
        A new ArithmeticStateDict with the element-wise sum.
    """
    if isinstance(other, (int, float, Number)):
        # Scalar addition
        result_dict = state_dict_add_scalar(self, other)
        return ArithmeticStateDict(result_dict)
    elif isinstance(other, (dict, OrderedDict)):
        # State dict addition
        result_dict = state_dict_add(self, other, strict=True)
        return ArithmeticStateDict(result_dict)
    else:
        raise TypeError(
            f"Cannot add ArithmeticStateDict with {type(other).__name__}"
        )

__iadd__(other)

In-place addition.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __iadd__(
    self, other: Union["ArithmeticStateDict", StateDictType, Number]
) -> "ArithmeticStateDict":
    """In-place addition."""
    if isinstance(other, (int, float, Number)):
        # Scalar addition
        for key in self:
            self[key] = self[key] + other
    elif isinstance(other, (dict, OrderedDict)):
        # State dict addition
        for key in self:
            if key in other:
                self[key] = self[key] + other[key]
    else:
        raise TypeError(f"Cannot add {type(other).__name__} to ArithmeticStateDict")
    return self

__imul__(scalar)

In-place multiplication.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __imul__(
    self, scalar: Union[Number, "ArithmeticStateDict", StateDictType]
) -> "ArithmeticStateDict":
    """In-place multiplication."""
    if isinstance(scalar, (int, float, Number)):
        for key in self:
            self[key] = self[key] * scalar
    elif isinstance(scalar, (dict, OrderedDict)):
        for key in self:
            if key in scalar:
                self[key] = self[key] * scalar[key]
    else:
        raise TypeError(
            f"Cannot multiply ArithmeticStateDict with {type(scalar).__name__}"
        )
    return self

__init__(*args, **kwargs)

Initialize ArithmeticStateDict with the same interface as OrderedDict.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __init__(self, *args, **kwargs):
    """Initialize ArithmeticStateDict with the same interface as OrderedDict."""
    super().__init__(*args, **kwargs)

__ipow__(exponent)

In-place power operation.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __ipow__(self, exponent: Number) -> "ArithmeticStateDict":
    """In-place power operation."""
    if not isinstance(exponent, (int, float, Number)):
        raise TypeError(
            f"Cannot raise ArithmeticStateDict to power of {type(exponent).__name__}"
        )

    for key in self:
        self[key] = self[key] ** exponent
    return self

__isub__(other)

In-place subtraction.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __isub__(
    self, other: Union["ArithmeticStateDict", StateDictType, Number]
) -> "ArithmeticStateDict":
    """In-place subtraction."""
    if isinstance(other, (int, float, Number)):
        # Scalar subtraction
        for key in self:
            self[key] = self[key] - other
    elif isinstance(other, (dict, OrderedDict)):
        # State dict subtraction
        for key in self:
            if key in other:
                self[key] = self[key] - other[key]
    else:
        raise TypeError(
            f"Cannot subtract {type(other).__name__} from ArithmeticStateDict"
        )
    return self

__itruediv__(scalar)

In-place division.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __itruediv__(self, scalar: Number) -> "ArithmeticStateDict":
    """In-place division."""
    if not isinstance(scalar, (int, float, Number)):
        raise TypeError(
            f"Cannot divide ArithmeticStateDict by {type(scalar).__name__}"
        )
    if scalar == 0:
        raise ZeroDivisionError("Cannot divide by zero")

    for key in self:
        self[key] = self[key] / scalar
    return self

__matmul__(other)

Hadamard product (element-wise multiplication) using @ operator.

Parameters:

• other (Union[ArithmeticStateDict, StateDictType]) –

  Another state dict for element-wise multiplication.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the Hadamard product.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __matmul__(
    self, other: Union["ArithmeticStateDict", StateDictType]
) -> "ArithmeticStateDict":
    """
    Hadamard product (element-wise multiplication) using @ operator.

    Args:
        other: Another state dict for element-wise multiplication.

    Returns:
        A new ArithmeticStateDict with the Hadamard product.
    """
    if not isinstance(other, (dict, OrderedDict)):
        raise TypeError(
            f"Cannot compute Hadamard product with {type(other).__name__}"
        )

    result_dict = state_dict_hadamard_product(self, other)
    return ArithmeticStateDict(result_dict)

__mul__(scalar)

Scalar multiplication or Hadamard product.

Parameters:

• scalar (Union[Number, ArithmeticStateDict, StateDictType]) –

  A scalar value for element-wise multiplication, or another state dict for Hadamard product.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the result.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __mul__(
    self, scalar: Union[Number, "ArithmeticStateDict", StateDictType]
) -> "ArithmeticStateDict":
    """
    Scalar multiplication or Hadamard product.

    Args:
        scalar: A scalar value for element-wise multiplication, or another state dict
               for Hadamard product.

    Returns:
        A new ArithmeticStateDict with the result.
    """
    if isinstance(scalar, (int, float, Number)):
        result_dict = state_dict_mul(self, scalar)
        return ArithmeticStateDict(result_dict)
    elif isinstance(scalar, (dict, OrderedDict)):
        # Hadamard product for dict-like objects
        result_dict = state_dict_hadamard_product(self, scalar)
        return ArithmeticStateDict(result_dict)
    else:
        raise TypeError(
            f"Cannot multiply ArithmeticStateDict with {type(scalar).__name__}"
        )

__pow__(exponent)

Element-wise power operation.

Parameters:

• exponent (Number) –

  The exponent to raise each element to.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with each element raised to the power.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __pow__(self, exponent: Number) -> "ArithmeticStateDict":
    """
    Element-wise power operation.

    Args:
        exponent: The exponent to raise each element to.

    Returns:
        A new ArithmeticStateDict with each element raised to the power.
    """
    if not isinstance(exponent, (int, float, Number)):
        raise TypeError(
            f"Cannot raise ArithmeticStateDict to power of {type(exponent).__name__}"
        )

    result_dict = state_dict_power(self, exponent)
    return ArithmeticStateDict(result_dict)

__radd__(other)

Right addition (other + self). Handles the case where sum() starts with 0 and scalar addition.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __radd__(
    self, other: Union["ArithmeticStateDict", StateDictType, Number]
) -> "ArithmeticStateDict":
    """
    Right addition (other + self).
    Handles the case where sum() starts with 0 and scalar addition.
    """
    if other == 0:  # sum() starts with 0 by default
        return self
    elif isinstance(other, (int, float, Number)):
        # Scalar addition is commutative
        return self.__add__(other)
    elif isinstance(other, (dict, OrderedDict)):
        return self.__add__(other)
    else:
        raise TypeError(
            f"Cannot add {type(other).__name__} with ArithmeticStateDict"
        )

__rmatmul__(other)

Right matrix multiplication (other @ self).

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __rmatmul__(
    self, other: Union["ArithmeticStateDict", StateDictType]
) -> "ArithmeticStateDict":
    """Right matrix multiplication (other @ self)."""
    return self.__matmul__(other)

__rmul__(scalar)

Right multiplication (scalar * self).

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __rmul__(
    self, scalar: Union[Number, "ArithmeticStateDict", StateDictType]
) -> "ArithmeticStateDict":
    """Right multiplication (scalar * self)."""
    return self.__mul__(scalar)

__rsub__(other)

Right subtraction (other - self).

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __rsub__(
    self, other: Union["ArithmeticStateDict", StateDictType, Number]
) -> "ArithmeticStateDict":
    """Right subtraction (other - self)."""
    if isinstance(other, (int, float, Number)):
        # Scalar - ArithmeticStateDict: subtract each element from scalar
        result = ArithmeticStateDict()
        for key, tensor in self.items():
            result[key] = other - tensor
        return result
    elif isinstance(other, (dict, OrderedDict)):
        result_dict = state_dict_sub(other, self, strict=True)
        return ArithmeticStateDict(result_dict)
    else:
        raise TypeError(
            f"Cannot subtract ArithmeticStateDict from {type(other).__name__}"
        )

__sub__(other)

Element-wise subtraction with another state dict or scalar.

Parameters:

• other (Union[ArithmeticStateDict, StateDictType, Number]) –

  Another state dict to subtract or a scalar to subtract from all elements.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the element-wise difference.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __sub__(
    self, other: Union["ArithmeticStateDict", StateDictType, Number]
) -> "ArithmeticStateDict":
    """
    Element-wise subtraction with another state dict or scalar.

    Args:
        other: Another state dict to subtract or a scalar to subtract from all elements.

    Returns:
        A new ArithmeticStateDict with the element-wise difference.
    """
    if isinstance(other, (int, float, Number)):
        # Scalar subtraction: subtract scalar from all elements
        result_dict = state_dict_add_scalar(self, -other)
        return ArithmeticStateDict(result_dict)
    elif isinstance(other, (dict, OrderedDict)):
        # State dict subtraction
        result_dict = state_dict_sub(self, other, strict=True)
        return ArithmeticStateDict(result_dict)
    else:
        raise TypeError(
            f"Cannot subtract {type(other).__name__} from ArithmeticStateDict"
        )

__truediv__(scalar)

Scalar division.

Parameters:

• scalar (Number) –

  A scalar value to divide by.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with each element divided by scalar.

Raises:

• ZeroDivisionError –

  If scalar is zero.

• TypeError –

  If scalar is not a number.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def __truediv__(self, scalar: Number) -> "ArithmeticStateDict":
    """
    Scalar division.

    Args:
        scalar: A scalar value to divide by.

    Returns:
        A new ArithmeticStateDict with each element divided by scalar.

    Raises:
        ZeroDivisionError: If scalar is zero.
        TypeError: If scalar is not a number.
    """
    if not isinstance(scalar, (int, float, Number)):
        raise TypeError(
            f"Cannot divide ArithmeticStateDict by {type(scalar).__name__}"
        )

    result_dict = state_dict_div(self, scalar)
    return ArithmeticStateDict(result_dict)

abs()

Element-wise absolute value.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with absolute values.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def abs(self) -> "ArithmeticStateDict":
    """
    Element-wise absolute value.

    Returns:
        A new ArithmeticStateDict with absolute values.
    """
    result = ArithmeticStateDict()
    for key, tensor in self.items():
        result[key] = torch.abs(tensor)
    return result

average(state_dicts) classmethod

Compute the average of multiple state dicts.

Parameters:

• state_dicts (List[Union[ArithmeticStateDict, StateDictType]]) –

  List of state dicts to average.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the average.

Source code in fusion_bench/utils/state_dict_arithmetic.py

@classmethod
def average(
    cls, state_dicts: List[Union["ArithmeticStateDict", StateDictType]]
) -> "ArithmeticStateDict":
    """
    Compute the average of multiple state dicts.

    Args:
        state_dicts: List of state dicts to average.

    Returns:
        A new ArithmeticStateDict with the average.
    """
    result_dict = state_dict_avg(state_dicts)
    return cls(result_dict)

clone()

Create a deep copy with cloned tensors.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with cloned tensors.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def clone(self) -> "ArithmeticStateDict":
    """
    Create a deep copy with cloned tensors.

    Returns:
        A new ArithmeticStateDict with cloned tensors.
    """
    result = ArithmeticStateDict()
    for key, tensor in self.items():
        result[key] = tensor.clone()
    return result

detach()

Detach all tensors from the computation graph.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with detached tensors.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def detach(self) -> "ArithmeticStateDict":
    """
    Detach all tensors from the computation graph.

    Returns:
        A new ArithmeticStateDict with detached tensors.
    """
    result = ArithmeticStateDict()
    for key, tensor in self.items():
        result[key] = tensor.detach()
    return result

from_state_dict(state_dict) classmethod

Create an ArithmeticStateDict from a regular state dict.

Parameters:

• state_dict (StateDictType) –

  A regular state dictionary.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the same data.

Source code in fusion_bench/utils/state_dict_arithmetic.py

@classmethod
def from_state_dict(cls, state_dict: StateDictType) -> "ArithmeticStateDict":
    """
    Create an ArithmeticStateDict from a regular state dict.

    Args:
        state_dict: A regular state dictionary.

    Returns:
        A new ArithmeticStateDict with the same data.
    """
    return cls(state_dict)

num_params()

Calculate the total number of parameters.

Returns:

• int –

  Total number of parameters in all tensors.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def num_params(self) -> int:
    """
    Calculate the total number of parameters.

    Returns:
        Total number of parameters in all tensors.
    """
    return sum(tensor.numel() for tensor in self.values())

sqrt()

Element-wise square root.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with square roots.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def sqrt(self) -> "ArithmeticStateDict":
    """
    Element-wise square root.

    Returns:
        A new ArithmeticStateDict with square roots.
    """
    result = ArithmeticStateDict()
    for key, tensor in self.items():
        result[key] = torch.sqrt(tensor)
    return result

sum()

Sum with other ArithmeticStateDicts using the + operator.

Parameters:

• *others –

  Other ArithmeticStateDicts to sum with.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the sum.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def sum(self) -> "ArithmeticStateDict":
    """
    Sum with other ArithmeticStateDicts using the + operator.

    Args:
        *others: Other ArithmeticStateDicts to sum with.

    Returns:
        A new ArithmeticStateDict with the sum.
    """
    # This is used for when sum() is called on a list of ArithmeticStateDicts
    return self

to_device(device, copy=False, inplace=False)

Move all tensors to the specified device.

Parameters:

• device (Union[device, str]) –

  Target device.

• copy (bool, default: False ) –

  Whether to force a copy.

• inplace (bool, default: False ) –

  Whether to modify in place.

Returns:

• ArithmeticStateDict –

  ArithmeticStateDict with tensors on the target device.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def to_device(
    self,
    device: Union[torch.device, str],
    copy: bool = False,
    inplace: bool = False,
) -> "ArithmeticStateDict":
    """
    Move all tensors to the specified device.

    Args:
        device: Target device.
        copy: Whether to force a copy.
        inplace: Whether to modify in place.

    Returns:
        ArithmeticStateDict with tensors on the target device.
    """
    if inplace:
        for key, tensor in self.items():
            self[key] = tensor.to(device, non_blocking=True, copy=copy)
        return self
    else:
        result = ArithmeticStateDict()
        for key, tensor in self.items():
            result[key] = tensor.to(device, non_blocking=True, copy=copy)
        return result

weighted_sum(state_dicts, weights) classmethod

Compute a weighted sum of multiple state dicts.

Parameters:

• state_dicts (List[Union[ArithmeticStateDict, StateDictType]]) –

  List of state dicts to combine.

• weights (List[float]) –

  List of weights for the combination.

Returns:

• ArithmeticStateDict –

  A new ArithmeticStateDict with the weighted sum.

Source code in fusion_bench/utils/state_dict_arithmetic.py

@classmethod
def weighted_sum(
    cls,
    state_dicts: List[Union["ArithmeticStateDict", StateDictType]],
    weights: List[float],
) -> "ArithmeticStateDict":
    """
    Compute a weighted sum of multiple state dicts.

    Args:
        state_dicts: List of state dicts to combine.
        weights: List of weights for the combination.

    Returns:
        A new ArithmeticStateDict with the weighted sum.
    """
    result_dict = state_dict_weighted_sum(state_dicts, weights)
    return cls(result_dict)

num_params_of_state_dict(state_dict)

Calculate the total number of parameters in a state dict.

Parameters:

• state_dict (StateDictType) –

  The state dict to count parameters in.

Returns:

• int –

  The total number of parameters in the state dict.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def num_params_of_state_dict(state_dict: StateDictType) -> int:
    """
    Calculate the total number of parameters in a state dict.

    Args:
        state_dict: The state dict to count parameters in.

    Returns:
        The total number of parameters in the state dict.
    """
    return sum(tensor.numel() for tensor in state_dict.values())

state_dict_add(a, b, strict=True, device=None, show_pbar=False)

Compute the element-wise sum of two state dicts.

Parameters:

• a (StateDictType) –

  The first state dict.

• b (StateDictType) –

  The second state dict.

• strict (bool, default: True ) –

  Whether to require exact key matching between state dicts.

• device (Optional[Union[device, str]], default: None ) –

  Optional device to move the result tensors to.

• show_pbar (bool, default: False ) –

  Whether to show a progress bar during computation.

Returns:

• StateDictType –

  A state dict containing the element-wise sums.

Raises:

• ValueError –

  If strict=True and the state dicts have different parameters.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_add(
    a: StateDictType,
    b: StateDictType,
    strict: bool = True,
    device: Optional[Union[torch.device, str]] = None,
    show_pbar: bool = False,
) -> StateDictType:
    """
    Compute the element-wise sum of two state dicts.

    Args:
        a: The first state dict.
        b: The second state dict.
        strict: Whether to require exact key matching between state dicts.
        device: Optional device to move the result tensors to.
        show_pbar: Whether to show a progress bar during computation.

    Returns:
        A state dict containing the element-wise sums.

    Raises:
        ValueError: If strict=True and the state dicts have different parameters.
    """
    result = OrderedDict()

    if strict:
        _validate_state_dict_same_keys([a, b])
        keys_to_process = a.keys()
    else:
        keys_to_process = set(a.keys()) & set(b.keys())

    keys_iter = (
        tqdm(keys_to_process, desc="Adding state dicts")
        if show_pbar
        else keys_to_process
    )

    for key in keys_iter:
        if key in b:  # This check is redundant when strict=True but harmless
            result[key] = a[key] + b[key]

    if device is not None:
        result = state_dict_to_device(result, device)

    return result

state_dict_add_scalar(state_dict, scalar)

Add a scalar value to all parameters in a state dict.

Parameters:

• state_dict (StateDictType) –

  The state dict to modify.

• scalar (Number) –

  The scalar value to add to each parameter.

Returns:

• StateDictType –

  A new state dict with the scalar added to each parameter.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_add_scalar(state_dict: StateDictType, scalar: Number) -> StateDictType:
    """
    Add a scalar value to all parameters in a state dict.

    Args:
        state_dict: The state dict to modify.
        scalar: The scalar value to add to each parameter.

    Returns:
        A new state dict with the scalar added to each parameter.
    """
    return OrderedDict((key, tensor + scalar) for key, tensor in state_dict.items())

state_dict_avg(state_dicts)

Calculate the element-wise average of a list of state dicts.

Parameters:

• state_dicts (List[StateDictType]) –

  List of state dicts to average.

Returns:

• StateDictType –

  A state dict containing the averaged parameters.

Raises:

• ValueError –

  If the list is empty or state dicts have different keys.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_avg(state_dicts: List[StateDictType]) -> StateDictType:
    """
    Calculate the element-wise average of a list of state dicts.

    Args:
        state_dicts: List of state dicts to average.

    Returns:
        A state dict containing the averaged parameters.

    Raises:
        ValueError: If the list is empty or state dicts have different keys.
    """
    _validate_state_dict_list_not_empty(state_dicts)
    _validate_state_dict_same_keys(state_dicts)

    num_state_dicts = len(state_dicts)
    avg_state_dict = OrderedDict()

    # Initialize with zeros_like for better performance
    for key in state_dicts[0]:
        avg_state_dict[key] = torch.zeros_like(state_dicts[0][key])

    # Accumulate all state dicts
    for state_dict in state_dicts:
        for key in avg_state_dict:
            avg_state_dict[key] += state_dict[key]

    # Divide by number of state dicts
    for key in avg_state_dict:
        avg_state_dict[key] /= num_state_dicts

    return avg_state_dict

state_dict_binary_mask(a, b, compare_fn='greater', strict=True, show_pbar=False)

Create binary masks by comparing elements in two state dicts.

Parameters:

• a (StateDictType) –

  The first state dict.

• b (StateDictType) –

  The second state dict.

• compare_fn (Union[Literal['greater', 'less', 'equal', 'not_equal'], Callable[[Tensor, Tensor], BoolTensor]], default: 'greater' ) –

  Comparison function to use. Can be a string literal ("greater", "less", "equal", "not_equal") or a callable that takes two tensors and returns a boolean tensor.

• strict (bool, default: True ) –

  Whether to require exact key matching between state dicts.

• show_pbar (bool, default: False ) –

  Whether to show a progress bar during computation.

Returns:

• BoolStateDictType –

  A dictionary containing boolean masks based on the comparison.

Raises:

• ValueError –

  If compare_fn is not a valid string or callable, or if strict=True and the state dicts have different keys or incompatible tensor shapes.

• TypeError –

  If tensors have incompatible types.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_binary_mask(
    a: StateDictType,
    b: StateDictType,
    compare_fn: Union[
        Literal["greater", "less", "equal", "not_equal"],
        Callable[[Tensor, Tensor], torch.BoolTensor],
    ] = "greater",
    strict: bool = True,
    show_pbar: bool = False,
) -> BoolStateDictType:
    """
    Create binary masks by comparing elements in two state dicts.

    Args:
        a: The first state dict.
        b: The second state dict.
        compare_fn: Comparison function to use. Can be a string literal
                   ("greater", "less", "equal", "not_equal") or a callable
                   that takes two tensors and returns a boolean tensor.
        strict: Whether to require exact key matching between state dicts.
        show_pbar: Whether to show a progress bar during computation.

    Returns:
        A dictionary containing boolean masks based on the comparison.

    Raises:
        ValueError: If compare_fn is not a valid string or callable, or if strict=True
                   and the state dicts have different keys or incompatible tensor shapes.
        TypeError: If tensors have incompatible types.
    """
    compare_fn_dict = {
        "greater": lambda x, y: x > y,
        "less": lambda x, y: x < y,
        "equal": lambda x, y: x == y,
        "not_equal": lambda x, y: x != y,
    }

    if isinstance(compare_fn, str):
        if compare_fn not in compare_fn_dict:
            raise ValueError(
                f"Invalid compare_fn string: {compare_fn}. Must be one of {list(compare_fn_dict.keys())}"
            )
        compare_fn = compare_fn_dict[compare_fn]
    elif not callable(compare_fn):
        raise ValueError(
            f"compare_fn must be a string or a callable, but got {type(compare_fn)}"
        )

    result = OrderedDict()

    if strict:
        _validate_state_dict_same_keys([a, b])
        keys_to_process = a.keys()
    else:
        keys_to_process = set(a.keys()) & set(b.keys())

    keys_iter = (
        tqdm(keys_to_process, desc="Creating binary masks")
        if show_pbar
        else keys_to_process
    )

    for key in keys_iter:
        result[key] = compare_fn(a[key], b[key])

    return result

state_dict_diff_abs(a, b)

Compute the element-wise absolute difference between two state dicts.

Parameters:

• a (StateDictType) –

  The first state dict.

• b (StateDictType) –

  The second state dict.

Returns:

• StateDictType –

  A state dict containing the absolute differences.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_diff_abs(a: StateDictType, b: StateDictType) -> StateDictType:
    """
    Compute the element-wise absolute difference between two state dicts.

    Args:
        a: The first state dict.
        b: The second state dict.

    Returns:
        A state dict containing the absolute differences.
    """
    diff = state_dict_sub(a, b)
    return OrderedDict((key, tensor.abs()) for key, tensor in diff.items())

state_dict_div(state_dict, scalar, show_pbar=False)

Divide all parameters in a state dict by a scalar.

Parameters:

• state_dict (StateDictType) –

  The state dict to divide.

• scalar (float) –

  The scalar value to divide each parameter by.

• show_pbar (bool, default: False ) –

  Whether to show a progress bar during computation.

Returns:

• StateDictType –

  A new state dict with each parameter divided by the scalar.

Raises:

• ZeroDivisionError –

  If scalar is zero.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_div(
    state_dict: StateDictType, scalar: float, show_pbar: bool = False
) -> StateDictType:
    """
    Divide all parameters in a state dict by a scalar.

    Args:
        state_dict: The state dict to divide.
        scalar: The scalar value to divide each parameter by.
        show_pbar: Whether to show a progress bar during computation.

    Returns:
        A new state dict with each parameter divided by the scalar.

    Raises:
        ZeroDivisionError: If scalar is zero.
    """
    if scalar == 0:
        raise ZeroDivisionError("Cannot divide state dict by zero")

    keys_iter = (
        tqdm(state_dict.keys(), desc="Dividing state dict")
        if show_pbar
        else state_dict.keys()
    )
    return OrderedDict((key, state_dict[key] / scalar) for key in keys_iter)

state_dict_flatten(state_dict)

Flatten all tensors in a state dict into a single 1D tensor.

Parameters:

• state_dict (StateDictType) –

  The state dict to flatten.

Returns:

• Tensor –

  A single flattened tensor containing all parameters.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_flatten(state_dict: StateDictType) -> Tensor:
    """
    Flatten all tensors in a state dict into a single 1D tensor.

    Args:
        state_dict: The state dict to flatten.

    Returns:
        A single flattened tensor containing all parameters.
    """
    return torch.cat([tensor.flatten() for tensor in state_dict.values()])

state_dict_hadamard_product(a, b)

Compute the Hadamard product (element-wise multiplication) of two state dicts.

Parameters:

• a (StateDictType) –

  The first state dict.

• b (StateDictType) –

  The second state dict.

Returns:

• StateDictType –

  A state dict containing the element-wise products.

Raises:

• ValueError –

  If the state dicts have different keys or incompatible tensor shapes.

• TypeError –

  If tensors have incompatible types.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_hadamard_product(a: StateDictType, b: StateDictType) -> StateDictType:
    """
    Compute the Hadamard product (element-wise multiplication) of two state dicts.

    Args:
        a: The first state dict.
        b: The second state dict.

    Returns:
        A state dict containing the element-wise products.

    Raises:
        ValueError: If the state dicts have different keys or incompatible tensor shapes.
        TypeError: If tensors have incompatible types.
    """
    _validate_state_dict_same_keys([a, b])
    return OrderedDict((key, a[key] * b[key]) for key in a)

state_dict_interpolation(state_dicts, scalars)

Interpolate between multiple state dicts using specified scalar weights.

Parameters:

• state_dicts (List[StateDictType]) –

  List of state dicts to interpolate between.

• scalars (List[float]) –

  List of scalar weights for interpolation.

Returns:

• StateDictType –

  A state dict containing the interpolated parameters.

Raises:

• ValueError –

  If the lists have different lengths or are empty, or if state dicts have different keys.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_interpolation(
    state_dicts: List[StateDictType], scalars: List[float]
) -> StateDictType:
    """
    Interpolate between multiple state dicts using specified scalar weights.

    Args:
        state_dicts: List of state dicts to interpolate between.
        scalars: List of scalar weights for interpolation.

    Returns:
        A state dict containing the interpolated parameters.

    Raises:
        ValueError: If the lists have different lengths or are empty, or if state dicts have different keys.
    """
    _validate_state_dict_list_not_empty(state_dicts)
    _validate_list_lengths_equal(state_dicts, scalars, "state_dicts", "scalars")
    _validate_state_dict_same_keys(state_dicts)

    interpolated_state_dict = OrderedDict()

    # Initialize with zeros
    for key in state_dicts[0]:
        interpolated_state_dict[key] = torch.zeros_like(state_dicts[0][key])

    # Accumulate weighted contributions
    for state_dict, scalar in zip(state_dicts, scalars):
        for key in interpolated_state_dict:
            interpolated_state_dict[key] += scalar * state_dict[key]

    return interpolated_state_dict

state_dict_mul(state_dict, scalar)

Multiply all parameters in a state dict by a scalar.

Parameters:

• state_dict (StateDictType) –

  The state dict to multiply.

• scalar (float) –

  The scalar value to multiply each parameter by.

Returns:

• StateDictType –

  A new state dict with each parameter multiplied by the scalar.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_mul(state_dict: StateDictType, scalar: float) -> StateDictType:
    """
    Multiply all parameters in a state dict by a scalar.

    Args:
        state_dict: The state dict to multiply.
        scalar: The scalar value to multiply each parameter by.

    Returns:
        A new state dict with each parameter multiplied by the scalar.
    """
    return OrderedDict((key, scalar * tensor) for key, tensor in state_dict.items())

state_dict_power(state_dict, p)

Raise all parameters in a state dict to a power.

Parameters:

• state_dict (StateDictType) –

  The state dict to raise to a power.

• p (float) –

  The exponent to raise each parameter to.

Returns:

• StateDictType –

  A new state dict with each parameter raised to the power p.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_power(state_dict: StateDictType, p: float) -> StateDictType:
    """
    Raise all parameters in a state dict to a power.

    Args:
        state_dict: The state dict to raise to a power.
        p: The exponent to raise each parameter to.

    Returns:
        A new state dict with each parameter raised to the power p.
    """
    return OrderedDict((key, tensor**p) for key, tensor in state_dict.items())

state_dict_sub(a, b, strict=True, device=None)

Compute the element-wise difference between two state dicts (a - b).

Parameters:

• a (StateDictType) –

  The first state dict (minuend).

• b (StateDictType) –

  The second state dict (subtrahend).

• strict (bool, default: True ) –

  Whether to require exact key matching between state dicts.

• device (Optional[Union[device, str]], default: None ) –

  Optional device to move the result tensors to.

Returns:

• StateDictType –

  A state dict containing the element-wise differences.

Raises:

• ValueError –

  If strict=True and the state dicts have different keys or incompatible tensor shapes.

• TypeError –

  If tensors have incompatible types.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_sub(
    a: StateDictType,
    b: StateDictType,
    strict: bool = True,
    device: Optional[Union[torch.device, str]] = None,
) -> StateDictType:
    """
    Compute the element-wise difference between two state dicts (a - b).

    Args:
        a: The first state dict (minuend).
        b: The second state dict (subtrahend).
        strict: Whether to require exact key matching between state dicts.
        device: Optional device to move the result tensors to.

    Returns:
        A state dict containing the element-wise differences.

    Raises:
        ValueError: If strict=True and the state dicts have different keys or incompatible tensor shapes.
        TypeError: If tensors have incompatible types.
    """
    result = OrderedDict()

    if strict:
        _validate_state_dict_same_keys([a, b])
        keys_to_process = a.keys()
    else:
        keys_to_process = set(a.keys()) & set(b.keys())

    for key in keys_to_process:
        result_tensor = a[key] - b[key]
        if device is not None:
            result_tensor = result_tensor.to(device, non_blocking=True)
        result[key] = result_tensor

    return result

state_dict_sum(state_dicts)

Compute the element-wise sum of multiple state dicts.

Parameters:

• state_dicts (List[StateDictType]) –

  List of state dicts to sum.

Returns:

• StateDictType –

  A state dict containing the element-wise sums.

Raises:

• ValueError –

  If the list is empty or state dicts have different keys.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_sum(state_dicts: List[StateDictType]) -> StateDictType:
    """
    Compute the element-wise sum of multiple state dicts.

    Args:
        state_dicts: List of state dicts to sum.

    Returns:
        A state dict containing the element-wise sums.

    Raises:
        ValueError: If the list is empty or state dicts have different keys.
    """
    _validate_state_dict_list_not_empty(state_dicts)
    _validate_state_dict_same_keys(state_dicts)

    sum_state_dict = OrderedDict()

    # Initialize with zeros
    for key in state_dicts[0]:
        sum_state_dict[key] = torch.zeros_like(state_dicts[0][key])

    # Accumulate all state dicts
    for state_dict in state_dicts:
        for key in sum_state_dict:
            sum_state_dict[key] += state_dict[key]

    return sum_state_dict

state_dict_to_device(state_dict, device, copy=False, inplace=False)

Move state dict tensors to the specified device.

Parameters:

• state_dict (StateDictType) –

  The state dictionary to move.

• device (Union[device, str]) –

  Target device for the tensors.

• copy (bool, default: False ) –

  Whether to force a copy even when the tensor is already on the target device.

• inplace (bool, default: False ) –

  Whether to modify the input state dict in place.

Returns:

• StateDictType –

  State dict with tensors moved to the specified device.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_to_device(
    state_dict: StateDictType,
    device: Union[torch.device, str],
    copy: bool = False,
    inplace: bool = False,
) -> StateDictType:
    """
    Move state dict tensors to the specified device.

    Args:
        state_dict: The state dictionary to move.
        device: Target device for the tensors.
        copy: Whether to force a copy even when the tensor is already on the target device.
        inplace: Whether to modify the input state dict in place.

    Returns:
        State dict with tensors moved to the specified device.
    """
    if inplace:
        ret_state_dict = state_dict
    else:
        ret_state_dict = OrderedDict()

    for key, tensor in state_dict.items():
        ret_state_dict[key] = cast(Tensor, tensor).to(
            device, non_blocking=True, copy=copy
        )
    return ret_state_dict

state_dict_weighted_sum(state_dicts, weights, device=None)

Compute the weighted sum of multiple state dicts.

Parameters:

• state_dicts (List[StateDictType]) –

  List of state dicts to combine.

• weights (List[float]) –

  List of weights for the weighted sum.

• device (Optional[Union[device, str]], default: None ) –

  Optional device to move the result tensors to.

Returns:

• StateDictType –

  A state dict containing the weighted sum of parameters.

Raises:

• ValueError –

  If the lists have different lengths or are empty, or if state dicts have different keys.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dict_weighted_sum(
    state_dicts: List[StateDictType],
    weights: List[float],
    device: Optional[Union[torch.device, str]] = None,
) -> StateDictType:
    """
    Compute the weighted sum of multiple state dicts.

    Args:
        state_dicts: List of state dicts to combine.
        weights: List of weights for the weighted sum.
        device: Optional device to move the result tensors to.

    Returns:
        A state dict containing the weighted sum of parameters.

    Raises:
        ValueError: If the lists have different lengths or are empty, or if state dicts have different keys.
    """
    _validate_state_dict_list_not_empty(state_dicts)
    _validate_list_lengths_equal(state_dicts, weights, "state_dicts", "weights")
    _validate_state_dict_same_keys(state_dicts)

    weighted_sum_state_dict = OrderedDict()

    # Single pass initialization and computation for better performance
    for key in state_dicts[0]:
        # Get reference tensor and handle sparse tensors
        ref_tensor = state_dicts[0][key]
        is_sparse = ref_tensor.is_sparse if hasattr(ref_tensor, "is_sparse") else False

        # Initialize result tensor
        if is_sparse:
            # For sparse tensors, start with zeros in dense format for efficient accumulation
            result_tensor = torch.zeros_like(ref_tensor).to_dense()
        else:
            result_tensor = torch.zeros_like(ref_tensor)

        # Accumulate weighted contributions in a single loop
        for state_dict, weight in zip(state_dicts, weights):
            tensor = state_dict[key]

            # Optimize for common cases
            if weight == 0.0:
                continue  # Skip zero weights
            elif weight == 1.0:
                result_tensor += tensor  # Avoid multiplication for unit weights
            else:
                # Use in-place operations when possible for memory efficiency
                if is_sparse and hasattr(tensor, "is_sparse") and tensor.is_sparse:
                    result_tensor += weight * tensor.to_dense()
                else:
                    result_tensor += weight * tensor

        # Move to target device if specified (do this once per tensor, not per operation)
        if device is not None:
            result_tensor = result_tensor.to(device, non_blocking=True)

        # Convert back to sparse if original was sparse and result is suitable
        if is_sparse and hasattr(result_tensor, "to_sparse"):
            try:
                # Only convert back to sparse if it would be memory efficient
                # (i.e., if the result has sufficient sparsity)
                if result_tensor.numel() > 0:
                    sparsity_ratio = (result_tensor == 0).float().mean().item()
                    if sparsity_ratio > 0.5:  # Convert back if >50% zeros
                        result_tensor = result_tensor.to_sparse()
            except (RuntimeError, AttributeError):
                # If conversion fails, keep as dense
                pass

        weighted_sum_state_dict[key] = result_tensor

    return weighted_sum_state_dict

state_dicts_check_keys(state_dicts)

Check that all state dictionaries have the same keys.

Parameters:

• state_dicts (List[StateDictType]) –

  A list of state dictionaries to check.

Raises:

• ValueError –

  If the state dictionaries have different keys or the list is empty.

Source code in fusion_bench/utils/state_dict_arithmetic.py

def state_dicts_check_keys(state_dicts: List[StateDictType]) -> None:
    """
    Check that all state dictionaries have the same keys.

    Args:
        state_dicts: A list of state dictionaries to check.

    Raises:
        ValueError: If the state dictionaries have different keys or the list is empty.
    """
    _validate_state_dict_list_not_empty(state_dicts)
    _validate_state_dict_same_keys(state_dicts)

Lazy Model Loading

fusion_bench.utils.lazy_state_dict.LazyStateDict

Bases: Mapping[str, Tensor], Generic[TorchModelType]

A dictionary-like object that lazily loads tensors from model checkpoints.

Source code in fusion_bench/utils/lazy_state_dict.py

class LazyStateDict(Mapping[str, torch.Tensor], Generic[TorchModelType]):
    """
    A dictionary-like object that lazily loads tensors from model checkpoints.
    """

    _local_path: str
    """Local path to the checkpoint."""
    _state_dict_cache: Optional[Dict]
    """Cache for the state dict, if enabled."""
    _index_filename: Optional[str]
    _checkpoint_files: Optional[List[str]]
    _index: Optional[Dict[str, str]]
    """Mapping of parameter names to checkpoint files."""

    meta_module: TorchModelType = None
    meta_module_class: Optional[Type[TorchModelType]] = None

    def __init__(
        self,
        checkpoint: str,
        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",
        hf_revision: Optional[str] = None,
        hf_cache_dir: Optional[str] = None,
        hf_proxies: Optional[Dict] = None,
    ):
        """
        Initialize LazyStateDict with a checkpoint path.

        Args:
            checkpoint (str): Path to the checkpoint file or directory.
            meta_module_class (Type[nn.Module], optional): Class of the meta module to instantiate.
            meta_module (nn.Module, optional): Pre-initialized meta module.
            cache_state_dict (bool): Whether to cache the state dict in memory.
            torch_dtype (torch.dtype, optional): The dtype to use for the tensors.
            device (str): The device to load the tensors onto.
            hf_revision (str, optional): The revision of the model to download from Hugging Face Hub.
            hf_cache_dir (str, optional): The cache directory for Hugging Face models.
            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

        # Instantiate meta module if class provided
        if self.meta_module_class is not None:
            with init_empty_weights():
                self.meta_module = self.meta_module_class.from_pretrained(
                    checkpoint,
                    torch_dtype=torch_dtype,
                    revision=hf_revision,
                    cache_dir=hf_cache_dir,
                    proxies=hf_proxies,
                )

        # Store original checkpoint path and resolve to local path
        self._checkpoint = checkpoint
        self._local_path = resolve_checkpoint_path(
            checkpoint,
            hf_revision=hf_revision,
            hf_cache_dir=hf_cache_dir,
            hf_proxies=hf_proxies,
        )

        # Detect checkpoint file type and set up indexing
        self._index, self._index_filename, self._checkpoint_files = (
            self._resolve_checkpoint_files(self._local_path)
        )

        # Set up based on checkpoint type
        if self._index is not None:
            # if meta_module is provided, remove the keys that are not in the meta_module
            if self.meta_module is not None:
                meta_module_state_dict = self.meta_module.state_dict()
                for key in tuple(self._index.keys()):
                    if key not in meta_module_state_dict:
                        self._index.pop(key)
            if cache_state_dict:
                self._state_dict_cache = {}
            else:
                self._state_dict_cache = None
        elif len(self._checkpoint_files) == 1 and self._checkpoint_files[0].endswith(
            SAFE_WEIGHTS_NAME
        ):
            # SafeTensors file: create index mapping all keys to this file
            with safe_open(
                self._checkpoint_files[0], framework="pt", device=device
            ) as f:
                self._index = {key: self._checkpoint_files[0] for key in f.keys()}
                if cache_state_dict:
                    self._state_dict_cache = {}
                else:
                    self._state_dict_cache = None
        elif len(self._checkpoint_files) == 1 and self._checkpoint_files[0].endswith(
            WEIGHTS_NAME
        ):
            # PyTorch .bin file: load entire state dict immediately
            log.info(f"Loading full state dict from {WEIGHTS_NAME}")
            self._state_dict_cache = torch.load(self._checkpoint_files[0])
            # if meta_module is provided, remove the keys that are not in the meta_module
            if self.meta_module is not None:
                meta_module_state_dict = self.meta_module.state_dict()
                for key in tuple(self._state_dict_cache.keys()):
                    if key not in meta_module_state_dict:
                        self._state_dict_cache.pop(key)
        else:
            # Unsupported checkpoint format
            raise ValueError(
                f"Cannot determine the type of checkpoint, please provide a checkpoint path to a file containing a whole state dict with file name {WEIGHTS_NAME} or {SAFE_WEIGHTS_NAME}, or the index of a sharded checkpoint ending with `.index.json`."
            )

        self._torch_dtype = parse_dtype(torch_dtype)
        self._device = device

    @property
    def checkpoint(self) -> str:
        return self._checkpoint

    @property
    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 deepcopy(self)

    def _resolve_checkpoint_files(self, checkpoint: str):
        """
        Detect and resolve checkpoint files based on the checkpoint path.

        Handles single files, directories with state dict files, and sharded checkpoints.

        Returns:
            Tuple of (index_dict, index_filename, checkpoint_files)
        """
        # Reference: https://huggingface.co/docs/accelerate/v0.17.1/en/usage_guides/big_modeling
        checkpoint_files = None
        index_filename = None
        if os.path.isfile(checkpoint):
            # Single file: check if it's an index or a state dict
            if str(checkpoint).endswith(".json"):
                index_filename = checkpoint
            else:
                checkpoint_files = [checkpoint]
        elif os.path.isdir(checkpoint):
            # check if the whole state dict is present
            potential_state_bin = [
                f for f in os.listdir(checkpoint) if f == WEIGHTS_NAME
            ]
            potential_state_safetensor = [
                f for f in os.listdir(checkpoint) if f == SAFE_WEIGHTS_NAME
            ]
            if len(potential_state_bin) == 1:
                checkpoint_files = [os.path.join(checkpoint, potential_state_bin[0])]
            elif len(potential_state_safetensor) == 1:
                checkpoint_files = [
                    os.path.join(checkpoint, potential_state_safetensor[0])
                ]
            else:
                # Check for sharded checkpoints
                potential_index = [
                    f for f in os.listdir(checkpoint) if f.endswith(".index.json")
                ]
                if len(potential_index) == 0:
                    raise ValueError(
                        f"{checkpoint} is not a folder containing a `.index.json` file or a {WEIGHTS_NAME} or a {SAFE_WEIGHTS_NAME} file"
                    )
                elif len(potential_index) == 1:
                    index_filename = os.path.join(checkpoint, potential_index[0])
                else:
                    raise ValueError(
                        f"{checkpoint} containing more than one `.index.json` file, delete the irrelevant ones."
                    )
        else:
            # Invalid checkpoint path
            raise ValueError(
                "`checkpoint` should be the path to a file containing a whole state dict, or the index of a sharded "
                f"checkpoint, or a folder containing a sharded checkpoint or the whole state dict, but got {checkpoint}."
            )

        # Load index file if present
        if index_filename is not None:
            checkpoint_folder = os.path.split(index_filename)[0]
            with open(index_filename) as f:
                index = json.loads(f.read())

            # Extract weight_map if present (standard format)
            if "weight_map" in index:
                index = index["weight_map"]
            # Get list of unique checkpoint files
            checkpoint_files = sorted(list(set(index.values())))
            checkpoint_files = [
                os.path.join(checkpoint_folder, f) for f in checkpoint_files
            ]
        else:
            index = None
        return index, index_filename, checkpoint_files

    def _load_tensor_from_checkpoint_file(
        self, checkpoint_file: str, key: str, update_cache: bool = True
    ) -> torch.Tensor:
        """
        Load a tensor from the checkpoint file.
        For safetensors, loads only the requested tensor.
        For PyTorch files, loads the entire state dict on first access.
        """
        if checkpoint_file.endswith(".safetensors"):
            with safe_open(checkpoint_file, framework="pt", device=self._device) as f:
                tensor = f.get_tensor(key)
                if self._torch_dtype is not None:
                    tensor = tensor.to(self._torch_dtype)
                if update_cache and self._state_dict_cache is not None:
                    self._state_dict_cache[key] = tensor
                return tensor
        else:
            # PyTorch .bin file: load entire state dict
            state_dict = torch.load(checkpoint_file, map_location=self._device)
            if update_cache:
                if self._state_dict_cache is not None:
                    self._state_dict_cache.update(state_dict)
                else:
                    log.warning(
                        f"Load full state dict from file {checkpoint_file}, but state dict cache is disabled."
                    )
            return state_dict[key]

    def __getitem__(self, key: str) -> torch.Tensor:
        if self._state_dict_cache is not None and key in self._state_dict_cache:
            return self._state_dict_cache[key]

        if self._index is None:
            if len(self._checkpoint_files) == 1 and os.path.isfile(
                self._checkpoint_files[0]
            ):
                checkpoint_file = self._checkpoint_files[0]
                tensor = self._load_tensor_from_checkpoint_file(
                    checkpoint_file, key, update_cache=True
                )
                return tensor
            else:
                if len(self._checkpoint_files) > 1:
                    raise RuntimeError(
                        "Get multiple checkpoint files, but index is not provided."
                    )
                if not os.path.isfile(self._checkpoint_files[0]):
                    raise FileNotFoundError(
                        f"Checkpoint file {self._checkpoint_files[0]} not found."
                    )
                raise RuntimeError("Unexpected error.")
        else:
            if key not in self._index:
                raise KeyError(f"Key {key} not found in index.")
            checkpoint_file = os.path.join(self._local_path, self._index[key])
            if not os.path.isfile(checkpoint_file):
                raise FileNotFoundError(f"Checkpoint file {checkpoint_file} not found.")
            tensor = self._load_tensor_from_checkpoint_file(
                checkpoint_file, key, update_cache=True
            )
            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.
        """
        assert key in list(
            self.keys()
        ), "KeyError: Cannot set a tensor for a key that does not exist in the LazyStateDict."
        if self._state_dict_cache is not None:
            self._state_dict_cache[key] = value
        else:
            log.warning("State dict cache is disabled, initializing the cache.")
            self._state_dict_cache = {key: value}

    def __contains__(self, key: str) -> bool:
        if self._state_dict_cache is not None and key in self._state_dict_cache:
            return True
        if self._index is not None and key in self._index:
            return True
        if len(self._checkpoint_files) == 1 and os.path.isfile(
            self._checkpoint_files[0]
        ):
            try:
                tensor = self._load_tensor_from_checkpoint_file(
                    self._checkpoint_files[0], key, update_cache=False
                )
                return tensor is not None
            except (KeyError, FileNotFoundError, RuntimeError, EOFError):
                return False
        return False

    def __len__(self) -> int:
        if self._index is not None:
            return len(self._index)
        if len(self._checkpoint_files) == 1 and os.path.isfile(
            self._checkpoint_files[0]
        ):
            checkpoint_file = self._checkpoint_files[0]
            if checkpoint_file.endswith(".safetensors"):
                with safe_open(checkpoint_file, framework="pt", device="cpu") as f:
                    return len(tuple(f.keys()))
            else:
                return len(
                    tuple(torch.load(checkpoint_file, map_location="cpu").keys())
                )
        raise RuntimeError(
            "Unexpected error: cannot determine the number of keys in the state dict."
        )

    def __iter__(self) -> Iterator[str]:
        if self._index is not None:
            return iter(self._index)
        elif self._state_dict_cache is not None:
            return iter(self._state_dict_cache)
        else:
            raise RuntimeError(
                "Unexpected error: cannot determine the keys in the state dict."
            )

    def keys(self) -> Iterator[str]:
        for key in self:
            yield key

    def values(self) -> Iterator[torch.Tensor]:
        for key in self:
            yield self[key]

    def items(self) -> Iterator[Tuple[str, torch.Tensor]]:
        for key in self:
            yield key, self[key]

    def __repr__(self) -> str:
        if self._index is not None:
            return f"{self.__class__.__name__}(keys={list(self.keys())})"
        else:
            return (
                f"{self.__class__.__name__}(checkpoint_files={self._checkpoint_files})"
            )

    def get_parameter(self, target: str) -> torch.Tensor:
        return self[target]

    def get_submodule(self, target: str) -> nn.Module:
        if self.meta_module is not None:
            module: nn.Module = deepcopy(self.meta_module.get_submodule(target))
            module.to_empty(device=self._device)
            state_dict = {}
            for name, _ in module.named_parameters():
                state_dict[name] = self[f"{target}.{name}"]
            module.load_state_dict(state_dict)
            return module
        else:
            raise RuntimeError(
                "Cannot get submodule because meta_module is not provided."
            )

    def load_state_dict(
        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.

        Args:
            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:
            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():
            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}'"
        )

dtype property

torch.dtype: The dtype of the module (assuming that all the module parameters have the same dtype).

__init__(checkpoint, meta_module_class=None, meta_module=None, cache_state_dict=False, torch_dtype=None, device='cpu', hf_revision=None, hf_cache_dir=None, hf_proxies=None)

Initialize LazyStateDict with a checkpoint path.

Parameters:

• checkpoint (str) –

  Path to the checkpoint file or directory.

• meta_module_class (Type[Module], default: None ) –

  Class of the meta module to instantiate.

• meta_module (Module, default: None ) –

  Pre-initialized meta module.

• cache_state_dict (bool, default: False ) –

  Whether to cache the state dict in memory.

• torch_dtype (dtype, default: None ) –

  The dtype to use for the tensors.

• device (str, default: 'cpu' ) –

  The device to load the tensors onto.

• hf_revision (str, default: None ) –

  The revision of the model to download from Hugging Face Hub.

• hf_cache_dir (str, default: None ) –

  The cache directory for Hugging Face models.

• hf_proxies (Dict, default: None ) –

  Proxies to use for downloading from Hugging Face Hub.

Source code in fusion_bench/utils/lazy_state_dict.py

def __init__(
    self,
    checkpoint: str,
    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",
    hf_revision: Optional[str] = None,
    hf_cache_dir: Optional[str] = None,
    hf_proxies: Optional[Dict] = None,
):
    """
    Initialize LazyStateDict with a checkpoint path.

    Args:
        checkpoint (str): Path to the checkpoint file or directory.
        meta_module_class (Type[nn.Module], optional): Class of the meta module to instantiate.
        meta_module (nn.Module, optional): Pre-initialized meta module.
        cache_state_dict (bool): Whether to cache the state dict in memory.
        torch_dtype (torch.dtype, optional): The dtype to use for the tensors.
        device (str): The device to load the tensors onto.
        hf_revision (str, optional): The revision of the model to download from Hugging Face Hub.
        hf_cache_dir (str, optional): The cache directory for Hugging Face models.
        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

    # Instantiate meta module if class provided
    if self.meta_module_class is not None:
        with init_empty_weights():
            self.meta_module = self.meta_module_class.from_pretrained(
                checkpoint,
                torch_dtype=torch_dtype,
                revision=hf_revision,
                cache_dir=hf_cache_dir,
                proxies=hf_proxies,
            )

    # Store original checkpoint path and resolve to local path
    self._checkpoint = checkpoint
    self._local_path = resolve_checkpoint_path(
        checkpoint,
        hf_revision=hf_revision,
        hf_cache_dir=hf_cache_dir,
        hf_proxies=hf_proxies,
    )

    # Detect checkpoint file type and set up indexing
    self._index, self._index_filename, self._checkpoint_files = (
        self._resolve_checkpoint_files(self._local_path)
    )

    # Set up based on checkpoint type
    if self._index is not None:
        # if meta_module is provided, remove the keys that are not in the meta_module
        if self.meta_module is not None:
            meta_module_state_dict = self.meta_module.state_dict()
            for key in tuple(self._index.keys()):
                if key not in meta_module_state_dict:
                    self._index.pop(key)
        if cache_state_dict:
            self._state_dict_cache = {}
        else:
            self._state_dict_cache = None
    elif len(self._checkpoint_files) == 1 and self._checkpoint_files[0].endswith(
        SAFE_WEIGHTS_NAME
    ):
        # SafeTensors file: create index mapping all keys to this file
        with safe_open(
            self._checkpoint_files[0], framework="pt", device=device
        ) as f:
            self._index = {key: self._checkpoint_files[0] for key in f.keys()}
            if cache_state_dict:
                self._state_dict_cache = {}
            else:
                self._state_dict_cache = None
    elif len(self._checkpoint_files) == 1 and self._checkpoint_files[0].endswith(
        WEIGHTS_NAME
    ):
        # PyTorch .bin file: load entire state dict immediately
        log.info(f"Loading full state dict from {WEIGHTS_NAME}")
        self._state_dict_cache = torch.load(self._checkpoint_files[0])
        # if meta_module is provided, remove the keys that are not in the meta_module
        if self.meta_module is not None:
            meta_module_state_dict = self.meta_module.state_dict()
            for key in tuple(self._state_dict_cache.keys()):
                if key not in meta_module_state_dict:
                    self._state_dict_cache.pop(key)
    else:
        # Unsupported checkpoint format
        raise ValueError(
            f"Cannot determine the type of checkpoint, please provide a checkpoint path to a file containing a whole state dict with file name {WEIGHTS_NAME} or {SAFE_WEIGHTS_NAME}, or the index of a sharded checkpoint ending with `.index.json`."
        )

    self._torch_dtype = parse_dtype(torch_dtype)
    self._device = device

__setitem__(key, value)

Set a tensor in the LazyStateDict. This will update the state dict cache if it is enabled.

Source code in fusion_bench/utils/lazy_state_dict.py

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.
    """
    assert key in list(
        self.keys()
    ), "KeyError: Cannot set a tensor for a key that does not exist in the LazyStateDict."
    if self._state_dict_cache is not None:
        self._state_dict_cache[key] = value
    else:
        log.warning("State dict cache is disabled, initializing the cache.")
        self._state_dict_cache = {key: value}

load_state_dict(state_dict, strict=True)

Load a state dict into this LazyStateDict. This method is only for compatibility with nn.Module and it overrides the cache of LazyStateDict.

Parameters:

• state_dict (Dict[str, Tensor]) –

  The state dict to load.

• strict (bool, default: True ) –

  Whether to enforce that all keys in the state dict are present in this LazyStateDict.

Source code in fusion_bench/utils/lazy_state_dict.py

def load_state_dict(
    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.

    Args:
        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:
        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():
        if key in self:  # Only set keys that exist in this LazyStateDict
            self[key] = value

    return _IncompatibleKeys(missing_keys, unexpected_keys)

state_dict(keep_vars=False)

Parameters:

• keep_vars (bool, default: False ) –

  Ignored, as LazyStateDict does not support keep_vars. Just for compatibility.

Source code in fusion_bench/utils/lazy_state_dict.py

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 deepcopy(self)