Repository: pytorch
Issue: [Bug]: some parameters' grad is None when using FSDP with torch2.1.0
Labels: oncall: distributed,triaged,module: fsdp

### 🐛 Describe the bug
To reproduce the problem: Training [InternLM](https://github.com/InternLM/InternLM/tree/develop) with [config fsdp=True](https://github.com/InternLM/InternLM/blob/develop/configs/7B_sft.py#L157):
```shell
srun -p llm -n8 --ntasks-per-node=8 --cpus-per-task=4 --gpus-per-task=1 python train.py --config ./configs/7B_sft.py
```

```python
zero1 parallel (dict):
    1. size: int
        * if size <= 0, the size of the zero process group is equal to the size of the dp process group,
            so parameters will be divided within the range of dp.
        * if size == 1, zero is not used, and all dp groups retain the full amount of model parameters.
        * if size > 1 and size <= dp world size, the world size of zero is a subset of dp world size.
        For smaller models, it is usually a better choice to split the parameters within nodes with a setting <= 8.
    2. fsdp: bool, enable/disable torch's fully sharded data parallel, defaults to False.
"""
parallel = dict(
    zero1=dict(size=-1, fsdp=True),
    tensor=1,
    pipeline=dict(size=1, interleaved_overlap=True),
    sequence_parallel=True,
)
```

The error message is shown as follows:
```shell
2023-10-19 14:58:44,720 ERROR train.py:318 in <module> -- Raise exception from SH-IDC1-10-140-1-139 with rank id: 0
Traceback (most recent call last):
  File "/mnt/petrelfs/huangting.p/workspace/InternLM/train.py", line 316, in <module>
    main(args)
  File "/mnt/petrelfs/huangting.p/workspace/InternLM/train.py", line 240, in main
    trainer_result = trainer.step()
  File "/mnt/petrelfs/huangting.p/workspace/InternLM/internlm/core/trainer.py", line 195, in step
    return self._engine.step()
  File "/mnt/petrelfs/huangting.p/workspace/InternLM/internlm/core/engine.py", line 118, in step
    success, grad_norm = self.optimizer.step()
  File "/mnt/petrelfs/share_data/llm_env/miniconda3-py39_4/envs/llm-torch2.1-flash2/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 68, in wrapper
    return wrapped(*args, **kwargs)
  File "/mnt/petrelfs/huangting.p/workspace/InternLM/internlm/solver/optimizer/fsdp_optimizer.py", line 118, in step
    norm_group = self._compute_norm_with_fsdp_flatten(group_idx)
  File "/mnt/petrelfs/huangting.p/workspace/InternLM/internlm/solver/optimizer/fsdp_optimizer.py", line 92, in _compute_norm_with_fsdp_flatten
    norm_group = compute_norm(gradients=gradients, parameters=params, last_stage=True)
  File "/mnt/petrelfs/huangting.p/workspace/InternLM/internlm/solver/optimizer/utils.py", line 270, in compute_norm
    tensor_parallel_grads.append(g.data.float())
AttributeError: 'NoneType' object has no attribute 'data'
```

### Versions
```shell
# torch2.1.0
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
```
![企业微信截图_c513d224-c444-4f69-ba6f-732f4e3b5574](https://github.com/pytorch/pytorch/assets/20810277/15d9d339-e1f4-4cde-803d-38c6d4f29905)


**Note that when using torch2.0.1, the error will not happen.**

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin