From e0d89a06e137a0829edddeaedbc5b773aa0bfc36 Mon Sep 17 00:00:00 2001
From: RunguoLi <li19107254665@gmail.com>
Date: Mon, 21 Sep 2026 23:14:25 -0500
Subject: [PATCH] [checkpoint] unpad padded parameters when saving sharded MoE
 checkpoints

`MoECheckpointIO._model_sharder` is a copy of
`HybridParallelCheckpointIO._model_sharder` without the
`to_unpadded_tensor` step, so `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 cannot be loaded by transformers:

    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]).

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

Unpad the parameters like HybridParallelCheckpointIO does, and run the MoE
checkpoint test with a padded vocab and tp_size=2 as well.
---
 colossalai/checkpoint_io/moe_checkpoint.py | 3 +++
 tests/test_moe/test_moe_checkpoint.py      | 9 +++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/colossalai/checkpoint_io/moe_checkpoint.py b/colossalai/checkpoint_io/moe_checkpoint.py
index 85e36f7c6336..fda3b3ce6953 100644
--- a/colossalai/checkpoint_io/moe_checkpoint.py
+++ b/colossalai/checkpoint_io/moe_checkpoint.py
@@ -34,6 +34,7 @@
 )
 from colossalai.interface import ModelWrapper, OptimizerWrapper
 from colossalai.tensor.moe_tensor.api import is_moe_tensor
+from colossalai.tensor.padded_tensor import is_padded_tensor, to_unpadded_tensor
 
 try:
     from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX
@@ -87,6 +88,8 @@ def _model_sharder(
                 continue
             # Gather tensor pieces when using tensor parallel.
             param_ = gather_distributed_param(param, keep_vars=False)
+            if is_padded_tensor(param_):
+                param_ = to_unpadded_tensor(param_)
             block, block_size = state_dict_sharder.append_param(prefix + name, param_)
             if block is not None:
                 yield block, block_size
diff --git a/tests/test_moe/test_moe_checkpoint.py b/tests/test_moe/test_moe_checkpoint.py
index f3f109192756..c408388d9a65 100644
--- a/tests/test_moe/test_moe_checkpoint.py
+++ b/tests/test_moe/test_moe_checkpoint.py
@@ -82,12 +82,15 @@ def check_optimizer_snapshot_equal(snapshot1, snapshot2, param2name, moe_dp_grou
                 num_attention_heads=2,
                 num_key_value_heads=2,
                 num_hidden_layers=2,
+                # not a multiple of make_vocab_size_divisible_by * tp_size, so the embedding is padded with tp
+                vocab_size=1000,
             ),
             MixtralForCausalLM,
         ],
     ],
 )
-def check_moe_checkpoint(test_config):
+@parameterize("plugin_config", [{"pp_size": 2, "ep_size": 2, "tp_size": 1}, {"pp_size": 2, "ep_size": 1, "tp_size": 2}])
+def check_moe_checkpoint(test_config, plugin_config):
     dtype, precision = torch.float16, "fp16"
     config, model_cls = test_config
     torch.cuda.set_device(dist.get_rank())
@@ -106,9 +109,7 @@ def check_moe_checkpoint(test_config):
         seed_all(10086)
         model = deepcopy(orig_model)
         optimizer = SGD(model.parameters(), lr=1e-3)
-        plugin = MoeHybridParallelPlugin(
-            pp_size=2, ep_size=2, tp_size=1, microbatch_size=1, zero_stage=1, precision=precision
-        )
+        plugin = MoeHybridParallelPlugin(**plugin_config, microbatch_size=1, zero_stage=1, precision=precision)
         booster = Booster(plugin=plugin)
         model, optimizer, *_ = booster.boost(model=model, optimizer=optimizer)
         # initialize grads