Compatibility issues between FusedAdam and DTensor

Describe the bug

Recent modifications to FusedAdam have made it incompatible with DTensor. Specifically, in the optimizer state initialization section, the optimizer state is now created according to the global shape of the DTensor instead of creating a DTensor optimizer state with the same shape as the parameters.

To maintain compatibility with DTensor, the state tensors should be initialized using zeros_like(param) or empty_like(param) instead of zeros(param.shape) or empty(param.shape).

Steps/Code to reproduce bug

Training FSDP2 or Megatron-FSDP with FusedAdam will reproduce the bug.

Expected behavior

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

Environment overview (please complete the following information)

• Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
• Method of Transformer Engine install: [pip install or from source]. Please specify exact commands you used to install.
• If method of install is [Docker], provide docker pull & docker run commands used

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• OS version
• PyTorch version
• Python version
• Transformer Engine version
• CUDA version
• CUDNN version

Device details

• GPU model

Additional context

Add any other context about the problem here.

Start at FusedAdam's optimizer state initialization and compare the state shapes with the parameter shapes when parameters are DTensors. Reproduce the issue through an FSDP2 or Megatron-FSDP training setup, then verify that the optimizer state remains DTensor-compatible and matches the parameter shape.