From 393c31da614c8b676f61bdf53e9a66eada186637 Mon Sep 17 00:00:00 2001
From: flybird11111 <1829166702@qq.com>
Date: Mon, 11 Nov 2024 18:01:02 +0800
Subject: [PATCH] fix 6108

---
 colossalai/checkpoint_io/utils.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py
index b3917bd9d..786d2df86 100644
--- a/colossalai/checkpoint_io/utils.py
+++ b/colossalai/checkpoint_io/utils.py
@@ -111,11 +111,6 @@ def search_tp_partition_dim(current_shape: torch.Size, original_shape: torch.Siz
         if length > current_shape[dim]:
             partition_dim = dim
             break
-    if partition_dim is not None:
-        assert (
-            original_shape[partition_dim] == tp_size * current_shape[partition_dim]
-        ), f"The parameter isn't evenly distributed among tensor parallel group: \
-                shape before sharding {original_shape}, shape after sharding {current_shape}"
 
     return partition_dim