Megatron-FSDP checkpoint resumption functional test currently does not support FSDP DTensor checkpointing.
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
- Checkpointing test strategy `ckpt-resume` doesn't work for Megatron-FSDP, we have other DCP unit tests and local proof that checkpointing is fine, but we need to update the CI test to treat DCP preprocessing correctly.
**Steps/Code to reproduce bug**
Discovered via CI/CD in this PR: https://github.com/NVIDIA/Megatron-LM/pull/4027 cc @dimapihtar @cspades
```
TEST_TYPE: regular # Usually ckpt-resume, but as a WAR to #513 set to regular
```
which causes:
```
ValueError: Size mismatch between saved torch.Size([6561, 512]) and current: torch.Size([5533, 512]) for optimizer.state.module.module.module.embedding.word_embeddings.weight.exp_avg
2026-04-03T12:16:00.5223655Z 4-task-1-0/0 [default7]:[rank7]: Traceback (most recent call last): (RANK 7)
2026-04-03T12:16:00.5224278Z 4-task-1-0/0 [default7]:[rank7]: File "/usr/local/lib/python3.12/dist-packages/torch/distributed/checkpoint/utils.py", line 193, in reduce_scatter
2026-04-03T12:16:00.5224868Z 4-task-1-0/0 [default7]:[rank7]: local_data = map_fun()
2026-04-03T12:16:00.5225183Z 4-task-1-0/0 [default7]:[rank7]: ^^^^^^^^^
2026-04-03T12:16:00.5225724Z 4-task-1-0/0 [default7]:[rank7]: File "/usr/local/lib/python3.12/dist-packages/torch/distributed/checkpoint/logger.py", line 90, in wrapper
2026-04-03T12:16:00.5226309Z 4-task-1-0/0 [default7]:[rank7]: result = func(*args, **kwargs)
2026-04-03T12:16:00.5226852Z 4-task-1-0/0 [default7]:[rank7]: ^^^^^^^^^^^^^^^^^^^^^
2026-04-03T12:16:00.5227465Z 4-task-1-0/0 [default7]:[rank7]: File "/usr/local/lib/python3.12/dist-packages/torch/distributed/checkpoint/state_dict_loader.py", line 269, in local_step
2026-04-03T12:16:00.5228123Z 4-task-1-0/0 [default7]:[rank7]: local_plan = planner.create_local_plan()
2026-04-03T12:16:00.5228515Z 4-task-1-0/0 [default7]:[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
2026-04-03T12:16:00.5229163Z 4-task-1-0/0 [default7]:[rank7]: File "/usr/local/lib/python3.12/dist-packages/torch/distributed/checkpoint/default_planner.py", line 352, in create_local_plan
2026-04-03T12:16:00.5229840Z 4-task-1-0/0 [default7]:[rank7]: return create_default_local_load_plan(
2026-04-03T12:16:00.5230222Z 4-task-1-0/0 [default7]:[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2026-04-03T12:16:00.5230902Z 4-task-1-0/0 [default7]:[rank7]: File "/usr/local/lib/python3.12/dist-packages/torch/distributed/checkpoint/default_planner.py", line 485, in create_default_local_load_plan
2026-04-03T12:16:00.5231697Z 4-task-1-0/0 [default7]:[rank7]: raise ValueError(
2026-04-03T12:16:00.5232477Z 4-task-1-0/0 [default7]:[rank7]: ValueError: Size mismatch between saved torch.Size([6561, 512]) and current: torch.Size([5533, 512]) for optimizer.state.module.module.module.embedding.word_embeddings.weight.exp_avg
```
A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.
**Expected behavior**
- This test should pass without issues, but the checkpoint load fails.
**Additional context**
Add any other context about the problem here.
- Most likely due to not calling Megatron-FSDP uneven DTensor utilities.
Contributor guide
Assessment
This issue has not been assessed yet.