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

[zero] yet an improved sharded param (#311)

Browse Source
This commit is contained in:
Jiarui Fang
2022-03-04 15:49:23 +08:00
committed by Frank Lee
parent c9e7d9582d
commit 90d3aef62c
3 changed files with 82 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
tests/test_zero_data_parallel/test_shard_param.py
View File

@@ -1,9 +1,11 @@
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
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
@@ -11,7 +13,7 @@ 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
from tests.test_zero_data_parallel.common import Net, CONFIG, allclose
def run_shard_tensor(rank, world_size, port):
@@ -36,28 +38,33 @@ def test_shard_tensor():
mp.spawn(run_func, nprocs=world_size)
def run_init_shard_param(rank, world_size, port):
def _run_shard_param_v2(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))
sparam = ShardedParam(param, None, True)
payload = sparam.payload(torch.device('cuda'))
assert (list(payload.shape) == [3])
del sparam
param_shape = (2, 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 = torch.nn.Parameter(torch.randn(2, 3))
param_ref = deepcopy(param)
sparam = ShardedParamV2(param=param, process_group=None)
param_shape = (2, 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])
allclose(sparam.data, param_ref.data)
assert (param.data.numel() == 1)
def run_shard_param_check(rank, world_size, port):
@pytest.mark.dist
def test_shard_param_v2():
world_size = 2
run_func = partial(_run_shard_param_v2, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
def _run_test_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(torch.randn(2, 3))
param_ref = deepcopy(param)
sparam = ShardedParamV2(param=param, process_group=None)
print(sparam.data)
print(param_ref.data)
logger = get_dist_logger()
model = Net()
@@ -77,12 +84,31 @@ def run_shard_param_check(rank, world_size, port):
@pytest.mark.dist
def test_shard_shape():
def test_shard_param():
world_size = 2
run_func = partial(run_shard_param_check, world_size=world_size, port=free_port())
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):
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))
sparam = ShardedParam(param, None, True)
payload = sparam.payload(torch.device('cuda'))
assert (list(payload.shape) == [3])
del sparam
param_shape = (2, 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)
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])
@pytest.mark.dist
def test_init_shard_param():
world_size = 2
@@ -92,5 +118,6 @@ def test_init_shard_param():
if __name__ == '__main__':
test_shard_tensor()
test_shard_shape()
test_shard_param()
test_shard_param_v2()
test_init_shard_param()