FSDP cannot consolidate optimizer state dict with flatten params is False

[ShenglongZ]

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

[min-xu-ai]

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

[syorami]

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]

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