From 001ca624dd4b96c68489827826b9730a074ac704 Mon Sep 17 00:00:00 2001 From: ver217 Date: Fri, 4 Mar 2022 11:49:02 +0800 Subject: [PATCH] impl shard optim v2 and add unit test --- .../engine/ophooks/_shard_param_ophook.py | 15 +++- colossalai/zero/sharded_optim/__init__.py | 3 +- .../zero/sharded_optim/sharded_optim_v2.py | 17 ++++- .../test_sharded_optim_v2.py | 68 +++++++++++++++++++ 4 files changed, 97 insertions(+), 6 deletions(-) create mode 100644 tests/test_zero_data_parallel/test_sharded_optim_v2.py diff --git a/colossalai/engine/ophooks/_shard_param_ophook.py b/colossalai/engine/ophooks/_shard_param_ophook.py index 5bee3f9a4..801370f6b 100644 --- a/colossalai/engine/ophooks/_shard_param_ophook.py +++ b/colossalai/engine/ophooks/_shard_param_ophook.py @@ -1,12 +1,16 @@ import torch -from . import BaseOpHook +import torch.distributed as dist from colossalai.registry import OPHOOKS +from . import BaseOpHook + + @OPHOOKS.register_module class ShardParamHook(BaseOpHook): """ A hook to process sharded param before and afther FWD and BWD operator executing. """ + def __init__(self): super().__init__() @@ -17,25 +21,32 @@ class ShardParamHook(BaseOpHook): for param in module.parameters(): assert hasattr(param, 'ca_attr') param.ca_attr.gather() + if dist.get_rank() == 0: + print(f'{param._name} pre fwd shape {param.ca_attr.payload("cpu").shape}') def post_fwd_exec(self, module: torch.nn.Module, *args): for param in module.parameters(): assert hasattr(param, 'ca_attr') param.ca_attr.shard() + if dist.get_rank() == 0: + print(f'{param._name} post fwd shape {param.ca_attr.payload("cpu").shape}') def pre_bwd_exec(self, module: torch.nn.Module, input, output): for param in module.parameters(): assert hasattr(param, 'ca_attr') param.ca_attr.gather() + if dist.get_rank() == 0: + print(f'{param._name} pre bwd shape {param.ca_attr.payload("cpu").shape}') def post_bwd_exec(self, module: torch.nn.Module, input): for param in module.parameters(): assert hasattr(param, 'ca_attr') param.ca_attr.shard() + if dist.get_rank() == 0: + print(f'{param._name} post bwd shape {param.ca_attr.payload("cpu").shape}') def pre_iter(self): pass def post_iter(self): pass - diff --git a/colossalai/zero/sharded_optim/__init__.py b/colossalai/zero/sharded_optim/__init__.py index 07681531c..94ef14b95 100644 --- a/colossalai/zero/sharded_optim/__init__.py +++ b/colossalai/zero/sharded_optim/__init__.py @@ -1,3 +1,4 @@ from .sharded_optim import ShardedOptimizer +from .sharded_optim_v2 import ShardedOptimizerV2 -__all__ = ['ShardedOptimizer'] \ No newline at end of file +__all__ = ['ShardedOptimizer', 'ShardedOptimizerV2'] diff --git a/colossalai/zero/sharded_optim/sharded_optim_v2.py b/colossalai/zero/sharded_optim/sharded_optim_v2.py index dc8b48db8..e07b886bb 100644 --- a/colossalai/zero/sharded_optim/sharded_optim_v2.py +++ b/colossalai/zero/sharded_optim/sharded_optim_v2.py @@ -14,6 +14,7 @@ from torch.distributed import ProcessGroup from torch.nn.parameter import Parameter from torch.optim import Optimizer +from ..sharded_model._zero3_utils import free_storage from ._utils import has_inf_or_nan @@ -62,6 +63,8 @@ class ShardedOptimizerV2(ColossalaiOptimizer): if hasattr(p, 'ca_attr'): assert p.ca_attr.is_sharded, 'ShardedAdam can be only used with sharded model' self.master_params[p] = p.ca_attr.payload(self.device) + if dist.get_rank() == 0: + print(f'load payload {p._name} {self.master_params[p].shape}') else: self.master_params[p] = p.data.to(device=self.device) if torch.is_floating_point(self.master_params[p]) and self.master_params[p].dtype != torch.float: @@ -84,19 +87,27 @@ class ShardedOptimizerV2(ColossalaiOptimizer): for p in group['params']: p.data = self.master_params[p] ret = self.optim.step(*args, **kwargs) - # Write master param to payload and set p.data to None + # Write master param to payload for group in self.optim.param_groups: for p in group['params']: if hasattr(p, 'ca_attr'): - # TODO: update payload - p.data = None + if dist.get_rank() == 0: + print(f'write {p._name} {p.shape} orig_shape {p.ca_attr._origin_shape} \ + payload shape {p.ca_attr._param_payload.shape} sharded {p.ca_attr.is_sharded}') + p.ca_attr.set_payload(p.data) + # We cannot set p.data to None directly, so we free storage + free_storage(p.data) return ret def backward(self, loss: Tensor) -> None: loss = self.loss_scale * loss self.optim_state = OptimState.SCALED if self.model_is_sharded: + if dist.get_rank() == 0: + print('sharded model backward') self.model.backward(loss) + if dist.get_rank() == 0: + print('sharded model backward done') else: super().backward(loss) diff --git a/tests/test_zero_data_parallel/test_sharded_optim_v2.py b/tests/test_zero_data_parallel/test_sharded_optim_v2.py new file mode 100644 index 000000000..b5f8f098f --- /dev/null +++ b/tests/test_zero_data_parallel/test_sharded_optim_v2.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python +# -*- encoding: utf-8 -*- + +import copy +from functools import partial + +import colossalai +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.sharded_model import ShardedModelV2 +from colossalai.zero.sharded_optim import ShardedOptimizerV2 +from torch.optim import Adam + +from common import (CONFIG, Net, check_grads, check_grads_padding, check_params, check_params_padding) + + +def run_step(model, optimizer, x, enable_autocast=False): + model.train() + with torch.cuda.amp.autocast(enabled=enable_autocast): + y = model(x) + loss = y.sum() + loss = loss.float() + if isinstance(model, ShardedModelV2): + optimizer.backward(loss) + for p in model.parameters(): + assert p.ca_attr.is_sharded + else: + loss.backward() + optimizer.step() + + +def run_dist(rank, world_size, port): + colossalai.launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl') + + model = Net(checkpoint=True).cuda() + zero_model = copy.deepcopy(model) + zero_model = ShardedModelV2(zero_model, process_group=gpc.get_group(ParallelMode.DATA)) + for n, p in zero_model.named_parameters(): + p._name = n + optim = Adam(model.parameters(), lr=1e-3) + sharded_optim = ShardedOptimizerV2(Adam(zero_model.parameters(), lr=1e-3), zero_model) + + for _ in range(2): + x = torch.rand(2, 5).cuda() + run_step(zero_model, sharded_optim, x, False) + run_step(model, optim, x, False) + if dist.get_world_size() > 1: + check_grads_padding(model, zero_model) + check_params_padding(model, zero_model) + else: + check_grads(model, zero_model) + check_params(model, zero_model) + + +@pytest.mark.dist +def test_sharded_optim_v2(): + world_size = 2 + run_func = partial(run_dist, world_size=world_size, port=free_port()) + mp.spawn(run_func, nprocs=world_size) + + +if __name__ == '__main__': + test_sharded_optim_v2()