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_checkpoint_io/test_gemini_checkpoint_io.py b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
index ba60de0bd1c1..3649b18ec8bb 100644
--- a/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
+++ b/tests/test_checkpoint_io/test_gemini_checkpoint_io.py
@@ -23,7 +23,7 @@
 
 ASYNC_MODES = [False, True] if HAS_TENSORNVME else [False]
 try:
-    import apex.amp
+    pass
 
     APEX_AVAILABLE = True
 except ImportError:
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
diff --git a/tests/test_zero/test_gemini/test_optim.py b/tests/test_zero/test_gemini/test_optim.py
index b89c591882d6..a0970e43db9e 100644
--- a/tests/test_zero/test_gemini/test_optim.py
+++ b/tests/test_zero/test_gemini/test_optim.py
@@ -15,7 +15,7 @@
 from tests.kit.model_zoo import model_zoo, run_fwd_bwd
 
 try:
-    import apex.amp
+    pass
 
     APEX_AMP_AVAILABLE = True
 except ImportError: