search

facebookresearch / fairscale

PyTorch extensions for high performance and large scale training.
Other
3.17k stars 277 forks source link

FSDP cannot consolidate optimizer state dict with flatten params is False #1100

Open ShenglongZ opened 1 year ago

ShenglongZ commented 1 year ago

I'm now training a large model with 2.5B parameters with AdamW optimizer. Due to the known issue about FSDP and activation checkpointing, I'm using FSDP with flatten params = False. When saving the training checkpoint, the model has state_dict() and local_state_dict() two methods which distinguish saving full or sharded model states. Is it possible to save all full (not sharded ) optimizer states in a single file as well?

I saw the gather_full_optim_state_dict method but there is an assertion that requires flatten_parameters=True image

min-xu-ai commented 1 year ago

Yeah, we have not find people and time to work on this. Did you try pytorch's version of FSDP?

syorami commented 1 year ago

hi @min-xu-ai I've tried pytorch version's FSDP but it has problems with the usage of checkpointing and offload. It seems that the work is still in progress. Is there any clean implementation of saving optimizer states with flatten_parameters=False?

min-xu-ai commented 1 year ago

Unfortunately. There isn’t any yet. Feel free to hack it. 😃