You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
HookRegistry.remove_hook drops the next hook's new_forward (breaks FirstBlockCache with group offloading) #15193
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.
Activity