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

Merged

vthumbe1503 merged 17 commits intoNVIDIA:mainfrom

cspades:cye/fused-adam-dcp

Apr 4, 2026

Commits on Mar 31, 2026

If model parameters are DTensors, optimizer state should also be DTensor.
cspades
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
cspades
committed
Unpack DTensor in FusedAdam.step().
cspades
committed
Apply suggestions from code review

cspades
and
greptile-apps[bot]
committed
Revert erroneous Greptile diff.
cspades
committed
Add DTensor parity check to FusedAdam.step().
cspades
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
cspades
committed
Add DTensor handling in state_dict and load_state_dict, and add a DCP re-sharding test.
cspades
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
cspades
committed
Update test commentary.
cspades
committed
Filter out DCP resharding tests from the 2 GPU FusedAdam test matrix, as those tests need to be run in sequence.
cspades
committed
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed

Commits on Apr 3, 2026

Commits on Apr 4, 2026