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

[autoparallel] handled illegal sharding strategy (#1728)

Browse Source 
* [autoparallel] handled illegal sharding strategy

* polish code
This commit is contained in:
Frank Lee
2022-10-19 12:53:06 +08:00
committed by GitHub
parent cbe9a4cb45
commit eee84908d4
36 changed files with 459 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/test_tensor/test_sharding_spec.py
View File

@@ -1,6 +1,7 @@
import torch
from colossalai.tensor.sharding_spec import _DimSpec, ShardingSpec
from colossalai.device.device_mesh import DeviceMesh
from colossalai.tensor.sharding_spec import ShardingSpec, _DimSpec
def test_sharding_spec():
@@ -11,7 +12,7 @@ def test_sharding_spec():
# [8, 9, 10,11],
# [12,13,14,15]]
device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
entire_shape = torch.Size((4, 8, 6))
entire_shape = torch.Size((16, 8, 6))
dim_partition_dict = {0: [0, 1]}
# DistSpec:
# shard_sequence: S01,R,R