github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix sharded param hook and unit test

pull/394/head
ver217 2022-03-04 13:40:48 +08:00 committed by Frank Lee
parent 001ca624dd
commit 36f9a74ab2
6 changed files with 49 additions and 66 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
colossalai/engine/ophooks/_shard_param_ophook.py
View File

@ -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

7
colossalai/zero/sharded_model/sharded_model_v2.py
View File

@ -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]:

13
colossalai/zero/sharded_optim/sharded_optim_v2.py
View File

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

14
colossalai/zero/sharded_param/sharded_param.py
View File

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

61
tests/test_zero_data_parallel/common.py
View File

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

7
tests/test_zero_data_parallel/test_sharded_optim_v2.py
View File

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