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

Code before:
    def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
        # If there wasn't an onload_leader assigned, we assume that the submodule that first called its forward
        # method is the onload_leader of the group.
        if self.group.onload_leader is None:
            self.group.onload_leader = module

        # If the current module is the onload_leader of the group, we onload the group if it is supposed
        # to onload itself. In the case of using prefetching with streams, we onload the next group if
        # it is not supposed to onload itself.
        if self.group.onload_leader == module:
            if self.group.onload_self:
                self.group.onload_()

            should_onload_next_group = self.next_group is not None and not self.next_group.onload_self
            if should_onload_next_group:
                self.next_group.onload_()

            should_synchronize = (
                not self.group.onload_self and self.group.stream is not None and not should_onload_next_group
            )
            if should_synchronize:
                # If this group didn't onload itself, it means it was asynchronously onloaded by the
                # previous group. We need to synchronize the side stream to ensure parameters
                # are completely loaded to proceed with forward pass.
                # Also, we should only do this synchronize if we don't already do it from the sync call in
                # self.next_group.onload_, hence the `not should_onload_next_group` check.
                self.group.stream.synchronize()

        args = send_to_device(args, self.group.onload_device, non_blocking=self.group.non_blocking)
        kwargs = send_to_device(kwargs, self.group.onload_device, non_blocking=self.group.non_blocking)
        return args, kwargs

    def post_forward(self, module: torch.nn.Module, output):
        if self.group.offload_leader == module:
            self.group.offload_()
        return output


class LazyPrefetchGroupOffloadingHook(ModelHook):
    r"""
    A hook, used in conjunction with GroupOffloadingHook, that applies lazy prefetching to groups of torch.nn.Module.
    This hook is used to determine the order in which the layers are executed during the forward pass. Once the layer
    invocation order is known, assignments of the next_group attribute for prefetching can be made, which allows
    prefetching groups in the correct order.
    """

    _is_stateful = False

    def __init__(self):
        self.execution_order: List[Tuple[str, torch.nn.Module]] = []
        self._layer_execution_tracker_module_names = set()

Review comment (suggestion):
(nit):

```suggestion
                # Also, we should only do this synchronization if we don't already do it from the sync call in
                # self.next_group.onload_, hence the `not should_onload_next_group` check.
```

Code after:
    def pre_forward(self, module: torch.nn.Module, *args, **kwargs):
        # If there wasn't an onload_leader assigned, we assume that the submodule that first called its forward
        # method is the onload_leader of the group.
        if self.group.onload_leader is None:
            self.group.onload_leader = module

        # If the current module is the onload_leader of the group, we onload the group if it is supposed
        # to onload itself. In the case of using prefetching with streams, we onload the next group if
        # it is not supposed to onload itself.
        if self.group.onload_leader == module:
            if self.group.onload_self:
                self.group.onload_()

            should_onload_next_group = self.next_group is not None and not self.next_group.onload_self
            if should_onload_next_group:
                self.next_group.onload_()

            should_synchronize = (
                not self.group.onload_self and self.group.stream is not None and not should_onload_next_group
            )
            if should_synchronize:
                # If this group didn't onload itself, it means it was asynchronously onloaded by the
                # previous group. We need to synchronize the side stream to ensure parameters
                # are completely loaded to proceed with forward pass. Without this, uninitialized
                # weights will be used in the computation, leading to incorrect results
                # Also, we should only do this synchronization if we don't already do it from the sync call in
                # self.next_group.onload_, hence the `not should_onload_next_group` check.
                self.group.stream.synchronize()

        args = send_to_device(args, self.group.onload_device, non_blocking=self.group.non_blocking)
        kwargs = send_to_device(kwargs, self.group.onload_device, non_blocking=self.group.non_blocking)
        return args, kwargs

    def post_forward(self, module: torch.nn.Module, output):
        if self.group.offload_leader == module:
            self.group.offload_()
        return output


class LazyPrefetchGroupOffloadingHook(ModelHook):
    r"""
    A hook, used in conjunction with GroupOffloadingHook, that applies lazy prefetching to groups of torch.nn.Module.
    This hook is used to determine the order in which the layers are executed during the forward pass. Once the layer
    invocation order is known, assignments of the next_group attribute for prefetching can be made, which allows
    prefetching groups in the correct order.
    """

    _is_stateful = False

    def __init__(self):
        self.execution_order: List[Tuple[str, torch.nn.Module]] = []