[FSDP] Cannot enable `offload_to_cpu=True` when using `LOCAL_STATE_DICT`

HAOCHENYE: <p>When I configure <code>state_dict_type</code> like this:</p>
<pre><code class="lang-python">        with fsdp_model.state_dict_type(
            fsdp_model,
            StateDictType.LOCAL_STATE_DICT,
            LocalStateDictConfig(offload_to_cpu=True),
            LocalOptimStateDictConfig(offload_to_cpu=True)
        ):
</code></pre>
<p>and then call</p>
<pre><code class="lang-python">model_state_dict = fsdp_model.state_dict()
</code></pre>
<p>an error will be raised <a href="https://github.com/pytorch/pytorch/blob/63f2f9fb0d0b7560a18117987b89c286b8f6b994/torch/distributed/_shard/sharded_tensor/api.py#L945" rel="noopener nofollow ugc">here</a>:</p>
<pre><code class="lang-bash">Local shards' tensor requires_grad property is incompatible with tensor property on rank 1: tensor property requires_grad=True, local shard tensor requires_grad=False. 
</code></pre>
<p>After debugging, I found that the <code>ShardedTensor</code> will create a CPU tensor with <code>requires_grad=False</code> <a href="https://github.com/pytorch/pytorch/blob/63f2f9fb0d0b7560a18117987b89c286b8f6b994/torch/distributed/_shard/sharded_tensor/api.py#LL490C5-L490C5" rel="noopener nofollow ugc">here</a> by default. However, the <code>requires_grads</code> is True in the sharded tensor of my model. This kind of inconsistency triggered the error.</p>
<p>I’m just curious about that is this case a bug or a feature? If it is a feathre, I just want to know why and how can I enable <code>cpu_offload=True</code> when using <code>LOCAL_STATE_DICT</code></p>
<p>The minimum code to reproduce the error</p>
<pre><code class="lang-python">import os.path as osp
import tempfile

import torch
import torch.distributed as dist
import torchvision
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import LocalStateDictConfig, StateDictType
from torch.distributed.fsdp.fully_sharded_data_parallel import \
    LocalOptimStateDictConfig
from torch.distributed.fsdp.wrap import _module_wrap_policy
from torch.optim import SGD


def conv_wrap_policy(*args, module_classes=torch.nn.Conv2d, **kwargs):
    return _module_wrap_policy(*args, module_classes=(module_classes, ), **kwargs)



def test_save_sharded_state_dict():
    # Save a checkpoint
    model = torchvision.models.resnet50()
    rank = dist.get_rank()
    fsdp_model = FSDP(model, use_orig_params=False, device_id=rank, auto_wrap_policy=conv_wrap_policy)
    # fsdp_model.get_state_dict_type(fsdp_model)
    optim = SGD(fsdp_model.parameters(), lr=0.001, momentum=0.9)
    loss = fsdp_model(torch.ones(3, 3, 224, 224).cuda().float())
    loss.backward()
    optim.step()
    with tempfile.TemporaryDirectory() as tmpdir:
        with fsdp_model.state_dict_type(
            fsdp_model,
            StateDictType.LOCAL_STATE_DICT,
            LocalStateDictConfig(offload_to_cpu=True),
            LocalOptimStateDictConfig(offload_to_cpu=True)
        ):
            # optim_state_dict = fsdp_model.full_optim_state_dict(fsdp_model, optim, rank0_only=True)
            optim_state_dict = fsdp_model.optim_state_dict(fsdp_model, optim)
            model_state_dict = fsdp_model.state_dict()
            torch.save(dict(model=model_state_dict, optim=optim_state_dict), osp.join(tmpdir, f'ckpt_{rank}.pth'))

            for p in model_state_dict.values():
                assert p.device == torch.device('cpu')
            ckpt = torch.load(f'ckpt_{rank}.pth')
            # check state dict is the original format
            fsdp_model.load_state_dict(ckpt['model'])
            optim.load_state_dict(ckpt['optim'])



if __name__ == '__main__':
    dist.init_process_group(backend='nccl')
    torch.cuda.set_device(dist.get_rank())
    # test_save_full_state_dict()
    test_save_sharded_state_dict()
</code></pre>