fsdp: improve state saving determinism and tree reconstruction

[coder-jayp]

Addresses #3401.

Changes:

• Made parameter grouping deterministic by sorting keys returned from tree_flatten(). This fixes the non-deterministic grouping issue that could cause incorrect state loading.
• Added intermediate reconstruction of the original parameter tree before calling optimizer.apply_gradients.
  This ensures optimizer state is no longer tied to the number of communication groups or communication_size.

This change improves the reliability of saving and loading FSDP optimizer states while preserving the performance benefits of grouped reduce-scatter / all-gather operations.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[coder-jayp]

Hi @nastya236, it would be great to get some feedback. Thanks.

[nastya236]

Thanks for your contribution! Will look at it as soon as possible.

[coder-jayp]

Hi @nastya236,

Just a quick follow-up on this PR. Happy to make any changes if needed. Thanks.