github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-10-25 01:40:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

update sharded optim and fix zero init ctx (#457)

Browse Source
This commit is contained in:
ver217
2022-03-18 15:44:47 +08:00
committed by GitHub
parent e2e9f82588
commit 642846d6f9
11 changed files with 162 additions and 162 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
tests/test_zero_data_parallel/test_zero_engine.py
View File

@@ -1,21 +1,24 @@
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import copy
from functools import partial
from colossalai.zero.sharded_model.sharded_model_v2 import ShardedModelV2
import pytest
import colossalai
from colossalai.utils import free_port
from colossalai.zero.sharded_optim._utils import has_inf_or_nan
import torch.multiprocessing as mp
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from colossalai.context.parallel_mode import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.utils import free_port
from colossalai.zero.init_ctx import ZeroInitContext
from colossalai.zero.sharded_model.utils import col_model_deepcopy
from colossalai.zero.sharded_optim._utils import has_inf_or_nan
from tests.components_to_test.registry import non_distributed_component_funcs
from torch.nn.parallel import DistributedDataParallel as DDP
from tests.components_to_test.registry import non_distributed_component_funcs
from common import check_sharded_params_padding, ZERO_PARALLEL_CONFIG, MP_PARALLEL_CONFIG, check_params
from common import (MP_PARALLEL_CONFIG, ZERO_PARALLEL_CONFIG, check_params,
check_sharded_params_padding)
def run_dist(rank, world_size, port, parallel_config):
@@ -30,10 +33,16 @@ def run_dist(rank, world_size, port, parallel_config):
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
with ZeroInitContext(convert_fp16=hasattr(gpc.config, 'fp16'),
target_device=torch.cuda.current_device(),
shard_strategy=gpc.config.zero.model_config.shared_strategy(
gpc.get_group(ParallelMode.DATA)),
shard_param=True):
colo_model = model_builder(checkpoint=True)
colo_model = model_builder(checkpoint=True)
torch_model = copy.deepcopy(colo_model).cuda()
torch_model.train()
torch_model = model_builder(checkpoint=True).half()
col_model_deepcopy(colo_model, torch_model)
torch_model = torch_model.cuda().float()
engine, train_dataloader, _, _ = colossalai.initialize(colo_model,
optimizer=optimizer_class,
criterion=criterion,
@@ -82,6 +91,10 @@ def run_dist(rank, world_size, port, parallel_config):
check_sharded_params_padding(torch_model, colo_model, loose=True)
# FIXME: enable this test in next PR
@pytest.mark.skip
@pytest.mark.dist
@pytest.mark.parametrize("world_size", [2, 4])
def test_mp_engine(world_size):
@@ -89,6 +102,7 @@ def test_mp_engine(world_size):
mp.spawn(run_func, nprocs=world_size)
@pytest.mark.skip
@pytest.mark.dist
@pytest.mark.parametrize("world_size", [1, 2])
def test_zero_engine(world_size):