Repository: pytorch
Issue: PyTorch's Distributed Checkpoint Cannot Save a Parameter of Size 1
Labels: oncall: distributed,triaged,module: distributed_checkpoint

### 🐛 Describe the bug

When a module has a parameter which is a tensor of size 1 and you try to save its FSDP with torch.distributed.checkpoint, you get the following exception:

```
NotImplementedError: Only single local shard is supported.
```

The following is the code that produces this error.

```
import os

import torch
import torch.distributed as dist
import torch.distributed.checkpoint
import torch.distributed.checkpoint.state_dict
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

        self.small_linear = nn.Linear(1,1)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        zero = torch.zeros(1, x.device)
        return output + self.small_linear.forward(zero)


if __name__ == '__main__':
    world_size = int(os.environ['WORLD_SIZE'])
    rank = int(os.environ['RANK'])
    local_rank = int(os.environ['LOCAL_RANK'])

    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    model = Net().to(rank)
    model = FSDP(model)
    optimizer = optim.Adam(model.parameters(), lr=1e-4)

    module_state_dict, optimizer_state_dict = torch.distributed.checkpoint.state_dict.get_state_dict(model, optimizer)
    state_dict = {
        "module": module_state_dict,
        "optimizer": optimizer_state_dict
    }
    torch.distributed.checkpoint.save(state_dict, checkpoint_id="temp/checkpoint")

    dist.destroy_process_group()
```

Run this with

```
torchrun --nnodes 1 --nproc_per_node 2 test.py
```

This is the exact error messages that I got.

```
[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/pong/workspace/test-dcp/test.py", line 52, in <module>
[rank1]:     module_state_dict, optimizer_state_dict = torch.distributed.checkpoint.state_dict.get_state_dict(model, optimizer)
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/torch/distributed/checkpoint/state_dict.py", line 1109, in get_state_dict
[rank1]:     model_state_dict = _get_model_state_dict(model, info)
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/torch/distributed/checkpoint/state_dict.py", line 520, in _get_model_state_dict
[rank1]:     if torch.is_tensor(p) and p.is_meta:
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/typing_extensions.py", line 2853, in wrapper
[rank1]:     return arg(*args, **kwargs)
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/torch/distributed/_shard/sharded_tensor/api.py", line 1140, in __torch_function__
[rank1]:     return dispatch(st_instance, func)
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/torch/distributed/_shard/sharded_tensor/api.py", line 1122, in dispatch
[rank1]:     return _SHARDED_OPS[func](types, args, kwargs, st._process_group)
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/torch/distributed/_shard/op_registry_utils.py", line 33, in wrapper
[rank1]:     return wrapped_func(types, args, kwargs, process_group)
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py", line 57, in st_is_meta
[rank1]:     return args[0].local_tensor().is_meta
[rank1]:   File "/home/pong/workspace/test-dcp/venv/lib/python3.10/site-packages/torch/distributed/_shard/sharded_tensor/api.py", line 1099, in local_tensor
[rank1]:     raise NotImplementedError("Only single local shard is supported.")
[rank1]: NotImplementedError: Only single local shard is supported.
```

The same issue occurs when you use other modules with parameters of size 1.

For example, if I change ```self.scalar``` to

```
        self.scalar = nn.Linear(1,1)
```

I still got the same error. However, once your parameter becomes larger. There are no errors any more. For examples, the following definitions of scalars did not resulted in exceptions being thrown.

```
        self.scalar = nn.Parameter(torch.Tensor([1.0, 2.0]))
```
```
        self.scalar = nn.Linear(100,100)
```

### Versions

Collecting environment information...
PyTorch version: 2.4.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.10.11 (main, Mar  7 2024, 16:52:06) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.153.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000

Nvidia driver version: 552.74
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
Address sizes:                      46 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             24
On-line CPU(s) list:                0-23
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Silver 4310 CPU @ 2.10GHz
CPU family:                         6
Model:                              106
Thread(s) per core:                 2
Core(s) per socket:                 12
Socket(s):                          1
Stepping:                           6
BogoMIPS:                           4190.15
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 tsc_reliable nonstop_tsc cpuid pni pclmulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx512vbmi umip avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq rdpid fsrm md_clear flush_l1d arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  Microsoft
Virtualization type:                full
L1d cache:                          576 KiB (12 instances)
L1i cache:                          384 KiB (12 instances)
L2 cache:                           15 MiB (12 instances)
L3 cache:                           18 MiB (1 instance)
Vulnerability Gather data sampling: Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
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; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==2.0.1
[pip3] torch==2.4.0+cu124
[pip3] triton==3.0.0

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @LucasLLC @MeetVadakkanchery @mhorowitz @pradeepfn