[checkpoint]: unpad padded parameters when saving sharded MoE checkpoints

[LiRunGuo]

📌 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 (ran the pinned hooks manually instead: black 24.10.0 --line-length=120, isort 5.13.2 --profile black, autoflake 2.3.1, all clean)

🚨 Issue number

No existing issue; found while adding Qwen3-MoE support. Described below.

📝 What does this PR do?

booster.save_model(..., shard=True) with MoeHybridParallelPlugin writes padded parameters. With tensor parallelism, VocabParallelEmbedding1D pads the vocab to a multiple of make_vocab_size_divisible_by * tp_size, so the saved embedding has the padded vocab size and transformers cannot load the checkpoint:

size mismatch for weight: copying a param with shape torch.Size([1024, 8]) from checkpoint,
the shape in current model is torch.Size([1000, 8]).

Cause. MoECheckpointIO._model_sharder is a copy of HybridParallelCheckpointIO._model_sharder without the to_unpadded_tensor step. (The unsharded path goes through model.state_dict(), which already unpads.)

This affects every MoE model whose vocab is not a multiple of 64 * tp_size, e.g. Qwen3 / Qwen3-MoE (151936). The existing test uses Mixtral's default vocab (32000), which is never padded, and tp_size=1.

Fix. Unpad like HybridParallelCheckpointIO does.

Test. tests/test_moe/test_moe_checkpoint.py now uses a vocab of 1000 and runs with {pp 2, ep 2, tp 1} (as before) and {pp 2, ep 1, tp 2}.

Results (4x H200, torch 2.5.1+cu124)

test_moe_checkpoint.py	main	this PR
pp 2, ep 2, tp 1	passes	passes
pp 2, ep 1, tp 2, vocab 1000	size mismatch for weight: [1024, 8] vs [1000, 8]	passes

💥 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.