github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-01 01:06:00 +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

27
tests/test_zero_data_parallel/common.py
View File

@@ -2,11 +2,10 @@ from functools import partial
import torch
import torch.distributed as dist
from colossalai.logging import get_dist_logger
from colossalai.utils import checkpoint
from colossalai.zero.shard_utils import TensorShardStrategy
from colossalai.zero.sharded_model import ShardedModelV2
from colossalai.nn.optimizer import CPUAdam
LOGGER = get_dist_logger('zero_test')
@@ -16,20 +15,18 @@ _ZERO_MODEL_CONFIG = dict(reduce_scatter_bucket_size_mb=25,
fp32_reduce_scatter=False,
offload_config=None,
gradient_predivide_factor=1.0,
shard_param=True,
use_memory_tracer=False)
use_memory_tracer=False,
shard_strategy=TensorShardStrategy)
_ZERO_OPTIMIZER_CONFIG = dict(
optimizer_class=torch.optim.Adam, #CPUAdam
cpu_offload=False,
initial_scale=2**5,
min_scale=1,
growth_factor=2,
backoff_factor=0.5,
growth_interval=1000,
hysteresis=2,
max_scale=2**32,
lr=1e-3)
_ZERO_OPTIMIZER_CONFIG = dict(cpu_offload=False,
initial_scale=2**5,
min_scale=1,
growth_factor=2,
backoff_factor=0.5,
growth_interval=1000,
hysteresis=2,
max_scale=2**32,
lr=1e-3)
ZERO_PARALLEL_CONFIG = dict(fp16=dict(mode=None,),
zero=dict(

43
tests/test_zero_data_parallel/test_shard_model_v2.py
View File

@@ -1,15 +1,13 @@
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import copy
from asyncio.log import logger
from functools import partial
import colossalai
import pytest
import torch
import torch.multiprocessing as mp
from colossalai.logging import get_dist_logger
from colossalai.testing import parameterize
from colossalai.utils import free_port
from colossalai.zero.init_ctx import ZeroInitContext
from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
@@ -20,36 +18,30 @@ from tests.components_to_test.registry import non_distributed_component_funcs
from torch.nn.parallel import DistributedDataParallel as DDP
from common import CONFIG, check_grads_padding, run_fwd_bwd
from colossalai.testing import parameterize
@parameterize("enable_autocast", [True])
@parameterize("use_zero_init_ctx", [True])
@parameterize("shard_strategy", [TensorShardStrategy, BucketTensorShardStrategy])
def run_model_test(enable_autocast, use_zero_init_ctx, shard_strategy, logger):
@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
def run_model_test(enable_autocast, shard_strategy_class):
test_models = ['repeated_computed_layers', 'resnet18', 'bert']
shard_strategy = shard_strategy()
shard_strategy = shard_strategy_class()
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, _, _, criterion = get_components_func()
rm_torch_payload_on_the_fly = False
if use_zero_init_ctx:
with ZeroInitContext(convert_fp16=True,
target_device=torch.device(f'cpu:0'),
shard_strategy=shard_strategy,
shard_param=True,
rm_torch_payload_on_the_fly=rm_torch_payload_on_the_fly):
zero_model = model_builder(checkpoint=True)
zero_model = ShardedModelV2(zero_model, shard_strategy, use_memory_tracer=True)
with ZeroInitContext(convert_fp16=True,
target_device=torch.cuda.current_device(),
shard_strategy=shard_strategy,
shard_param=True,
rm_torch_payload_on_the_fly=rm_torch_payload_on_the_fly):
zero_model = model_builder(checkpoint=True)
zero_model = ShardedModelV2(zero_model, shard_strategy, use_memory_tracer=True)
model = model_builder(checkpoint=True).half()
col_model_deepcopy(zero_model, model)
model = model.cuda()
else:
model = model_builder(checkpoint=True).half().cuda()
zero_model = ShardedModelV2(copy.deepcopy(model), shard_strategy)
model = model_builder(checkpoint=True).half()
col_model_deepcopy(zero_model, model)
model = model.cuda()
model = DDP(model)
@@ -63,15 +55,10 @@ def run_model_test(enable_autocast, use_zero_init_ctx, shard_strategy, logger):
check_grads_padding(model, zero_model, loose=True)
# logger.debug('overall cuda ', zero_model._memstats_collector._overall_cuda)
# logger.debug('model cuda ', zero_model._memstats_collector._model_data_cuda)
def run_dist(rank, world_size, port):
colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
logger = get_dist_logger()
logger.set_level('DEBUG')
run_model_test(logger=logger)
run_model_test()
@pytest.mark.dist

50
tests/test_zero_data_parallel/test_sharded_optim_v2.py
View File

@@ -1,4 +1,3 @@
import copy
from functools import partial
import colossalai
@@ -6,15 +5,18 @@ import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from colossalai.nn.optimizer import CPUAdam
from colossalai.testing import parameterize
from colossalai.utils import free_port
from colossalai.zero.init_ctx import ZeroInitContext
from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
from colossalai.zero.sharded_model import ShardedModelV2
from colossalai.zero.sharded_model.utils import col_model_deepcopy
from colossalai.zero.sharded_optim import ShardedOptimizerV2
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 colossalai.nn.optimizer import CPUAdam
from colossalai.zero.sharded_optim._utils import has_inf_or_nan
from colossalai.testing import parameterize
from common import CONFIG, check_sharded_params_padding
@@ -38,36 +40,42 @@ def _run_step(model, optimizer, data, label, criterion, enable_autocast=False):
@parameterize("cpu_offload", [True, False])
@parameterize("use_cpuadam", [True, False])
@parameterize("shard_strategy", [TensorShardStrategy, BucketTensorShardStrategy])
def _run_test_sharded_optim_v2(cpu_offload, shard_strategy, use_cpuadam):
@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
def _run_test_sharded_optim_v2(cpu_offload, shard_strategy_class, use_cpuadam):
test_models = ['repeated_computed_layers', 'resnet18', 'bert']
shard_strategy = shard_strategy()
shard_strategy = shard_strategy_class()
if use_cpuadam and cpu_offload is False:
return
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model, train_dataloader, _, optimizer_class, criterion = get_components_func()
model = model(checkpoint=True).cuda()
zero_model = ShardedModelV2(copy.deepcopy(model),
model_builder, train_dataloader, _, optimizer_class, criterion = get_components_func()
with ZeroInitContext(convert_fp16=True,
target_device=torch.device(f'cpu:0'),
shard_strategy=shard_strategy,
shard_param=True,
rm_torch_payload_on_the_fly=False):
zero_model = model_builder(checkpoint=True)
zero_model = ShardedModelV2(zero_model,
shard_strategy,
offload_config=dict(device='cpu') if cpu_offload else None)
model = model_builder(checkpoint=True).half()
col_model_deepcopy(zero_model, model)
model = model.cuda().float()
if dist.get_world_size() > 1:
model = DDP(model)
lr = 1e-3
if use_cpuadam:
optim = torch.optim.Adam(model.parameters(), lr=lr)
sharded_optim = ShardedOptimizerV2(zero_model, CPUAdam, cpu_offload=cpu_offload, initial_scale=2**5, lr=lr)
else:
optim = optimizer_class(model.parameters(), lr=lr)
sharded_optim = ShardedOptimizerV2(zero_model,
optimizer_class,
cpu_offload=cpu_offload,
initial_scale=2**5,
lr=lr)
optimizer_class = CPUAdam
optim = optimizer_class(model.parameters(), lr=1e-3)
sharded_optim = optimizer_class(zero_model.parameters(), lr=1e-3)
sharded_optim = ShardedOptimizerV2(zero_model, sharded_optim, cpu_offload=cpu_offload, initial_scale=2**5)
for i, (data, label) in enumerate(train_dataloader):
#FIXME() if i > 5, the unittest will fail
# FIXME() if i > 5, the unittest will fail
if i > 3:
break
data, label = data.cuda(), label.cuda()

14
tests/test_zero_data_parallel/test_sharded_optim_with_sync_bn.py
View File

@@ -6,12 +6,12 @@ from functools import partial
import colossalai
import pytest
import torch
import torch.multiprocessing as mp
from colossalai.utils import free_port
from colossalai.core import global_context as gpc
from colossalai.context.parallel_mode import ParallelMode
from torchvision.models import resnet50
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 torchvision.models import resnet50
def run_dist(rank, world_size, port):
@@ -64,6 +64,10 @@ def run_dist(rank, world_size, port):
'expected the output from different ranks to be the same, but got different values'
# FIXME: enable this test in next PR
@pytest.mark.skip
@pytest.mark.dist
def test_sharded_optim_with_sync_bn():
"""

27
tests/test_zero_data_parallel/test_state_dict.py
View File

@@ -8,24 +8,37 @@ import colossalai
import pytest
import torch
import torch.multiprocessing as mp
from colossalai.testing import parameterize
from colossalai.utils import free_port
from colossalai.zero.init_ctx import ZeroInitContext
from colossalai.zero.shard_utils import (BucketTensorShardStrategy, TensorShardStrategy)
from colossalai.zero.sharded_model import ShardedModelV2
from colossalai.zero.sharded_model.utils import col_model_deepcopy
from tests.components_to_test.registry import non_distributed_component_funcs
from colossalai.testing import parameterize
from common import CONFIG
@parameterize("shard_strategy", [TensorShardStrategy, BucketTensorShardStrategy])
def run_zero_state_dict(shard_strategy):
@parameterize("shard_strategy_class", [TensorShardStrategy, BucketTensorShardStrategy])
def run_zero_state_dict(shard_strategy_class):
test_models = ['repeated_computed_layers', 'resnet18']
shard_strategy = shard_strategy()
shard_strategy = shard_strategy_class()
for model_name in test_models:
get_components_func = non_distributed_component_funcs.get_callable(model_name)
model_builder, train_dataloader, test_dataloader, optimizer, criterion = get_components_func()
model = model_builder()
model = model.half().cuda()
zero_model = ShardedModelV2(deepcopy(model), shard_strategy)
with ZeroInitContext(convert_fp16=True,
target_device=torch.cuda.current_device(),
shard_strategy=shard_strategy,
shard_param=True,
rm_torch_payload_on_the_fly=False):
zero_model = model_builder(checkpoint=True)
zero_model = ShardedModelV2(zero_model, shard_strategy)
model = model_builder(checkpoint=True).half()
col_model_deepcopy(zero_model, model)
model = model.cuda()
zero_state_dict = zero_model.state_dict()
for key, val in model.state_dict().items():
assert torch.equal(val, zero_state_dict[key])

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):