Repository: huggingface/diffusers
File: src/diffusers/hooks/group_offloading.py (Python)

Code before:

    # Parameters and Buffers of the top-level module need to be offloaded/onloaded separately
    # when the forward pass of this module is called. This is because the top-level module is not
    # part of any group (as doing so would lead to no VRAM savings).
    parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
    buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)
    parameters = [param for _, param in parameters]
    buffers = [buffer for _, buffer in buffers]

    # Create a group for the unmatched submodules of the top-level module so that they are on the correct
    # device when the forward pass is called.
    unmatched_modules = [unmatched_module for _, unmatched_module in unmatched_modules]
    unmatched_group = ModuleGroup(
        modules=unmatched_modules,
        offload_device=offload_device,
        onload_device=onload_device,
        offload_to_disk_path=offload_to_disk_path,
        offload_leader=module,
        onload_leader=module,
        parameters=parameters,
        buffers=buffers,
        non_blocking=False,
        stream=None,
        record_stream=False,
        onload_self=True,
        _group_id=_GROUP_ID_UNMATCHED_GROUP,
    )
    if stream is None:
        _apply_group_offloading_hook(module, unmatched_group, None)
    else:
        _apply_lazy_group_offloading_hook(module, unmatched_group, None)


def _apply_group_offloading_leaf_level(
    module: torch.nn.Module,
    offload_device: torch.device,
    onload_device: torch.device,
    non_blocking: bool,
    stream: Union[torch.cuda.Stream, torch.Stream, None] = None,
    record_stream: Optional[bool] = False,
    low_cpu_mem_usage: bool = False,
    offload_to_disk_path: Optional[str] = None,
) -> None:
    r"""
    This function applies offloading to groups of leaf modules in a torch.nn.Module. This method has minimal memory
    requirements. However, it can be slower compared to other offloading methods due to the excessive number of device
    synchronizations. When using devices that support streams to overlap data transfer and computation, this method can
    reduce memory usage without any performance degradation.

    Args:
        module (`torch.nn.Module`):

Review comment (suggestion):
For unmatched group we could also include the module class name e.g. `f"{module.__class__.__name__}_unmatched_group"`

Code after:
    for i, group in enumerate(matched_module_groups):
        for group_module in group.modules:
            _apply_group_offloading_hook(group_module, group, None, config=config)

    # Parameters and Buffers of the top-level module need to be offloaded/onloaded separately
    # when the forward pass of this module is called. This is because the top-level module is not
    # part of any group (as doing so would lead to no VRAM savings).
    parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
    buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)
    parameters = [param for _, param in parameters]
    buffers = [buffer for _, buffer in buffers]

    # Create a group for the unmatched submodules of the top-level module so that they are on the correct
    # device when the forward pass is called.
    unmatched_modules = [unmatched_module for _, unmatched_module in unmatched_modules]
    unmatched_group = ModuleGroup(
        modules=unmatched_modules,
        offload_device=config.offload_device,
        onload_device=config.onload_device,
        offload_to_disk_path=config.offload_to_disk_path,
        offload_leader=module,
        onload_leader=module,
        parameters=parameters,
        buffers=buffers,
        non_blocking=False,
        stream=None,
        record_stream=False,
        onload_self=True,
        group_id=f"{module.__class__.__name__}_unmatched_group",
    )
    if config.stream is None:
        _apply_group_offloading_hook(module, unmatched_group, None, config=config)
    else:
        _apply_lazy_group_offloading_hook(module, unmatched_group, None, config=config)


def _apply_group_offloading_leaf_level(module: torch.nn.Module, config: GroupOffloadingConfig) -> None:
    r"""
    This function applies offloading to groups of leaf modules in a torch.nn.Module. This method has minimal memory
    requirements. However, it can be slower compared to other offloading methods due to the excessive number of device
    synchronizations. When using devices that support streams to overlap data transfer and computation, this method can
    reduce memory usage without any performance degradation.
    """
    # Create module groups for leaf modules and apply group offloading hooks
    modules_with_group_offloading = set()
    for name, submodule in module.named_modules():
        if not isinstance(submodule, _SUPPORTED_PYTORCH_LAYERS):
            continue
        group = ModuleGroup(
            modules=[submodule],
            offload_device=config.offload_device,