[Lite] Optimizer state has not been consolidated on this rank

[hiyyg]

When using Lite and trying to save the optimizer state_dict using 'ddp_sharded', I got the error below:

RuntimeError: Optimizer state has not been consolidated on this rank.                 Please call `consolidate_state_dict()` on all ranks beforehand if you meant to save the global state

Any guide on how to solve this issue?

cc @carmocca @justusschock @awaelchli @SeanNaren

[awaelchli]

Hi @hiyyg
You can work around this issue for now by calling optimizer.consolidate_state_dict(). Let me know if that works.

Not yet sure if we will be able to find an elegant solution that does not require the user to manually make this call. This is where Lightning can shine though, it can take care of this automatically :)

[hiyyg]

Thanks @awaelchli , it now works. However, after using this optimizer.consolidate_state_dict(), the used gpu memory is larger than ddp, is that normal?

[awaelchli]

Yes, that's normal. "ddp_sharded" means the optimizer state gets "split" across processes/devices. This means you save memory but if you actually want to save the state of the optimizer, you need to transfer the individual shards and assemble the whole state. This is what optimizer.consolidate_state_dict() does and the reason why you need to call it before saving the state. The full state occupies more memory than a single shard.