github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-05-03 17:37:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

using pytest parametrize

Browse Source
This commit is contained in:
jiaruifang
2022-03-08 12:03:35 +08:00
committed by Frank Lee
parent da6bfb1427
commit 4d07bffd77
3 changed files with 43 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
tests/test_zero_data_parallel/test_shard_param.py
View File

@@ -4,19 +4,21 @@
from copy import deepcopy
from functools import partial
import colossalai
from colossalai.zero.sharded_param.sharded_param import ShardedParamV2
import pytest
import torch
import torch.multiprocessing as mp
import colossalai
from colossalai.zero.sharded_param.sharded_param import ShardedParamV2
from colossalai.zero.shard_utils import TensorShardStrategy
from colossalai.zero.sharded_param import ShardedTensor, ShardedParam
from colossalai.utils import free_port
from colossalai.logging import get_dist_logger, disable_existing_loggers
from tests.test_zero_data_parallel.common import Net, CONFIG, allclose
def run_shard_tensor(rank, world_size, port):
def _run_shard_tensor(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
t = ShardedTensor(tensor=torch.randn(world_size * 2, 3))
assert list(t.origin_shape) == [world_size * 2, 3]
@@ -32,9 +34,9 @@ def run_shard_tensor(rank, world_size, port):
@pytest.mark.dist
def test_shard_tensor():
world_size = 2
run_func = partial(run_shard_tensor, world_size=world_size, port=free_port())
@pytest.mark.parametrize("world_size", [1, 2])
def test_shard_tensor(world_size):
run_func = partial(_run_shard_tensor, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
@@ -52,8 +54,8 @@ def _run_shard_param_v2(rank, world_size, port):
@pytest.mark.dist
def test_shard_param_v2():
world_size = 2
@pytest.mark.parametrize("world_size", [1, 2])
def test_shard_param_v2(world_size):
run_func = partial(_run_shard_param_v2, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
@@ -86,40 +88,40 @@ def _run_test_shard_param(rank, world_size, port):
@pytest.mark.dist
def test_shard_param():
world_size = 2
@pytest.mark.parametrize("world_size", [1, 2])
def test_shard_param(world_size):
run_func = partial(_run_test_shard_param, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
def run_init_shard_param(rank, world_size, port):
def _run_init_shard_param(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
param = torch.nn.Parameter(data=torch.rand(2, 3))
param = torch.nn.Parameter(data=torch.rand(world_size, 3))
sparam = ShardedParam(param, None, True)
payload = sparam.payload(torch.device('cuda'))
assert (list(payload.shape) == [3])
del sparam
param_shape = (2, 3)
param_shape = (world_size, 3)
sparam = ShardedParam(param_shape, process_group=None, is_sharded=True, device=torch.device('cpu'))
payload = sparam.payload(torch.device('cuda'))
assert (list(payload.shape) == [3])
param_shape = (2, 3)
param_shape = (world_size, 3)
sparam = ShardedParam(param_shape, process_group=None, is_sharded=False, device=torch.device('cpu'))
payload = sparam.payload(torch.device('cuda'))
assert (list(payload.shape) == [2, 3])
assert (list(payload.shape) == [world_size, 3])
@pytest.mark.dist
def test_init_shard_param():
world_size = 2
run_func = partial(run_init_shard_param, world_size=world_size, port=free_port())
@pytest.mark.parametrize("world_size", [1, 4])
def test_init_shard_param(world_size):
run_func = partial(_run_init_shard_param, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
if __name__ == '__main__':
test_shard_tensor()
test_shard_param()
test_shard_param_v2()
test_init_shard_param()
test_shard_tensor(2)
test_shard_param(2)
test_shard_param_v2(2)
test_init_shard_param(4)