Function z3_leaf_module

deepspeed/utils/z3_leaf_module.py:15–23

Returns whether a module in `model` has been flagged as a 'leaf' module. See `set_z3_leaf_modules` for more details. Args: model (torch.nn.Module): The model to which the leaf module flag will be applied. Returns: bool: Whether the module has been flag

(model: torch.nn.Module)

def z3_leaf_module(model: torch.nn.Module) -> bool:
    """Returns whether a module in `model` has been flagged as a 'leaf' module.
        See `set_z3_leaf_modules` for more details.
        Args:
            model (torch.nn.Module): The model to which the leaf module flag will be applied.
        Returns:
            bool: Whether the module has been flagged as a 'leaf' module.
    """
    return hasattr(model, '_z3_leaf') and model._z3_leaf


def z3_leaf_parameter(model: torch.nn.Parameter) -> bool: