mirror of https://github.com/hpcaitech/ColossalAI
60 lines
2.2 KiB
Python
60 lines
2.2 KiB
Python
import torch
|
|
from colossalai.auto_parallel.solver.op_handler.broadcast import is_broadcastable, get_broadcast_shape, recover_sharding_spec_for_broadcast_shape
|
|
from colossalai.tensor.sharding_spec import ShardingSpec
|
|
from colossalai.device.device_mesh import DeviceMesh
|
|
|
|
|
|
def test_is_broadcastable():
|
|
x1 = torch.rand(4, 4, 8)
|
|
x2 = torch.rand(1, 8)
|
|
assert is_broadcastable(x1.shape, x2.shape)
|
|
|
|
x1 = torch.rand(4, 2, 8)
|
|
x2 = torch.rand(2, 8)
|
|
assert is_broadcastable(x1.shape, x2.shape)
|
|
|
|
x1 = torch.rand(4, 2, 8)
|
|
x2 = torch.rand(4, 8)
|
|
assert not is_broadcastable(x1.shape, x2.shape)
|
|
|
|
|
|
def test_get_broadcast_shape():
|
|
x1 = torch.rand(4, 4, 8)
|
|
x2 = torch.rand(1, 8)
|
|
assert get_broadcast_shape(x1.shape, x2.shape) == [4, 4, 8]
|
|
|
|
x1 = torch.rand(4, 2, 8)
|
|
x2 = torch.rand(2, 8)
|
|
assert get_broadcast_shape(x1.shape, x2.shape) == [4, 2, 8]
|
|
|
|
x1 = torch.rand(4, 2, 8)
|
|
x2 = torch.rand(8)
|
|
assert get_broadcast_shape(x1.shape, x2.shape) == [4, 2, 8]
|
|
|
|
|
|
def test_recover_sharding_spec_for_broadcast_shape():
|
|
x1 = torch.rand(4, 1, 8)
|
|
x2 = torch.rand(2, 8)
|
|
|
|
physical_mesh_id = torch.arange(0, 4)
|
|
mesh_shape = (2, 2)
|
|
# [[0, 1]
|
|
# [2, 3]]
|
|
device_mesh = DeviceMesh(physical_mesh_id, mesh_shape)
|
|
|
|
broadcast_shape = get_broadcast_shape(x1.shape, x2.shape)
|
|
logical_sharding_spec_for_x1 = ShardingSpec(device_mesh=device_mesh,
|
|
dim_partition_dict={
|
|
0: [0],
|
|
1: [1]
|
|
},
|
|
entire_shape=broadcast_shape)
|
|
physical_sharding_spec_for_x1 = recover_sharding_spec_for_broadcast_shape(logical_sharding_spec_for_x1,
|
|
broadcast_shape, x1.shape)
|
|
print(physical_sharding_spec_for_x1)
|
|
|
|
assert physical_sharding_spec_for_x1.entire_shape == x1.shape
|
|
# dim 1 for the physical tensor is of broadcast type MULTIPLE, so should ignore
|
|
assert physical_sharding_spec_for_x1.dim_partition_dict == {0: [0]}
|
|
assert physical_sharding_spec_for_x1.sharding_sequence == ['S0', 'R', 'R']
|