Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Avoid overwrite of compiled module wrapper attributes #5549

Merged

Conversation

deepcharm
Copy link
Contributor

Fix overwriting of the compiled wrapper class attributes by those of the wrapped class itself: Copy only those attributes which are not already present in the wrapper.

In the current implementation of the CompiledModuleWrapper the wrapper attributes (eg forward method) are overwritten by self._dict_ = module._dict_.copy():

def CompiledModuleWrapper(mod, compile_config: Union[CompileConfig, None] = None):
     class wrapper(mod.__class__):
         def __init__(self, module, compile_config: Union[CompileConfig, None] = None):
             self.__dict__ = module.__dict__.copy()

This causes the wrapper's forward method not being called and, consequently, the wrapped module not compiled. Instead, the wrapped module forward method is being called as illustrated in the diagram below (a real scenario from Deespeed-Chat):

compiled_module_wrapper_bug

The proposed fix copies only those attributes which are not present in the wrapper class, thus implementing the desired inheritance quality of the wrapper.

Attached is a simple reproducer of the problem.
compiled_module_wrapper_bug.zip

Fix overwriting of the compiled wrapper class atributtes
by those of the wrapped class itself: Copy only those attributes
which are not already present in the wrapper.
@loadams loadams requested a review from tohtana May 20, 2024 17:19
Copy link
Contributor

@tohtana tohtana left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the fix!

@tjruwase tjruwase added this pull request to the merge queue May 21, 2024
Merged via the queue into microsoft:master with commit 5b314f4 May 21, 2024
12 checks passed
@deepcharm deepcharm deleted the compiled-wrapper-attrs-overwrite-fix branch May 22, 2024 07:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

4 participants