github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 19:13:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

[autoparallel] fixed wrong sharding strategy in conv handler (#1747)

Browse Source 
* [autoparallel] fixed wrong sharding strategy in conv handler

* polish code
This commit is contained in:
Frank Lee
2022-10-20 16:12:39 +08:00
committed by GitHub
parent 8b8937d901
commit 474111ecb5
6 changed files with 75 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/auto_parallel/tensor_shard/utils/misc.py
View File

@@ -68,4 +68,5 @@ def check_sharding_spec_validity(sharding_spec: ShardingSpec, tensor: torch.Tens
f'The dimension at index {i} has value {dim_size}, but it is sharded over {num_devices} devices.'
# make sure the entire shape matches the physical tensor shape
assert sharding_spec.entire_shape == tensor.shape
assert sharding_spec.entire_shape == tensor.shape, \
f'The entire_shape of the sharding spec {sharding_spec.entire_shape} does not match the tensor shape {tensor.shape}'