Skip to content

HookRegistry.remove_hook drops the next hook's new_forward (breaks FirstBlockCache with group offloading) #15193

Open

Description

@theontho

Describe the bug

HookRegistry.remove_hook corrupts the hook chain when the removed hook is not the last one and the next hook defines new_forward.

When removing hook i, remove_hook assigns the removed hook's old_forward to self._fn_refs[i + 1].forward. A hook that implements new_forward keeps its new_forward in fn_ref.forward and reaches the rest of the chain through fn_ref.original_forward. Overwriting forward therefore drops that hook's new_forward: the hook is still registered, but it no longer runs.

Where we hit it

We hit this while running Wan 2.2 with FirstBlockCache and lazy group offloading (enable_group_offload(..., use_stream=True) / layer tracking). Group offloading temporarily registers and removes its tracking hook around the blocks. After that removal, the FirstBlockCache tail hook's new_forward is bypassed. The cache's shared state then never gets its tail residuals, and a later step fails with:

TypeError: unsupported operand type(s) for +: 'Tensor' and 'NoneType'   # tail_block_residuals is None

It's intermittent in real pipelines because it depends on hook registration order, but it reproduces deterministically with two hooks.

Reproduction

import torch
from diffusers.hooks import HookRegistry, ModelHook

class Add(ModelHook):
    def pre_forward(self, module, *args, **kwargs):
        return tuple(a + 1 for a in args), kwargs

class PassThrough(ModelHook):
    calls = 0
    def new_forward(self, module, *args, **kwargs):
        PassThrough.calls += 1
        return self.fn_ref.original_forward(*args, **kwargs)

model = torch.nn.Linear(4, 4)
registry = HookRegistry.check_if_exists_or_initialize(model)
registry.register_hook(Add(), "add")
registry.register_hook(PassThrough(), "pass")
registry.remove_hook("add")
model(torch.ones(1, 4))
print(PassThrough.calls)  # expected 1, got 0: the remaining hook no longer runs

Expected behavior

After removing a hook, every remaining hook still runs, and the removed hook no longer runs.

System info

diffusers 0.42.x and current main; reproduces on CPU.

A fix with regression tests is in #15143.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions