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

[checkpoint]support generalized scheduler (#1222)

Browse Source
This commit is contained in:
Yi Zhao
2022-07-07 18:16:38 +08:00
committed by GitHub
parent a98319f023
commit 04537bf83e
4 changed files with 85 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
tests/test_utils/test_colo_checkpoint.py
View File

@@ -8,6 +8,8 @@ from functools import partial
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.optim.lr_scheduler import MultiplicativeLR
import colossalai
from colossalai.testing import rerun_if_address_is_in_use
@@ -102,10 +104,14 @@ def remove(path):
raise ValueError("file {} is not a file or dir.".format(path))
def run_checkpoint(init_spec_func, use_ddp, test_epoch, pg):
def run_checkpoint(init_spec_func, use_ddp, test_epoch, test_scheduler, pg):
num_epoch = 5
warmup_epoch = 2
batch = 3
feature = 32
category = 16
train_dataloader = DummyDataLoader(batch, category, feature, length=16)
with ColoInitContext(device=get_current_device()):
model = MLP(feature, category)
@@ -129,14 +135,25 @@ def run_checkpoint(init_spec_func, use_ddp, test_epoch, pg):
weight_decay=0)
optimizer_ref = torch.optim.Adam(model_ref.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, total_steps=20, warmup_steps=5)
lr_scheduler_reload = CosineAnnealingWarmupLR(optimizer=optimizer_reload, total_steps=20, warmup_steps=5)
lr_scheduler_ref = CosineAnnealingWarmupLR(optimizer=optimizer_ref, total_steps=20, warmup_steps=5)
if test_scheduler == 'colossalai_cosine_warmup':
lr_scheduler = CosineAnnealingWarmupLR(optimizer=optimizer, total_steps=num_epoch, warmup_steps=warmup_epoch)
lr_scheduler_reload = CosineAnnealingWarmupLR(optimizer=optimizer_reload,
total_steps=num_epoch,
warmup_steps=warmup_epoch)
elif test_scheduler == 'torch_cosine':
lr_scheduler = CosineAnnealingLR(optimizer=optimizer, T_max=num_epoch)
lr_scheduler_reload = CosineAnnealingLR(optimizer=optimizer_reload, T_max=num_epoch)
elif test_scheduler == 'torch_lambda':
lr_lambda = lambda epoch: 0.95
lr_scheduler = MultiplicativeLR(optimizer=optimizer, lr_lambda=lr_lambda)
lr_scheduler_reload = MultiplicativeLR(optimizer=optimizer_reload, lr_lambda=lr_lambda)
init_spec_func(model, pg)
init_spec_func(model_ref, pg)
for epoch in range(0, 20):
for epoch in range(0, num_epoch):
if epoch <= test_epoch:
for i, image_dict in enumerate(train_dataloader):
if use_ddp:
@@ -155,7 +172,6 @@ def run_checkpoint(init_spec_func, use_ddp, test_epoch, pg):
for ref_p, p in zip(model_ref.parameters(), model.parameters()):
ref_p.data.copy_(p)
optimizer_ref = copy.deepcopy(optimizer)
lr_scheduler_ref = copy.deepcopy(lr_scheduler)
check_param_equal(model, model_ref)
save_checkpoint('./checkpoint', epoch, model, optimizer, lr_scheduler)
@@ -189,28 +205,34 @@ def run_checkpoint(init_spec_func, use_ddp, test_epoch, pg):
check_param_equal(model_ref, model_reload)
def run_dist(rank, world_size, port, use_ddp, test_epoch):
def run_dist(rank, world_size, port, use_ddp, test_epoch, test_scheduler):
if use_ddp and world_size == 1:
return
tp_world_size = world_size // 2 if use_ddp else world_size
config = dict(parallel=dict(tensor=dict(mode="1d", size=tp_world_size),))
colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
pg = ProcessGroup(tp_degree=world_size)
run_checkpoint(init_1d_row_for_linear_weight_spec, use_ddp, test_epoch, pg)
run_checkpoint(init_1d_row_for_linear_weight_spec, use_ddp, test_epoch, test_scheduler, pg)
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [4])
@pytest.mark.parametrize('use_ddp', [True])
@pytest.mark.parametrize('test_epoch', [1, 2, 3])
@pytest.mark.parametrize('test_scheduler', ['colossalai_cosine_warmup', 'torch_cosine', 'torch_lambda'])
@rerun_if_address_is_in_use()
def test_checkpoint(world_size, use_ddp, test_epoch):
def test_checkpoint(world_size, use_ddp, test_epoch, test_scheduler):
if not os.path.isdir('./checkpoint'):
os.mkdir('./checkpoint')
run_func = partial(run_dist, world_size=world_size, port=free_port(), use_ddp=use_ddp, test_epoch=test_epoch)
run_func = partial(run_dist,
world_size=world_size,
port=free_port(),
use_ddp=use_ddp,
test_epoch=test_epoch,
test_scheduler=test_scheduler)
mp.spawn(run_func, nprocs=world_size)
remove('./checkpoint')
if __name__ == '__main__':
test_checkpoint(4, True, 1)
test_checkpoint(4, True, 1, 1)