[FSDP][optim_state_dict] Make FSDP optim_state_dict aware of DDP pref…

…ix (#96415)

Summary: When wrapping FSDP within DDP, optimizer state_dict may be broken due to the prefix of DDP. This PR fixes the issue.

Test Plan: CI

Differential Revision: D43893609

Pull Request resolved: pytorch/pytorch#96415
Approved by: https://github.com/zhaojuanmao

torch/distributed/fsdp/_common_utils.py

@@ -299,6 +299,15 @@ def f(module: torch.nn.Module, prefix: str, *args, **kwargs):
                             f"submodule_name = {submodule_name}"
                         )
                         new_prefix = prefix
+                    elif submodule_name == "module":
+                        warnings.warn(
+                            "An unexpected prefix is detected. This case "
+                            " should only happen when DDP wraps the outer "
+                            " modules while FSDP wraps the inner ones."
+                            f"prefix = {prefix}, "
+                            f"submodule_name = {submodule_name}"
+                        )
+                        new_prefix = prefix
             f(submodule, new_prefix, *args, **kwargs)
 
     f(root_module, "", *args, **kwargs)