[Bug] state_dict returns wrong path when DMP is as a submodule

[JacoCheung]

Since torchrec register a state_hook. The path of is always prepended a long prefix:

for example, if I have a module structure like

A.B.C.ec = EmbeddingCollection() # A is not DMP wrapped, instead, ec is DMP wrapped.

# then I'll get a state dict key
A.B.C.ec.A.B.C.ec.embeddings.weight

I would like to keep A.B.C.ec.embeddings.weight only. I found the related source code.

For some reason, I have to use DMP as a submodule of the model. Is there any thing I can do?