Fix ModuleDict wrapping

[xmfan]

Stacked PRs:

• Make MoE axis_name patcheable #275
• ->Fix ModuleDict wrapping #260

---

Fix ModuleDict wrapping

[fmassa]

This generally LGTM and I was also thinking about doing something like that!

I wonder if we could (or should?) simplify/generalize the implementation to keep the original subclass information around as well?

[fmassa]

I wonder if we would want to keep the whole original class structure around (maybe with a nn.Module subclass indicating that the class has been AutoParallelized).
Something like

cls = type(ref_submod)
new_inst = ref_submod.__new__(cls)
new_inst.__dict__ = ref_submod.__dict__.copy()
setattr(curr_mod, attr_name, new_inst)

or if we want a subclass

cls = type(ref_submod)
new_cls = type(f"AutoP[{cls.__name__}]", (cls,), ref_submod.__dict__.copy())
new_inst = new_cls.__new__(new_cls)
new_inst.__dict__ = ref_submod.__dict__.copy()
setattr(curr_mod, attr_name, new_inst)

(but we need to cache those new classes to avoid creating too many redundant classes maybe?)

[xmfan]

i can give it a try

[xmfan]

this works for nn.Module subclasses without duplicating memory because all params/buffers are stored under __dict__._modules, and the dict is shallow copied

[fmassa]

LGTM, thanks!