(fsdp - maybe a bug) SHARDED_STATE_DICT returns tensor with no data
Labels: triaged, module: fsdp

### 🐛 Describe the bug

The huggingface `T5ForConditionalGeneration` model has the embedding layer named as `shared` at the top module, and passes a reference to it to the `T5Stack` ctor. In the `T5Stack` the passed embedding is stored as `embed_tokens`.
So `t5model.shared` and `t5model.encoder.embed_tokens` point to the same `nn.Embedding()` layer. Here's the two code references:

1. [`model.shared`](https://github.com/huggingface/transformers/blob/12298cb65c7e9d615b749dde935a0b4966f4ae49/src/transformers/models/t5/modeling_t5.py#L1541)
1. [`model.encoder.embed_tokens`](https://github.com/huggingface/transformers/blob/12298cb65c7e9d615b749dde935a0b4966f4ae49/src/transformers/models/t5/modeling_t5.py#L879)

and an excerpt of the code:

```
class T5ForConditionalGeneration(...): # nn.Module
    def __init__(...):
         # ... omitted...
          self.shared = nn.Embedding(...)
          self.encoder = T5Stack(encoder_config, self.shared)
         # ... omitted ...

class T5Stack(...): # nn.Module
     def __init__(..., embed_tokens):
         # ... omitted ...
         self.embed_tokens = embed_tokens
```

Wrapping this T5 model with FSDP and inspecting the details of embedding layer in a transformer (T5 model) model wrapped in FSDP with `ShardingStrategy.FULL_SHARD` and `auto_wrap_policy=transformer_auto_wrap_policy`, I'm observing that with `StateDictType.SHARDED_STATE_DICT`:

1.  `model.state_dict()["shared.weight"]`  is a `ShardedTensor` (as expected)
2. `model.state_dict()["encoder.embed_tokens.weight"]` is a Tensor with no data (see repro script's output below)

I'm guessing this has to do with this warning in the FSDP [docs page](
![image](https://github.com/pytorch/pytorch/assets/43595115/8c85bae5-fc7d-4a8a-b8fe-959c48f32560)
):
![image](https://github.com/pytorch/pytorch/assets/43595115/af228cbd-eb2e-4547-a23a-9ceb9a634ec8)

But I'm wrapping the `T5Block` (see my repro script below) not `T5Stack` in an FSDP unit so the outer module (`T5ForConditionalGeneration`) and `T5Block` belongs to the same (root) FSDP unit. 

Is this an expected behavior (I'm misunderstanding something) or is it a bug or should the documentation be updated?

Repro Script
--------------

To run:
```
$ pip install transformers, tabulate
$  torchrun --rdzv_backend c10d --rdzv_id 1 --nnodes 1 --nproc_per_node 2  repro_script.py sharded
```

Save this as `repro_script.py`

```python
import functools
import os
import sys

import torch
from tabulate import tabulate
import torch.distributed as dist
from torch.distributed._shard.sharded_tensor import ShardedTensor
from torch.distributed.elastic.multiprocessing.errors import record
from torch.distributed.fsdp import (
    FullyShardedDataParallel,
    ShardingStrategy,
    StateDictType,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers import T5ForConditionalGeneration, T5Config
from transformers.models.t5.modeling_t5 import T5Block


@record
def main():
    sdtype = sys.argv[1]
    state_dict_type = {
        "full": StateDictType.FULL_STATE_DICT,
        "local": StateDictType.LOCAL_STATE_DICT,
        "sharded": StateDictType.SHARDED_STATE_DICT,
    }[sdtype]

    rank = int(os.environ["RANK"])

    sharding_strategy = ShardingStrategy.FULL_SHARD

    t5_model = T5ForConditionalGeneration(
        T5Config(
            d_ff=512,
            d_kv=32,
            d_model=128,
            is_encoder_decoder=True,
            model_type="t5",
            n_positions=512,
            num_heads=2,
            num_layers=4,
            vocab_size=32128,
        )
    )

    fsdp_model = FullyShardedDataParallel(
        t5_model,
        auto_wrap_policy=functools.partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls={T5Block},
        ),
        sharding_strategy=sharding_strategy,
        device_id=device_id,
    )

    table = []
    layer_names = ["shared.weight", "encoder.embed_tokens.weight"]

    with FullyShardedDataParallel.state_dict_type(
        fsdp_model, state_dict_type=state_dict_type
    ):
        for layer_name in layer_names:
            state_dict = fsdp_model.state_dict()

            layer = state_dict.get(layer_name)
            tensor_type = type(layer)
            row = {
                "rank": rank,
                "sharding strategy": sharding_strategy.name,
                "state_dict_type": state_dict_type.name,
                "layer": layer_name,
            }

            if layer is None:
                continue

            row.update(
                {
                    "dtype": layer.dtype,
                    "shape": layer.shape,
                    "tensor type": tensor_type.__qualname__,
                }
            )

            if tensor_type != ShardedTensor:
                row["storage"] = layer.untyped_storage()

            table.append(row)

    if rank == 0:
        print(tabulate(table, headers="keys", stralign="left"))


if __name__ == "__main__":
    device_id = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(device_id)
    dist.init_process_group("nccl")

    try:
        main()
    finally:
        dist.destroy_process_group()
```

Output 
----------
```
  rank  sharding strategy    state_dict_type     layer                        dtype          shape                     tensor type    storage
------  -------------------  ------------------  ---------------------------  -------------  ------------------------  -------------  -------------------------------------------------------
     0  FULL_SHARD           SHARDED_STATE_DICT  shared.weight                torch.float32  torch.Size([32128, 128])  ShardedTensor
     0  FULL_SHARD           SHARDED_STATE_DICT  encoder.embed_tokens.weight  torch.float32  torch.Size([32128, 128])  Tensor         [torch.storage.UntypedStorage(device=cuda:0) of size 0]
```
### Versions

```
kiuk@ip-10-0-61-167% python collect_env.py                                                                                                                                                                                            ~/tmp
Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Amazon Linux 2 (x86_64)
GCC version: (GCC) 7.3.1 20180712 (Red Hat 7.3.1-15)
Clang version: 11.1.0 (Amazon Linux 2 11.1.0-1.amzn2.0.2)
CMake version: version 3.26.1
Libc version: glibc-2.26

Python version: 3.9.16 (main, Mar 31 2023, 16:44:31)  [GCC 7.3.1 20180712 (Red Hat 7.3.1-15)] (64-bit runtime)
Python platform: Linux-4.14.301-224.520.amzn2.x86_64-x86_64-with-glibc2.26
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 515.65.01
cuDNN version: Could not collect
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
CPU(s):              96
On-line CPU(s) list: 0-95
Thread(s) per core:  2
Core(s) per socket:  24
Socket(s):           2
NUMA node(s):        2
Vendor ID:           GenuineIntel
CPU family:          6
Model:               85
Model name:          Intel(R) Xeon(R) Platinum 8275CL CPU @ 3.00GHz
Stepping:            7
CPU MHz:             1187.921
BogoMIPS:            5999.99
Hypervisor vendor:   KVM
Virtualization type: full
L1d cache:           32K
L1i cache:           32K
L2 cache:            1024K
L3 cache:            36608K
NUMA node0 CPU(s):   0-23,48-71
NUMA node1 CPU(s):   24-47,72-95
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves ida arat pku ospke

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] mypy==0.991
[pip3] mypy-boto3-batch==1.26.103
[pip3] mypy-boto3-cloudwatch==1.26.127
[pip3] mypy-boto3-ec2==1.26.103
[pip3] mypy-boto3-iam==1.26.97
[pip3] mypy-boto3-s3==1.26.99
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.5
[pip3] pytorch-lightning==1.9.4
[pip3] torch==2.0.1
[pip3] torch-tb-profiler==0.4.1
[pip3] torchdata==0.6.1
[pip3] torchmetrics==0.11.4
[pip3] torchsnapshot-nightly==2023.3.15
[pip3] torchx-nightly==2023.5.9
[pip3] triton==2.0.0
[conda] No relevant packages
```

cc @zhaojuanmao @mrshenli @rohan-varma @awgu