`init_meta_context` fails `isinstance` call (fails xFormers)

[SeanNaren]

🐛 Bug

The below case fails for init_meta_context. This probably has to do something with the subclasses for MyModule being changed. This as a result fails when we use an xformers model factory within the context manager.

import torch

from pytorch_lightning.utilities.meta import init_meta_context


class BaseModule(torch.nn.Module):
    pass


class MyModule(BaseModule):
    pass


with init_meta_context():
    my_module = MyModule()
    assert isinstance(my_module, BaseModule)

cc @tchaton @blefaudeux

[SeanNaren]

After further inspection this will require a re-write (I think) of how the logic works. The issue boils down to this:

with init_meta_context():
    my_module = MyModule() This calls `__new__()` in the metaclass and returns an actual MyModule class
    assert isinstance(my_module, BaseModule) # this fails, as the actual MyModule does not inherit from the MetaClass BaseModule

I had a go at re-writing the code, but it's tough. The gist of the re-write is to monkey patch the __new__ function rather than introduce the MetaClass, cc @tchaton

[tchaton]

Hey @SeanNaren.

Yes, I made an attempt at re-writing, but it seems it is more challenging than I thought. The new version has been hackier and didn't seem to resolve your problem either. Furthermore, it would break if a user was to make a deepcopy over the module.

Best,
T.C