[FSDP2/Megatron-FSDP/DCP] If model parameters are DTensors, optimizer states should also be DTensors. - #2795

[FSDP2/Megatron-FSDP/DCP] If model parameters are DTensors, optimizer states should also be DTensors.#2795

vthumbe1503 merged 17 commits into

NVIDIA:mainNVIDIA/TransformerEngine:mainfrom

cspades:cye/fused-adam-dcpcspades/TransformerEngine:cye/fused-adam-dcpCopy head branch name to clipboard