mirror of https://github.com/hpcaitech/ColossalAI
[zero] alleviate memory usage in ZeRODDP state_dict (#1398)
HELSON
GitHub
parent
4f5f8f77d1
commit
4e98e938ce
|
|
@ -6,12 +6,13 @@ from colossalai.zero.utils.zero_hook_v2 import ZeROHookV2
|
|||
from colossalai.gemini.chunk import TensorState, Chunk
|
||||
from colossalai.tensor.param_op_hook import ParamOpHookManager
|
||||
from colossalai.gemini.gemini_mgr import GeminiManager
|
||||
from typing import Dict, Iterable, List, Optional
|
||||
from typing import Dict, Iterable, List, Optional, Set
|
||||
from colossalai.logging import get_dist_logger
|
||||
from collections import OrderedDict
|
||||
from colossalai.tensor.colo_parameter import ColoParameter
|
||||
from colossalai.tensor import ProcessGroup as ColoProcessGroup
|
||||
from .reducer import Reducer
|
||||
|
||||
try:
|
||||
from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX, _IncompatibleKeys
|
||||
except ImportError:
|
||||
|
|
@ -84,6 +85,18 @@ class ColoDDP(torch.nn.Module):
|
|||
def named_parameters(self, prefix: str = '', recurse: bool = True):
|
||||
return self.module.named_parameters(prefix, recurse)
|
||||
|
||||
def named_buffers(self, prefix: str = '', recurse: bool = True):
|
||||
return self.module.named_buffers(prefix, recurse)
|
||||
|
||||
def named_children(self):
|
||||
return self.module.named_children()
|
||||
|
||||
def named_modules(self,
|
||||
memo: Optional[Set[torch.nn.Module]] = None,
|
||||
prefix: str = '',
|
||||
remove_duplicate: bool = True):
|
||||
return self.module.named_modules(memo, prefix, remove_duplicate)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
self.module.zero_grad(set_to_none=True)
|
||||
return self.module(*args, **kwargs)
|
||||
|
|
@ -274,7 +287,7 @@ class ZeroDDP(ColoDDP):
|
|||
for tensor in chunk.get_tensors():
|
||||
self.grads_device[tensor] = device
|
||||
|
||||
def state_dict(self, destination=None, prefix='', keep_vars=False):
|
||||
def state_dict(self, destination=None, prefix='', keep_vars=False, only_rank_0: bool = True):
|
||||
r"""Returns a dictionary containing a whole state of the module.
|
||||
|
||||
Both parameters and persistent buffers (e.g. running averages) are
|
||||
|
|
@ -291,18 +304,22 @@ class ZeroDDP(ColoDDP):
|
|||
['bias', 'weight']
|
||||
|
||||
"""
|
||||
is_rank_0 = self.chunk_manager.process_group.dp_local_rank() == 0
|
||||
record_flag = (not only_rank_0) or is_rank_0
|
||||
|
||||
if destination is None:
|
||||
destination = OrderedDict()
|
||||
destination._metadata = OrderedDict()
|
||||
destination._metadata[prefix[:-1]] = local_metadata = dict(version=self._version)
|
||||
self._save_to_state_dict(destination, prefix, keep_vars)
|
||||
self._save_to_state_dict(destination, prefix, keep_vars, record_flag)
|
||||
|
||||
for hook in self._state_dict_hooks.values():
|
||||
hook_result = hook(self, destination, prefix, local_metadata)
|
||||
if hook_result is not None:
|
||||
destination = hook_result
|
||||
return destination
|
||||
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars):
|
||||
def _save_to_state_dict(self, destination, prefix, keep_vars, record_flag: bool = True):
|
||||
r"""Saves module state to `destination` dictionary, containing a state
|
||||
of the module, but not its descendants. This is called on every
|
||||
submodule in :meth:`~torch.nn.Module.state_dict`.
|
||||
|
|
@ -315,22 +332,36 @@ class ZeroDDP(ColoDDP):
|
|||
prefix (str): the prefix for parameters and buffers used in this
|
||||
module
|
||||
"""
|
||||
chunks = self.chunk_manager.get_chunks(self.fp32_params)
|
||||
chunks_orig_device_type = []
|
||||
for chunk in chunks:
|
||||
chunks_orig_device_type.append(chunk.device_type)
|
||||
# save parameters
|
||||
param_to_save_data = dict()
|
||||
chunk_list = self.chunk_manager.get_chunks(self.fp32_params)
|
||||
for chunk in chunk_list:
|
||||
# record the original device of the chunk
|
||||
org_chunk_dev_typ = chunk.device_type
|
||||
self.chunk_manager.access_chunk(chunk)
|
||||
|
||||
for tensor in chunk.get_tensors():
|
||||
rec_p = torch.empty([0])
|
||||
if record_flag:
|
||||
rec_p = tensor.cpu() # move the whole tensor to CPU mem
|
||||
assert tensor not in param_to_save_data
|
||||
param_to_save_data[tensor] = rec_p
|
||||
# release the actual memory of the chunk
|
||||
self.chunk_manager.release_chunk(chunk)
|
||||
if not chunk.is_empty and org_chunk_dev_typ == 'cpu':
|
||||
self.chunk_manager.move_chunk(chunk, torch.device('cpu'))
|
||||
|
||||
for (name, p), fp32_p in zip(self.named_parameters(), self.fp32_params):
|
||||
if p is not None:
|
||||
rec_p = fp32_p.clone() if fp32_p.device.type == 'cpu' else fp32_p.cpu()
|
||||
assert fp32_p in param_to_save_data, "Parameter '{}' is neglected in the chunk list".format(name)
|
||||
rec_p = param_to_save_data[fp32_p]
|
||||
destination[prefix + name] = rec_p if keep_vars else rec_p.detach()
|
||||
for orig_dvice_type, chunk in zip(chunks_orig_device_type, chunks):
|
||||
self.chunk_manager.release_chunk(chunk)
|
||||
if not chunk.is_empty and orig_dvice_type == 'cpu':
|
||||
self.chunk_manager.move_chunk(chunk, torch.device('cpu'))
|
||||
|
||||
# save all buffers
|
||||
for name, buf in self.named_buffers():
|
||||
if buf is not None and name not in self._non_persistent_buffers_set:
|
||||
destination[prefix + name] = buf if keep_vars else buf.detach()
|
||||
# save extra states
|
||||
extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX
|
||||
if getattr(self.__class__, "get_extra_state",
|
||||
torch.nn.Module.get_extra_state) is not torch.nn.Module.get_extra_state:
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import copy
|
||||
|
||||
import pytest
|
||||
import colossalai
|
||||
import torch
|
||||
|
|
@ -11,9 +13,9 @@ from functools import partial
|
|||
from tests.components_to_test.registry import non_distributed_component_funcs
|
||||
from colossalai.nn.parallel import ZeroDDP, ColoDDP
|
||||
from colossalai.gemini.gemini_mgr import GeminiManager
|
||||
from typing import Callable
|
||||
from collections import OrderedDict
|
||||
from colossalai.tensor import ProcessGroup, ColoParameter
|
||||
from colossalai.testing import parameterize
|
||||
|
||||
|
||||
def check_state_dict_equal(state_dict: OrderedDict, other_state_dict: OrderedDict):
|
||||
|
|
@ -25,7 +27,27 @@ def check_state_dict_equal(state_dict: OrderedDict, other_state_dict: OrderedDic
|
|||
else:
|
||||
temp_t2 = t2
|
||||
|
||||
assert torch.equal(t1, temp_t2)
|
||||
assert torch.equal(t1, temp_t2), "\t{}\n\t{}".format(t1, temp_t2)
|
||||
|
||||
|
||||
def check_model_equal(model_a, model_b, allow_empty: bool = False, same_dtype: bool = True):
|
||||
for (na, pa), (nb, pb) in zip(model_a.named_parameters(), model_b.named_parameters()):
|
||||
assert na == nb
|
||||
|
||||
if not allow_empty:
|
||||
assert pa.storage().size() > 0
|
||||
assert pb.storage().size() > 0
|
||||
else:
|
||||
if pa.storage().size() == 0 or pb.storage().size() == 0:
|
||||
continue
|
||||
|
||||
if same_dtype:
|
||||
assert pa.dtype == pb.dtype
|
||||
temp_pb = pb
|
||||
else:
|
||||
temp_pb = pb.to(pa.dtype)
|
||||
|
||||
assert torch.equal(pa, temp_pb), "Parameter '{}' is not equal.\n {} {}".format(na, pa, pb)
|
||||
|
||||
|
||||
def init_ddp(module: torch.nn.Module) -> ColoDDP:
|
||||
|
|
@ -33,22 +55,26 @@ def init_ddp(module: torch.nn.Module) -> ColoDDP:
|
|||
return ColoDDP(module, process_group=pg)
|
||||
|
||||
|
||||
def init_ddpv2(module: torch.nn.Module, use_chunk: bool = False, use_zero: bool = False) -> ZeroDDP:
|
||||
def init_ddpv2(module: torch.nn.Module,
|
||||
use_chunk: bool = False,
|
||||
use_zero: bool = False,
|
||||
placement_policy: str = 'cuda') -> ZeroDDP:
|
||||
pg = ProcessGroup()
|
||||
chunk_size = ChunkManager.search_chunk_size(module, 64, 4) if use_chunk else None
|
||||
chunk_manager = ChunkManager(chunk_size, pg, enable_distributed_storage=use_zero)
|
||||
gemini_manager = GeminiManager('cuda', chunk_manager)
|
||||
gemini_manager = GeminiManager(placement_policy, chunk_manager)
|
||||
return ZeroDDP(module, gemini_manager)
|
||||
|
||||
|
||||
def run_state_dict(ddp_init_func: Callable[[torch.nn.Module], ColoDDP]):
|
||||
get_components_func = non_distributed_component_funcs.get_callable('nested_model')
|
||||
def run_ddp_state_dict():
|
||||
get_components_func = non_distributed_component_funcs.get_callable('gpt2')
|
||||
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
|
||||
torch_model = model_builder().cuda()
|
||||
with ColoInitContext(device=get_current_device()):
|
||||
model = model_builder()
|
||||
model = ddp_init_func(model)
|
||||
model = init_ddp(model)
|
||||
torch_state_dict = torch_model.state_dict()
|
||||
|
||||
for param in model.parameters():
|
||||
if isinstance(param, ColoParameter):
|
||||
assert param.get_process_group() is not None
|
||||
|
|
@ -62,13 +88,44 @@ def run_state_dict(ddp_init_func: Callable[[torch.nn.Module], ColoDDP]):
|
|||
check_state_dict_equal(torch_state_dict, state_dict)
|
||||
|
||||
|
||||
@parameterize('use_chunk', [False, True])
|
||||
@parameterize('placement_policy', ['cuda', 'cpu'])
|
||||
@parameterize('use_zero', [False, True])
|
||||
@parameterize('only_rank_0', [False, True])
|
||||
def run_zero_state_dict(use_chunk, placement_policy, use_zero, only_rank_0):
|
||||
get_components_func = non_distributed_component_funcs.get_callable('gpt2')
|
||||
model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
|
||||
|
||||
torch_model = model_builder().cuda()
|
||||
org_torch_model = copy.deepcopy(torch_model)
|
||||
torch_state_dict = torch_model.state_dict()
|
||||
|
||||
with ColoInitContext(device=get_current_device()):
|
||||
model = model_builder()
|
||||
model = init_ddpv2(model, use_chunk, use_zero, placement_policy)
|
||||
|
||||
for param in model.parameters():
|
||||
if isinstance(param, ColoParameter):
|
||||
assert param.get_process_group() is not None
|
||||
|
||||
model.load_state_dict(torch_state_dict, strict=False)
|
||||
check_model_equal(model, torch_model, allow_empty=True, same_dtype=False)
|
||||
|
||||
for param in model.parameters():
|
||||
if isinstance(param, ColoParameter):
|
||||
assert param.get_process_group() is not None
|
||||
|
||||
pg = ProcessGroup()
|
||||
state_dict = model.state_dict(only_rank_0=only_rank_0)
|
||||
if not only_rank_0 or pg.dp_local_rank() == 0:
|
||||
torch_model.load_state_dict(state_dict, strict=False)
|
||||
check_model_equal(torch_model, org_torch_model, allow_empty=False, same_dtype=True)
|
||||
|
||||
|
||||
def run_dist(rank, world_size, port):
|
||||
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
|
||||
run_state_dict(init_ddp)
|
||||
run_state_dict(partial(init_ddpv2, use_chunk=False, use_zero=False))
|
||||
run_state_dict(partial(init_ddpv2, use_chunk=False, use_zero=True))
|
||||
run_state_dict(partial(init_ddpv2, use_chunk=True, use_zero=False))
|
||||
run_state_dict(partial(init_ddpv2, use_chunk=True, use_zero=True))
|
||||
run_ddp_state_dict()
|
||||
run_zero_state_dict()
|
||||
|
||||
|
||||
@pytest.mark.dist
|
||||
|
|
|
|||
Loading…
Reference in New Issue