[checkpointio] gather tensor before unpad it if the tensor is both padded and distributed

[Lemon-412]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs
• I have installed pre-commit: pip install pre-commit && pre-commit install

🚨 Issue number

fixed #6167

📝 What does this PR do?

To prevent #6167, gather the tensor before unpad it if the tensor is both padded and distributed.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

[Lemon-412]

request review from @flybird11111 @ver217 .

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

request review from @flybird11111 @ver217 .

[Lemon-412]

it seems like one unittest (test_dist_lamb) fails with OOM.
weird, should we rerun the test?

[ver217]

it seems like one unittest (test_dist_lamb) fails with OOM. weird, should we rerun the test?

This was fixed in main branch. Could you rebase the main branch?

[Lemon-412]

it seems like one unittest (test_dist_lamb) fails with OOM. weird, should we rerun the test?

This was fixed in main branch. Could you rebase the main branch?

Done. Approval is required again since we force push the code using rebase.