在提交新问题之前，请确保您已经在社区中搜索过相关问题，并使用了社区中提供的资源/工具后，仍未找到满意的解决方式。

⚠️ 安全信息提醒：请仔细检查提供的文本内容，确保其不包含敏感数据信息，包括但不限于：

• API 令牌或密钥
• 密码或身份验证凭证
• 私有网址或接口地址
• 个人或机密数据
• ...

在分享配置信息或代码示例时，请将敏感信息脱敏处理，或使用 <TOKEN> 等占位符替代原有内容。

环境信息

例如：
- 操作系统
- 昇腾硬件信息
- CANN软件版本
- 安装的对应软件版本

🐛 问题描述

pretrain_gpt.py训练时，加入下列参数，预期可以支持：

unset CUDA_DEVICE_MAX_CONNECTIONS
MEGATRON_FSDP_ARGS="
    --use-megatron-fsdp \
    --fsdp-double-buffer \
    --ckpt-format fsdp_dtensor \
    --use-distributed-optimizer \
"

实际执行有如下报错：

[rank2]: Traceback (most recent call last):
[rank2]:   File "/home/xxx/workspace/Megatron-LM/pretrain_gpt.py", line 523, in <module>
[rank2]:     pretrain(
[rank2]:   File "/home/xxx/workspace/Megatron-LM/megatron/training/training.py", line 1422, in pretrain
[rank2]:     iteration, num_floating_point_operations_so_far = train(
[rank2]:                                                       ^^^^^^
[rank2]:   File "/home/xxx/workspace/Megatron-LM/megatron/training/training.py", line 3769, in train
[rank2]:     ) = train_step(
[rank2]:         ^^^^^^^^^^^
[rank2]:   File "/home/xxx/workspace/Megatron-LM/megatron/training/training.py", line 2428, in train_step
[rank2]:     update_successful, grad_norm, num_zeros_in_grad = optimizer.step()
[rank2]:                                                       ^^^^^^^^^^^^^^^^
[rank2]:   File "/home/miniconda3/envs/fanlu_py312/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/home/xxx/workspace/Megatron-LM/megatron/core/optimizer/optimizer.py", line 831, in step
[rank2]:     success = self.step_with_ready_grads()
[rank2]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/home/miniconda3/envs/fanlu_py312/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/home/xxx/workspace/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 3096, in step_with_ready_grads
[rank2]:     update_successful = super().step_with_ready_grads()
[rank2]:                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/home/miniconda3/envs/fanlu_py312/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/home/xxx/workspace/Megatron-LM/megatron/core/optimizer/optimizer.py", line 779, in step_with_ready_grads
[rank2]:     self.optimizer.step()
[rank2]:   File "/home/miniconda3/envs/fanlu_py312/lib/python3.12/site-packages/torch/optim/optimizer.py", line 485, in wrapper
[rank2]:     out = func(*args, **kwargs)
[rank2]:           ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/home/miniconda3/envs/fanlu_py312/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank2]:     return func(*args, **kwargs)
[rank2]:            ^^^^^^^^^^^^^^^^^^^^^
[rank2]:   File "/home/xxx/workspace/TransformerEngineNPU/transformer_engine/pytorch/optimizers/fused_adam.py", line 789, in step
[rank2]:     self._check_param_group_supported(group)
[rank2]:   File "/home/xxx/workspace/TransformerEngineNPU/transformer_engine/pytorch/optimizers/fused_adam.py", line 349, in _check_param_group_supported
[rank2]:     self._check_unsupported_param_type(param)
[rank2]:   File "/home/xxx/workspace/TransformerEngineNPU/transformer_engine/pytorch/optimizers/fused_adam.py", line 330, in _check_unsupported_param_type
[rank2]:     raise NotImplementedError(
[rank2]: NotImplementedError: DTensor parameters are not supported by NPU FusedAdam fallback yet.

欢迎加入社区，感谢您对社区的贡献 🎉!