From e60afb566df102d4c5c4ae7c2a4b0f795b5a8c52 Mon Sep 17 00:00:00 2001 From: RunguoLi 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 7813b9133dca..67a75f773909 100644 --- a/colossalai/checkpoint_io/moe_checkpoint.py +++ b/colossalai/checkpoint_io/moe_checkpoint.py @@ -37,6 +37,7 @@ from colossalai.interface import ModelWrapper, OptimizerWrapper from colossalai.tensor.d_tensor import is_customized_distributed_tensor from colossalai.tensor.moe_tensor.api import is_moe_tensor, is_moe_tensor_ep_fused +from colossalai.tensor.padded_tensor import is_padded_tensor, to_unpadded_tensor try: from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX @@ -114,6 +115,8 @@ def maybe_pin_tensor(name: str, tensor: torch.Tensor) -> torch.Tensor: dist.gather(param_.contiguous(), gather_list=gathered, dst=ep_root, group=ep_group) if ep_rank == 0: param_ = torch.cat(gathered, dim=0) + if is_padded_tensor(param_): + param_ = to_unpadded_tensor(param_) param_ = maybe_pin_tensor(name, param_) block, block_size = state_dict_sharder.append_param(prefix + name, param_) if block is not None: 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