fsdp with gradient checkpointing

[LifuWang-66]

Hi Tianwei,

When I create an accelerate config file, it asked "Do you want to enable FSDP activation checkpointing?". I assume it meant that fsdp is compatible with gradient checkpointing. Why is it not supported in DMD2?

Thanks for your help!

[tianweiy]

this one is only a flag. we also need to modify the code to support it. Some discussions are in #23 (comment)

However, DMD training is a little complex with different optimizers / models updated at different frequencies and some models running in a no_grad mode. This would make implementing the correct logic for gradient accumulation a little challenging. Therefore we don't have the bandwidth to look into this yet.