`init_meta` function return type mismatch

[nninept]

Incomplete Code

This is related with #13763 PR
and pytorch_lightning/utilities/meta.py

Problem

def init_meta(module_fn: Callable[..., Module], *args, **kwargs) -> Module:
       ...
       return module

class _MaterializerModule(subclass, metaclass=_IsinstanceMetaclass):
       ...
       def __new__(cls, *args, **kwargs):
              subclass = cls.__bases__[0]
              cls.add_subclasses(subclass)
              with cls.instantiation_context()
                      obj = init_meta(subclass, *args, **kwargs)

              obj.materialize = partial(cls.materialize, materialize_fn=obj.materialize)
              return obj

1. Return Type
   init_meta function returns torch.nn.Module type. But __new__ method assumes that obj is _MaterizerModule or at least its subclass.

2. Parameter Type Mismatch
   init_meta function takes Callable as first parameter. But in __new__ method, subclass variable is a class.

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Lite: enables pure PyTorch users to scale their existing code on any kind of device while retaining full control over their own loops and optimization logic.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, fine-tuning, and solving problems with deep learning.

• Bolts: Pretrained SOTA Deep Learning models, callbacks, and more for research and production with PyTorch Lightning and PyTorch.

• Lightning Transformers: Flexible interface for high-performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

cc @Borda @otaj

[tchaton]

@carmocca I think this code could be replaced with PyTorch built-in no ?

[carmocca]

I don't think the meta_init feature ever made it to PyTorch upstream (pytorch/pytorch#66317). At least I couldn't any of the code merged. cc @cbalioglu

I still think this is too unusable to have in the repository. There are some issues related to it too: #13660. Do we know of any users successfully using this?

I still lean on removing this until PyTorch adds support upstream.

[cbalioglu]

@carmocca actually this work is available for a while now in torchdistX (https://github.com/pytorch/torchdistx). In PyTorch v1.12 FSDP has officially introduced support for it (I believe there will be a blog post soon). We haven't upstreamed it to PyTorch though since having a separate package gives us the chance to iterate faster based on the user feedback.

Note also that the implementation in torchdistX is significantly more advanced than the (mostly hacky) version that was copy/pasted to Lightning. There are already several internal teams at Meta successfully using it.

[carmocca]

Thanks for the reply! Then I think the best option is to drop our copied version, suggest using torchdistx's features instead, and replace our tests to use them too to do some early testing on the integration with PyTorch Lightning

[cbalioglu]

Makes sense. Let me know if you hit any blockers, I would be happy to help.

[carmocca]

Another issue with our hacky copy is that it doesn't work with PyTorch 1.12

>       torch._C._enter_python_mode(cls)
E       AttributeError: module 'torch._C' has no attribute '_enter_python_mode'

Let me know if you hit any blockers

I have a PR open already in #13868. It's not a blocker but it would be nice if torchdistx added pytorch/torchdistx#49