[BUG] In zero3 mode, how to set nn.Linear weight (some parameters can be updated, but some cannot).

Describe the bug
For zero3, all of the parameters are partitioned, so the weight shape of nn.Linear is 0, and grad cannot be set for some dimensions of weight.

To Reproduce

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-VL")
model.lm_head.weight.data[-1, :].requires_grad_(True)

RUN torchrun --nproc_per_node 8 code.py --deepspeed_config ds_config_zero3.json

Expected behavior
Code runs successfully in the mode of deepspeed zero3.

### Tasks

### Tasks

Start with the Qwen/Qwen-VL reproduction in code.py and ds_config_zero3.json, running it with torchrun on eight processes. Trace how zero3 partitions lm_head.weight and handles the attempted requires_grad_ update on a selected row. Done means the reproduction runs successfully under DeepSpeed zero3 with the intended parameter update behavior.