FSDP checkpoints don't load when run is restarted with greater world size

[darkmirage]

A checkpoint is saved from an 8-GPU run with dp_shard set to 8 and all other parallelisms set to 1. My understanding is that this is configured as an FSDP run.

The checkpoint is resumed from 16 GPUs with dp_shard now set to 16. When loading the checkpoint, we get this error:

[rank0]: Traceback (most recent call last): (RANK 15)                                                                                            [rank0]:   File "/app/.venv/lib/python3.10/site-packages/torch/distributed/checkpoint/utils.py", line 164, in reduce_scatter                     [rank0]:     local_data = map_fun()                                                                                                              [rank0]:   File "/app/.venv/lib/python3.10/site-packages/torch/distributed/checkpoint/logger.py", line 83, in wrapper                            
[rank0]:     result = func(*args, **kwargs)                                                                                                      
[rank0]:   File "/app/.venv/lib/python3.10/site-packages/torch/distributed/checkpoint/state_dict_loader.py", line 211, in local_step             
[rank0]:     local_plan = planner.create_local_plan()                                                                                            
[rank0]:   File "/app/.venv/lib/python3.10/site-packages/torch/distributed/checkpoint/default_planner.py", line 233, in create_local_plan        
[rank0]:     return create_default_local_load_plan(                                                                                              
[rank0]:   File "/app/.venv/lib/python3.10/site-packages/torch/distributed/checkpoint/default_planner.py", line 354, in create_default_local_load
[rank0]:     raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")                                                                 
[rank0]: RuntimeError: Missing key in checkpoint state_dict: dataloader.dp_rank_15.  

My understanding is that torch distributed checkpoints are supposed to support dynamic resharding at load time. Does this not work with torchtitan?

I was able to successfully resume a checkpoint going down from 32 GPUs to 16.

[fegin]

ye, the dataloader and learning rate scheduler do not support resharding but model and optimizer states do support resharding, we may be able to add selective resharding support.

[tianyu-l]

we may be able to add selective resharding support.

before supporting this, we should error out in data loader if world size becomes smaller, otherwise it's silent error.

[mori360]

We add error messages at #816 when loading under low dp_degree at checkpoint saved with high dp_degree.

[rank0]: RuntimeError: Missing key in checkpoint state_dict: dataloader.dp_rank_15.

For the error message at summary here, which loads under high dp_degree at checkpoint saved with low dp_degree, we would support optional checkpoint loading in the next step.

[fegin]

We can close this issue after #819 is landed.