github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-18 19:58:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

test pretrain loading on multi-process (#922)

Browse Source
This commit is contained in:
Jiarui Fang 2022-05-09 17:07:35 +08:00 committed by GitHub
parent c195d2814c
commit 534afb018a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 41 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
tests/test_tensor/test_model.py
View File

@ -278,26 +278,6 @@ def test_colo_optimizer():
if i > 5:
break
def _test_pretrained():
from _utils import check_equal
from transformers import BertForMaskedLM
set_seed(1)
model_pretrained = BertForMaskedLM.from_pretrained('bert-base-uncased')
with ColoInitContext(lazy_memory_allocate=False, device=get_current_device()):
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model_pretrained = model_pretrained.cuda()
model = model.cuda()
dict_pretrained = {}
dict_col = {}
for name, param in model_pretrained.named_parameters():
dict_pretrained[name] = param
for name, param in model.named_parameters():
dict_col[name] = param
for name, param in dict_pretrained.items():
check_equal(param, dict_col[name])
def run_1d_row_tp(model_name: str):
# A simple net with two stacked nn.Linear
@ -376,7 +356,29 @@ def run_1d_row_tp(model_name: str):
break
def run_dist(rank, world_size, port):
def _run_pretrain_load():
from _utils import check_equal
from transformers import BertForMaskedLM
set_seed(1)
model_pretrained = BertForMaskedLM.from_pretrained('bert-base-uncased')
with ColoInitContext(lazy_memory_allocate=False, device=get_current_device()):
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model_pretrained = model_pretrained.cuda()
model = model.cuda()
dict_pretrained = {}
dict_col = {}
for name, param in model_pretrained.named_parameters():
dict_pretrained[name] = param
for name, param in model.named_parameters():
dict_col[name] = param
for name, param in dict_pretrained.items():
check_equal(param, dict_col[name])
def run_model_dist(rank, world_size, port):
config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
for name in ['simple_net']:
@ -390,7 +392,23 @@ def run_dist(rank, world_size, port):
#@parameterize('world_size', [1, 4])
@rerun_if_address_is_in_use()
def test_model(world_size):
run_func = partial(run_dist, world_size=world_size, port=free_port())
run_func = partial(run_model_dist, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
def run_pretrain_load_dist(rank, world_size, port):
config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
_run_pretrain_load()
# The test case has to download huggingface pretrained models from the internet
# So we manually trigger the test.
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@rerun_if_address_is_in_use()
def _test_pretrain_load(world_size):
run_func = partial(run_pretrain_load_dist, world_size=world_size, port=free_port())
mp.spawn(run_func, nprocs=world_size)
@ -398,4 +416,4 @@ if __name__ == '__main__':
# test_model_parameters()
# test_colo_optimizer()
# test_model()
_test_pretrained()
_test_pretrain_load(4)