Apply transpose rule to safetensor saver function

[wang2yn84]

The current safetensor saver blindly saves merged tensor with tranpose. The transposition should be applied based on the huggingface vs tunix model implementation. This PR applies the transposition based on transpose rule instead, makes the logic more robust. The existing test covers saving and loading the checkpoint back.

Reference

Colab Notebook

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.