From 36f9a74ab26c1d6094bd71171b293783391eb48e Mon Sep 17 00:00:00 2001 From: ver217 Date: Fri, 4 Mar 2022 13:40:48 +0800 Subject: [PATCH] fix sharded param hook and unit test --- .../engine/ophooks/_shard_param_ophook.py | 13 ++-- .../zero/sharded_model/sharded_model_v2.py | 7 ++- .../zero/sharded_optim/sharded_optim_v2.py | 13 +--- .../zero/sharded_param/sharded_param.py | 14 +++-- tests/test_zero_data_parallel/common.py | 61 ++++++++----------- .../test_sharded_optim_v2.py | 7 +-- 6 files changed, 49 insertions(+), 66 deletions(-) diff --git a/colossalai/engine/ophooks/_shard_param_ophook.py b/colossalai/engine/ophooks/_shard_param_ophook.py index 801370f6b..57f76970c 100644 --- a/colossalai/engine/ophooks/_shard_param_ophook.py +++ b/colossalai/engine/ophooks/_shard_param_ophook.py @@ -1,5 +1,4 @@ import torch -import torch.distributed as dist from colossalai.registry import OPHOOKS from . import BaseOpHook @@ -21,29 +20,25 @@ 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}') + param.data = param.ca_attr.payload() 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}') + param.data = param.ca_attr.payload() 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}') + param.data = param.ca_attr.payload() 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}') + param.data = param.ca_attr.payload() def pre_iter(self): pass diff --git a/colossalai/zero/sharded_model/sharded_model_v2.py b/colossalai/zero/sharded_model/sharded_model_v2.py index 7dbc99c3d..9f1a8a95f 100644 --- a/colossalai/zero/sharded_model/sharded_model_v2.py +++ b/colossalai/zero/sharded_model/sharded_model_v2.py @@ -6,8 +6,7 @@ import torch.distributed as dist import torch.nn as nn from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc -from colossalai.engine.ophooks import (ShardGradHook, ShardParamHook, - register_ophooks_recursively) +from colossalai.engine.ophooks import (ShardGradHook, ShardParamHook, register_ophooks_recursively) from colossalai.engine.paramhooks import BaseParamHookMgr from colossalai.logging import get_dist_logger from colossalai.zero.sharded_model.reduce_scatter import ReduceScatterBucketer @@ -109,6 +108,10 @@ class ShardedModelV2(nn.Module): if not self._require_backward_grad_sync: continue p._sharded_grad.write_back() + # In case some post bwd hook is not fired + for p in self.module.parameters(): + if not p.ca_attr.is_sharded: + p.ca_attr.shard() @torch.no_grad() def _grad_post_backward_hook(self, param: Parameter, grad: torch.Tensor) -> Optional[torch.Tensor]: diff --git a/colossalai/zero/sharded_optim/sharded_optim_v2.py b/colossalai/zero/sharded_optim/sharded_optim_v2.py index e07b886bb..5b4cb1122 100644 --- a/colossalai/zero/sharded_optim/sharded_optim_v2.py +++ b/colossalai/zero/sharded_optim/sharded_optim_v2.py @@ -14,7 +14,6 @@ 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 @@ -63,8 +62,6 @@ 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: @@ -91,23 +88,15 @@ class ShardedOptimizerV2(ColossalaiOptimizer): for group in self.optim.param_groups: for p in group['params']: if hasattr(p, 'ca_attr'): - 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) + p.data = p.ca_attr.payload() 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/colossalai/zero/sharded_param/sharded_param.py b/colossalai/zero/sharded_param/sharded_param.py index 56f2382af..472487d6a 100644 --- a/colossalai/zero/sharded_param/sharded_param.py +++ b/colossalai/zero/sharded_param/sharded_param.py @@ -1,10 +1,11 @@ +from typing import Optional, Tuple, Union + +import numpy import torch import torch.distributed as dist from colossalai.context.parallel_mode import ParallelMode from colossalai.core import global_context as gpc from colossalai.zero.sharded_model._zero3_utils import get_shard -from typing import Union, Tuple, Optional -import numpy class ShardedParam(object): @@ -28,6 +29,7 @@ class ShardedParam(object): self.world_size = dist.get_world_size(self.process_group) self.local_rank = dist.get_rank(self.process_group) self.is_sharded = False + self.device = device # Hijack the data payload of param if isinstance(other, torch.nn.Parameter): @@ -50,17 +52,19 @@ class ShardedParam(object): self._payload_numel = None - def payload(self, target_device: torch.device): + def payload(self, target_device: Optional[torch.device] = None): r""" get the payload and move it to target device """ - return self._param_payload.to(target_device) + if target_device is not None: + return self._param_payload.to(target_device) + return self._param_payload def set_payload(self, data: torch.Tensor): r""" set payload as data """ - assert self._param_payload.numel() == data.numel() + assert self._param_payload.shape == data.shape self._param_payload.copy_(data) def shard(self): diff --git a/tests/test_zero_data_parallel/common.py b/tests/test_zero_data_parallel/common.py index 351831f97..8492f5225 100644 --- a/tests/test_zero_data_parallel/common.py +++ b/tests/test_zero_data_parallel/common.py @@ -3,37 +3,21 @@ from functools import partial import torch import torch.distributed as dist import torch.nn as nn -from colossalai.logging import disable_existing_loggers, get_dist_logger +from colossalai.logging import get_dist_logger from colossalai.utils import checkpoint LOGGER = get_dist_logger() -CONFIG = dict( - fp16=dict( - mode=None, - ), - zero=dict( - level=3, - verbose=False, - offload_optimizer_config=dict( - device='cpu', - pin_memory=True, - buffer_count=5, - fast_init=False - ), - offload_param_config=dict( - device='cpu', - pin_memory=True, - buffer_count=5, - buffer_size=1e8, - max_in_cpu=1e9 - ) - ), - parallel=dict( - pipeline=dict(size=1), - tensor=dict(size=1, mode=None) - ) -) +CONFIG = dict(fp16=dict(mode=None,), + zero=dict(level=3, + verbose=False, + offload_optimizer_config=dict(device='cpu', pin_memory=True, buffer_count=5, fast_init=False), + offload_param_config=dict(device='cpu', + pin_memory=True, + buffer_count=5, + buffer_size=1e8, + max_in_cpu=1e9)), + parallel=dict(pipeline=dict(size=1), tensor=dict(size=1, mode=None))) def checkpoint_wrapper(module, enable=True): @@ -43,6 +27,7 @@ def checkpoint_wrapper(module, enable=True): class Net(nn.Module): + def __init__(self, checkpoint=False) -> None: super().__init__() self.fc1 = nn.Linear(5, 5) @@ -50,13 +35,7 @@ class Net(nn.Module): self.fc3 = nn.Linear(5, 1) if checkpoint: self.fc1 = checkpoint_wrapper(self.fc1) - self.layers = [ - self.fc1, - self.fc2, - self.fc1, - self.fc2, - self.fc3 - ] + self.layers = [self.fc1, self.fc2, self.fc1, self.fc2, self.fc3] def forward(self, x): for layer in self.layers: @@ -111,3 +90,17 @@ def check_params_padding(model, zero_model, loose=False): zero_p = zero_p[:p.size(0)] assert p.dtype == zero_p.dtype assert allclose(p, zero_p, loose=loose) + + +def check_sharded_params_padding(model, zero_model, loose=False): + rank = dist.get_rank() + for p, zero_p in zip(model.parameters(), zero_model.parameters()): + zero_p = zero_p.ca_attr.payload(p.device) + chunks = torch.flatten(p).chunk(dist.get_world_size()) + if rank >= len(chunks): + continue + p = chunks[rank] + if zero_p.size(0) > p.size(0): + zero_p = zero_p[:p.size(0)] + assert p.dtype == zero_p.dtype + assert allclose(p, zero_p, loose=loose) diff --git a/tests/test_zero_data_parallel/test_sharded_optim_v2.py b/tests/test_zero_data_parallel/test_sharded_optim_v2.py index b5f8f098f..dfd612182 100644 --- a/tests/test_zero_data_parallel/test_sharded_optim_v2.py +++ b/tests/test_zero_data_parallel/test_sharded_optim_v2.py @@ -16,19 +16,18 @@ 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) +from common import (CONFIG, Net, check_grads, check_grads_padding, check_params, check_sharded_params_padding) def run_step(model, optimizer, x, enable_autocast=False): model.train() + optimizer.zero_grad() 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() @@ -51,7 +50,7 @@ def run_dist(rank, world_size, port): run_step(model, optim, x, False) if dist.get_world_size() > 1: check_grads_padding(model, zero_model) - check_params_padding(model, zero_model) + check_sharded_params_padding(model, zero_model) else: check_grads(model, zero_model) check_params(model, zero_model)