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

[Tensor] rename some APIs in TensorSpec and Polish view unittest (#1176)

Browse Source
This commit is contained in:
Jiarui Fang
2022-06-27 15:56:11 +08:00
committed by GitHub
parent dd0420909f
commit 0dd4e2bbfb
9 changed files with 26 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
tests/test_tensor/test_tensor.py
View File

@@ -63,13 +63,19 @@ def test_operand():
def _run_view(world_size):
t_ref = torch.randn(4, 5)
t = ColoTensor.from_torch_tensor(
t_ref, TensorSpec(distspec.shard(process_group=gpc.get_group(ParallelMode.DATA), dims=[0], num_partitions=[2])))
t_ref,
TensorSpec(distspec.shard(process_group=gpc.get_group(ParallelMode.DATA), dims=[0],
num_partitions=[world_size])))
assert t.size()[0] == 4 * world_size
assert t.size(1) == 5
assert t.size() == torch.Size([4 * world_size, 5])
t.view_base(4 * 5)
assert t.tensor_spec.dist_spec.placement.value == 's'
t = t.view(4 * 5 * world_size)
assert t.tensor_spec.dist_spec.placement.value == 'r'
assert t.shape == torch.Size([4 * 5 * world_size])
@@ -100,11 +106,10 @@ def run_dist_tests(rank, world_size, port):
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 2])
@rerun_if_address_is_in_use()
def _test_dist_cases(world_size):
def test_dist_cases(world_size):
run_func = partial(run_dist_tests, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
if __name__ == '__main__':
# _test_dist_init(4)
_test_dist_cases(2)
test_dist_cases(2)