github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-24 12:09:17 +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
commit 6a811810ba
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 0 additions and 5 deletions
Show Stats 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