mirror of https://github.com/hpcaitech/ColossalAI
107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
from typing import Optional
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import colossalai
|
|
from colossalai.device.device_mesh import DeviceMesh
|
|
from colossalai.tensor.d_tensor.layout import Layout
|
|
from colossalai.tensor.d_tensor.sharding_spec import ShardingSpec
|
|
from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
|
|
from colossalai.utils.common import print_rank_0
|
|
|
|
try:
|
|
from colossalai.lazy.lazy_init import LazyInitContext, LazyTensor, _MyTensor
|
|
except:
|
|
pass
|
|
from lazy_init_utils import SUPPORT_LAZY, assert_dist_model_equal, set_seed
|
|
|
|
from tests.kit.model_zoo import model_zoo
|
|
|
|
|
|
def find_shard_dim(shape: torch.Size) -> Optional[int]:
|
|
for dim, size in enumerate(shape):
|
|
if size % 2 == 0:
|
|
return dim
|
|
|
|
|
|
def make_layout(device_mesh: DeviceMesh, original_tensor: torch.Tensor) -> Layout:
|
|
shard_dim = find_shard_dim(original_tensor.shape)
|
|
dim_partition_dict = {shard_dim: [0]} if shard_dim is not None else {}
|
|
target_sharding_spec = ShardingSpec(dim_size=original_tensor.dim(), dim_partition_dict=dim_partition_dict)
|
|
layout = Layout(device_mesh=device_mesh,
|
|
device_type=torch.device('cuda'),
|
|
sharding_spec=target_sharding_spec,
|
|
entire_shape=original_tensor.shape)
|
|
return layout
|
|
|
|
|
|
def _get_current_name(prefix: str, name: str) -> str:
|
|
return f'{prefix}.{name}'.lstrip('.')
|
|
|
|
|
|
def generate_layout_dict(model: nn.Module, device_mesh: DeviceMesh) -> dict:
|
|
layout_dict = {}
|
|
|
|
@torch.no_grad()
|
|
def generate_recursively(module: nn.Module, prefix: str = ''):
|
|
# recursively initialize the module
|
|
for name, mod in module.named_children():
|
|
generate_recursively(mod, prefix=_get_current_name(prefix, name))
|
|
|
|
# initialize tensors directly attached to the current module
|
|
for name, param in module.named_parameters(recurse=False):
|
|
if isinstance(param, LazyTensor):
|
|
layout = make_layout(device_mesh, param)
|
|
layout_dict[_get_current_name(prefix, name)] = layout
|
|
|
|
for name, buf in module.named_buffers(recurse=False):
|
|
if isinstance(buf, LazyTensor):
|
|
layout = make_layout(device_mesh, buf)
|
|
layout_dict[_get_current_name(prefix, name)] = layout
|
|
|
|
generate_recursively(model)
|
|
|
|
return layout_dict
|
|
|
|
|
|
@parameterize('subset', ['torchvision', 'diffusers', 'timm', 'transformers', 'torchaudio', 'deepfm', 'dlrm'])
|
|
def run_dist_lazy_init(subset, seed: int = 42):
|
|
sub_model_zoo = model_zoo.get_sub_registry(subset)
|
|
device_mesh = DeviceMesh(torch.Tensor([0, 1, 2, 3]), (2, 2), init_process_group=True)
|
|
_MyTensor._pre_op_fn = lambda *args: set_seed(seed)
|
|
LazyTensor._pre_op_fn = lambda *args: set_seed(seed)
|
|
|
|
for name, entry in sub_model_zoo.items():
|
|
# TODO(ver217): lazy init does not support weight norm, skip these models
|
|
if name in ('torchaudio_wav2vec2_base', 'torchaudio_hubert_base'):
|
|
continue
|
|
print_rank_0(name)
|
|
model_fn, data_gen_fn, output_transform_fn, model_attr = entry
|
|
ctx = LazyInitContext(tensor_cls=_MyTensor)
|
|
with ctx:
|
|
model = model_fn()
|
|
ctx = LazyInitContext()
|
|
with ctx:
|
|
deferred_model = model_fn()
|
|
layout_dict = generate_layout_dict(deferred_model, device_mesh)
|
|
ctx.distribute(deferred_model, layout_dict, verbose=True)
|
|
assert_dist_model_equal(model, deferred_model, layout_dict)
|
|
|
|
|
|
def run_dist(rank, world_size, port) -> None:
|
|
colossalai.launch({}, rank=rank, world_size=world_size, host='localhost', port=port)
|
|
run_dist_lazy_init()
|
|
|
|
|
|
@pytest.mark.skipif(not SUPPORT_LAZY, reason='torch version should be >= 1.12.0')
|
|
@pytest.mark.dist
|
|
@rerun_if_address_is_in_use()
|
|
def test_dist_lazy_init():
|
|
spawn(run_dist, 4)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
test_dist_lazy_init()
|