[BUG]Higher peak gpu memory usage when using FP16 optimizer.

**Describe the bug**
When setting '--use-precision-aware-optimizer', the fp16 optimizer(master_param: fp32, exp_avg_sq: fp16, exp_avg: fp16) has a higher peak gpu memory usage when updating parameters than fp32 optimizer.

**fp16 optimizer gpu memory usage:**
peak gpu memory: 56.0GiB

**fp32 optimizer gpu memory usage:**
peak gpu memory: 49.0GiB

**To Reproduce**
fp16 optimizer:
--use-precision-aware-optimizer \
--exp-avg-dtype fp16 \
--exp-avg-sq-dtype fp16 \
fp32 optimizer:
--use-precision-aware-optimizer \

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

**Stack trace/logs**
We have observed that this is because when calling FusedAdam.step(), the fp16 exp_avg and exp_avg_sq are unscaled to float before calling adam_func. However, the fp32 optimizer does not require such a conversion. Why do we need to convert the momentum to float instead of directly using fp16 for updates, as mentioned in the DeepSeek-V3 paper?

`3965 Addr: b'2dfc5b0000_0, Size: 3.5GiB (3758096384 bytes) allocation, Total memory used after allocation: 56.0GiB (60116956817 bytes), stream 0, timestamp Thu Mar 13 2025 15:17:40 GMT+0800 (中国标准时间)
CUDACachingAllocator.cpp:0:c10::cuda::CUDACachingAllocator::Native::DeviceCachingAllocator::malloc(signed char, unsigned long, CUstream_st*)
:0:c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::malloc(void**, signed char, unsigned long, CUstream_st*)
:0:c10::cuda::CUDACachingAllocator::Native::NativeCachingAllocator::allocate(unsigned long)
:0:at::TensorBase at::detail::_empty_strided_generic >(c10::ArrayRef, c10::ArrayRef, c10::Allocator*, c10::DispatchKeySet, c10::ScalarType)
??:0:at::detail::empty_strided_generic(c10::ArrayRef, c10::ArrayRef, c10::Allocator*, c10::DispatchKeySet, c10::ScalarType)
??:0:at::detail::empty_strided_cuda(c10::ArrayRef, c10::ArrayRef, c10::ScalarType, std::optional)
??:0:at::detail::empty_strided_cuda(c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)
??:0:at::native::empty_strided_cuda(c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)
RegisterCUDA.cpp:0:at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA__empty_strided(c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)
RegisterCUDA.cpp:0:c10::impl::wrap_kernel_functor_unboxed_, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CUDA__empty_strided>, at::Tensor, c10::guts::typelist::typelist, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional > >, at::Tensor (c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)
??:0:at::_ops::empty_strided::redispatch(c10::DispatchKeySet, c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)
RegisterBackendSelect.cpp:0:c10::impl::wrap_kernel_functor_unboxed_, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional), &at::(anonymous namespace)::empty_strided>, at::Tensor, c10::guts::typelist::typelist, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional > >, at::Tensor (c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)
??:0:at::_ops::empty_strided::call(c10::ArrayRef, c10::ArrayRef, std::optional, std::optional, std::optional, std::optional)
:0:at::empty_strided(c10::ArrayRef, c10::ArrayRef, c10::TensorOptions)
??:0:at::native::_to_copy(at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
RegisterCompositeExplicitAutograd.cpp:0:c10::impl::wrap_kernel_functor_unboxed_, std::optional, std::optional, std::optional, bool, std::optional), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeExplicitAutograd___to_copy>, at::Tensor, c10::guts::typelist::typelist, std::optional, std::optional, std::optional, bool, std::optional > >, at::Tensor (at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
??:0:at::_ops::_to_copy::redispatch(c10::DispatchKeySet, at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
RegisterBackendSelect.cpp:0:c10::impl::wrap_kernel_functor_unboxed_, std::optional, std::optional, std::optional, bool, std::optional), &at::(anonymous namespace)::_to_copy>, at::Tensor, c10::guts::typelist::typelist, std::optional, std::optional, std::optional, bool, std::optional > >, at::Tensor (at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
??:0:at::_ops::_to_copy::redispatch(c10::DispatchKeySet, at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
VariableType_0.cpp:0:torch::autograd::VariableType::(anonymous namespace)::_to_copy(c10::DispatchKeySet, at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
VariableType_0.cpp:0:c10::impl::wrap_kernel_functor_unboxed_, std::optional, std::optional, std::optional, bool, std::optional), &torch::autograd::VariableType::(anonymous namespace)::_to_copy>, at::Tensor, c10::guts::typelist::typelist, std::optional, std::optional, std::optional, bool, std::optional > >, at::Tensor (c10::DispatchKeySet, at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
??:0:at::_ops::_to_copy::call(at::Tensor const&, std::optional, std::optional, std::optional, std::optional, bool, std::optional)
??:0:at::native::to(at::Tensor const&, c10::ScalarType, bool, bool, std::optional)
RegisterCompositeImplicitAutograd.cpp:0:c10::impl::wrap_kernel_functor_unboxed_), &at::(anonymous namespace)::(anonymous namespace)::wrapper_CompositeImplicitAutograd_dtype_to>, at::Tensor, c10::guts::typelist::typelist > >, at::Tensor (at::Tensor const&, c10::ScalarType, bool, bool, std::optional)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, c10::ScalarType, bool, bool, std::optional)
??:0:at::_ops::to_dtype::call(at::Tensor const&, c10::ScalarType, bool, bool, std::optional)
python_variable_methods.cpp:0:torch::autograd::THPVariable_to_type(_object*, c10::ScalarType, std::optional)
python_variable_methods.cpp:0:torch::autograd::THPVariable_float(_object*, _object*, _object*)
/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/optimizers/fused_adam.py:304:get_unscaled_state
/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/optimizers/fused_adam.py:556:step
/usr/local/lib/python3.10/dist-packages/torch/optim/optimizer.py:487:wrapper
/workspace/megatron-lm-bc/megatron/core/optimizer/optimizer.py:452:step_with_ready_grads
/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py:116:decorate_context
??:0:PyMethod_New
workspace/megatron-lm-bc/megatron/core/optimizer/distrib_optimizer.py:1897:step_with_ready_grads
/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py:116:decorate_context
/workspace/megatron-lm-bc/megatron/core/optimizer/optimizer.py:1026:step_with_ready_grads
/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py:116:decorate_context
/workspace/megatron-lm-bc/megatron/core/optimizer/optimizer.py:1067:step
`



**Environment (please complete the following information):**
- Megatron-LM: core_r0.10.0
- PyTorch version: 2.5.1
- CUDA version: 12.4
- transformer_engine version: 2.1.0+450146a

**Proposed fix**
If you have a proposal for how to fix the issue state it here or link to a PR.

**Additional context**