[BUG]Cannot Save optimizer param in RL training checkpoints

**Describe the bug**

A clear and concise description of what the bug is.

```
Traceback (most recent call last):
File "", line 198, in _run_module_as_main
File "", line 88, in _run_code
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/recipe/r1_ascend/main_ppo.py", line 149, in
main()
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/main.py", line 94, in decorated_main
_run_hydra(
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
_run_app(
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/_internal/utils.py", line 457, in _run_app
run_and_report(
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
raise ex
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
return func()
^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/_internal/utils.py", line 458, in
lambda: hydra.run(
^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/_internal/hydra.py", line 132, in run
_ = ret.return_value
^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/core/utils.py", line 260, in return_value
raise self._return_value
File "/usr/local/python3.11.13/lib/python3.11/site-packages/hydra/core/utils.py", line 186, in run_job
ret.return_value = task_function(task_cfg)
^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/recipe/r1_ascend/main_ppo.py", line 42, in main
run_ppo(config)
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/recipe/r1_ascend/main_ppo.py", line 86, in run_ppo
ray.get(runner.run.remote(config))
File "/usr/local/python3.11.13/lib/python3.11/site-packages/ray/_private/auto_init_hook.py", line 22, in auto_init_wrapper
return fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/ray/_private/client_mode_hook.py", line 104, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/ray/_private/worker.py", line 2882, in get
values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/python3.11.13/lib/python3.11/site-packages/ray/_private/worker.py", line 968, in get_objects
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AttributeError): ray::TaskRunner.run() (pid=201336, ip=172.16.6.170, actor_id=5dd86c54372b6a0f29d6671501000000, repr=)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/trainer/main_ppo.py", line 317, in run
trainer.fit()
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/trainer/ppo/ray_trainer.py", line 1176, in fit
self._save_checkpoint()
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/trainer/ppo/ray_trainer.py", line 797, in _save_checkpoint
self.actor_rollout_wg.save_checkpoint(
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/single_controller/ray/base.py", line 48, in __call__
output = ray.get(output)
^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ray.exceptions.RayTaskError(AttributeError): ray::WorkerDict.actor_rollout_save_checkpoint() (pid=203506, ip=172.16.6.170, actor_id=39c355553c94c6d37ecbfd7701000000, repr=)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/single_controller/ray/base.py", line 700, in func
return getattr(self.worker_dict[key], name)(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/single_controller/base/decorator.py", line 433, in inner
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/workers/megatron_workers.py", line 771, in save_checkpoint
self.checkpoint_mananager.save_checkpoint(
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/utils/checkpoint/megatron_checkpoint_manager.py", line 384, in save_checkpoint
state_dict = self.generate_state_dict(
^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/verl/utils/checkpoint/megatron_checkpoint_manager.py", line 255, in generate_state_dict
optimizer_sharded_states = self.optimizer.sharded_state_dict(state_dict)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/megatron/core/optimizer/optimizer.py", line 1019, in sharded_state_dict
return self.chained_optimizers[0].sharded_state_dict(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/megatron/core/optimizer/distrib_optimizer.py", line 1111, in sharded_state_dict
param_state = self.sharded_param_state_fs_model_space(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/megatron/core/optimizer/distrib_optimizer.py", line 1363, in sharded_param_state_fs_model_space
tensors = _get_param_state_sharded_tensors(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/ma-user/modelarts/user-job-dir/verl_cann_qwen3_32B_gitcode_sam_1031/rl_train/qwen3/megatron/core/optimizer/distrib_optimizer.py", line 1325, in _get_param_state_sharded_tensors
dtype=state_ten.dtype,
^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'dtype'
```

this bug I see in code https://github.com/NVIDIA/Megatron-LM/tree/main/megatron/core/optimizer/distrib_optimizer.py

> for state_key, state_ten in tensors.items():
> if state_key == 'step':
> # Note that step is a 0-dim tensor, unlike other
> # states have the same size as the parameter.
> # The optimizer state of STEP is handled
> # specifically and is read from param_groups.
> continue
> replace_kwargs = dict(
> key=f'{prefix}.{state_key}.{sharded_metadata.key}',
> data=state_ten,
> dtype=state_ten.dtype,
> flattened_range=item_slice,
> replica_id=replica_id,
> )
> if isinstance(sharded_metadata, ShardedTensorFactory):
> replace_kwargs.pop('dtype')
> tensors[state_key] = replace(sharded_metadata, **replace_kwargs)
> tensors[state_key].validate_metadata_integrity()
> return tensors

I fix this bug，with this

> for state_key, state_ten in tensors.items():
> if state_key == 'step':
> # Note that step is a 0-dim tensor, unlike other
> # states have the same size as the parameter.
> # The optimizer state of STEP is handled
> # specifically and is read from param_groups.
> continue
> if state_ten is None:
> return []
> replace_kwargs = dict(
> key=f'{prefix}.{state_key}.{sharded_metadata.key}',
> data=state_ten,
> dtype=state_ten.dtype,
> flattened_range=item_slice,
> replica_id=replica_id,
> )
> if isinstance(sharded_metadata, ShardedTensorFactory):
> replace_kwargs.pop('dtype')
> tensors[state_key] = replace(sharded_metadata, **replace_kwargs)
> tensors[state_key].validate_metadata_integrity()
> return tensors

A clear and concise description of what the bug is.

**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

**Expected behavior**

A clear and concise description of what you expected to happen.

**Additional context**

Add any other context about the problem here.