FSDP saving state improvement

[nastya236]

Currently in FSDP skips intermediate reconstruction of tree. This results in the following issues:

• grouping for the parameters might be non-deterministic -> loaded state might be incorrect
• state tight to the number of groups -- so to the size of the communication, which is wrong design

While current strategy works deterministically if everything is fixed from run to run, however improvement is needed.

[coder-jayp]

I would like to work on this.

[nastya236]

Go for it!