github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-12 20:54:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge 393c31da61 into 97f4bee9d8

Browse Source
This commit is contained in:
flybird11111
2025-06-19 18:02:15 +08:00
committed by GitHub
parent 97f4bee9d8 393c31da61
commit 6a811810ba
1 changed files with 0 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
colossalai/checkpoint_io/utils.py
View File

@@ -123,11 +123,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