Fix ModuleDict wrapping

[xmfan]

Stacked PRs:

• ->Fix ModuleDict wrapping #260

---

Claude'd

Before this PR, we directly use _assign_attr which defaults to always use nn.Module: https://github.com/pytorch/pytorch/blob/064f80dfa0482f6bd365a2f7db2e9c2f9f3ea88c/torch/export/unflatten.py#L102

For Torchtitan, the optimizer has some flag accessing pattern using.the ModuleDict APIs, we had to patch them until now (https://github.com/pytorch/torchtitan/blob/d54a6d4c92e063189c91b433e3b3d85d79dfb657/torchtitan/components/optimizer.py#L354-L361). After this PR, we'll use nn.ModuleDict when applicable.

[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?)