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

Merged

vthumbe1503 merged 17 commits into

NVIDIA: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
47d6d66
[pre-commit.ci] auto fixes from pre-commit.com hooks

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

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

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

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

Commits on Apr 3, 2026

Merge branch 'main' into cye/fused-adam-dcp
vthumbe1503
authored
d897a42
fix float8
vthumbe1503
committed
1f13911
xfail block scaling
vthumbe1503
committed
1c637ea

Commits on Apr 4, 2026