Repository: pytorch
Issue: FSDP `use_orig_params` + full sharding results in missing parameters in gathered state dict
Labels: oncall: distributed,triaged,module: fsdp

### 🐛 Describe the bug

## Description
There appears to be a bug in the `FullyShardedDataParallel` (FSDP) wrapper in PyTorch when accessing the inner module's state dict with `use_orig_params=True` and `sharding_strategy=ShardingStrategy.FULL_SHARD`. The inner module's state dict is missing some parameters for its child modules, while the wrapped model's state dict contains the expected parameter values.

## Reproduction Script (world size = 2)
```python
import torch
import torch.distributed as tdist
from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
from torch.distributed.fsdp import FullyShardedDataParallel, ShardingStrategy


class WrapperModel(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.model = CounterExampleModel()
        self.fc1 = self.model.fc1
        self.fc2 = self.model.fc2


class CounterExampleModel(torch.nn.Module):
    """Small classification model.

    Args:
        num_features (int): number of input features (default: 1)
        num_classes (int): number of classes (default: 2)
    """

    def __init__(
        self,
        num_features: int = 1,
        num_classes: int = 2,
        num_hidden: int = 8,
        device: str = 'cpu',
        bias: bool = True,
    ) -> None:

        self.num_features = num_features
        self.num_classes = num_classes

        fc1 = torch.nn.Linear(num_features, num_hidden, device=device, bias=bias)
        fc2 = torch.nn.Linear(num_hidden, num_classes, device=device, bias=bias)

        net = torch.nn.Sequential(
            torch.nn.AdaptiveAvgPool2d(1),
            torch.nn.Flatten(),
            fc1,
            torch.nn.ReLU(),
            fc2,
            torch.nn.Softmax(dim=-1),
        )
        super().__init__()

        self.net = net
        self.fc1 = fc1
        self.fc2 = fc2


if __name__ == '__main__':
    tdist.init_process_group(backend='gloo')

    wrapped_torch_module = WrapperModel()

    fsdp_config = {
        'use_orig_params': True,
        'state_dict_type': 'full',
    }

    wrapped_torch_module = FullyShardedDataParallel(
        wrapped_torch_module,
        use_orig_params=True,
        sharding_strategy=ShardingStrategy.FULL_SHARD,
        device_id = torch.cuda.current_device(),
    )

    wrapped_state_dict = get_model_state_dict(
        wrapped_torch_module,
        submodules=None,
        options=StateDictOptions(full_state_dict=True),
    )

    inner_state_dict = get_model_state_dict(
        wrapped_torch_module.model,
        submodules=None,
        options=StateDictOptions(full_state_dict=True),
    )

    if tdist.get_rank() == 0:
        print(f'{wrapped_torch_module=}')
        print(f"{inner_state_dict['fc2.bias']=}")
        print(f"{wrapped_state_dict['fc2.bias']=}")
        assert len(inner_state_dict['fc2.bias']) != 0 # fails
        assert len(wrapped_state_dict['fc2.bias']) != 0 # passes
```

## Stdout
```
wrapped_torch_module=FullyShardedDataParallel(
  (_fsdp_wrapped_module): WrapperModel(
    (model): CounterExampleModel(
      (net): Sequential(
        (0): AdaptiveAvgPool2d(output_size=1)
        (1): Flatten(start_dim=1, end_dim=-1)
        (2): Linear(in_features=1, out_features=8, bias=True)
        (3): ReLU()
        (4): Linear(in_features=8, out_features=2, bias=True)
        (5): Softmax(dim=-1)
      )
      (fc1): Linear(in_features=1, out_features=8, bias=True)
      (fc2): Linear(in_features=8, out_features=2, bias=True)
    )
    (fc1): Linear(in_features=1, out_features=8, bias=True)
    (fc2): Linear(in_features=8, out_features=2, bias=True)
  )
)
inner_state_dict['fc2.bias']=tensor([], device='cuda:0')
wrapped_state_dict['fc2.bias']=tensor([0.2590, 0.1419], device='cuda:0')
```

## Stderr
```
Traceback (most recent call last):
   File "/mnt/workdisk/iamroot/composer/counterexample.py", line 96, in <module>
     assert len(inner_state_dict['fc2.bias']) != 0
AssertionError
```

The assert statement `assert len(inner_state_dict['fc2.bias']) != 0` fails, indicating that the `'fc2.bias'` tensor in the inner module's state dict is empty. However, the same check on the wrapped model's state dict passes, showing that the `'fc2.bias'` tensor has the expected non-zero length.

Accessing the inner module's state dict directly when using FSDP with `use_orig_params=True` does not provide the expected results. The inner module's state dict is empty, while the wrapped model's state dict contains the expected parameter values. 

### Versions

Collecting environment information...
PyTorch version: 2.3.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.31

Python version: 3.10.13 (main, Aug 25 2023, 13:20:03) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-91-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB

Nvidia driver version: 535.129.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Byte Order:                         Little Endian
Address sizes:                      48 bits physical, 48 bits virtual
CPU(s):                             64
On-line CPU(s) list:                0-63
Thread(s) per core:                 1
Core(s) per socket:                 32
Socket(s):                          2
NUMA node(s):                       8
Vendor ID:                          AuthenticAMD
CPU family:                         25
Model:                              1
Model name:                         AMD EPYC 7513 32-Core Processor
Stepping:                           1
Frequency boost:                    enabled
CPU MHz:                            2600.000
CPU max MHz:                        3681.6399
CPU min MHz:                        1500.0000
BogoMIPS:                           5199.55
Virtualization:                     AMD-V
L1d cache:                          2 MiB
L1i cache:                          2 MiB
L2 cache:                           32 MiB
L3 cache:                           256 MiB
NUMA node0 CPU(s):                  0-7
NUMA node1 CPU(s):                  8-15
NUMA node2 CPU(s):                  16-23
NUMA node3 CPU(s):                  24-31
NUMA node4 CPU(s):                  32-39
NUMA node5 CPU(s):                  40-47
NUMA node6 CPU(s):                  48-55
NUMA node7 CPU(s):                  56-63
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Mitigation; safe RET
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 invpcid_single hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold v_vmsave_vmload vgif v_spec_ctrl umip pku ospke vaes vpclmulqdq rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] onnx==1.16.0
[pip3] onnxruntime==1.17.3
[pip3] pytorch-ranger==0.1.1
[pip3] torch==2.3.0
[pip3] torch-optimizer==0.3.0
[pip3] torchmetrics==1.0.3
[pip3] torchvision==0.16.2
[pip3] triton==2.3.0
[conda] Could not collect

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @rohan-varma