github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 18:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

[legacy] clean up legacy code (#4743)

Browse Source 
* [legacy] remove outdated codes of pipeline (#4692)

* [legacy] remove cli of benchmark and update optim (#4690)

* [legacy] remove cli of benchmark and update optim

* [doc] fix cli doc test

* [legacy] fix engine clip grad norm

* [legacy] remove outdated colo tensor (#4694)

* [legacy] remove outdated colo tensor

* [test] fix test import

* [legacy] move outdated zero to legacy (#4696)

* [legacy] clean up utils (#4700)

* [legacy] clean up utils

* [example] update examples

* [legacy] clean up amp

* [legacy] fix amp module

* [legacy] clean up gpc (#4742)

* [legacy] clean up context

* [legacy] clean core, constants and global vars

* [legacy] refactor initialize

* [example] fix examples ci

* [example] fix examples ci

* [legacy] fix tests

* [example] fix gpt example

* [example] fix examples ci

* [devops] fix ci installation

* [example] fix examples ci
This commit is contained in:
Hongxin Liu
2023-09-18 16:31:06 +08:00
committed by GitHub
parent 32e7f99416
commit b5f9e37c70
342 changed files with 2919 additions and 4182 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tests/test_tensor/common_utils/__init__.py
View File

@@ -1 +0,0 @@
from ._utils import *

91
tests/test_tensor/common_utils/_utils.py
View File

@@ -1,91 +0,0 @@
import os
import random
import numpy as np
import torch
import torch.distributed as dist
from torch.testing import assert_close
from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.tensor import ComputePattern, ComputeSpec, ShardSpec
def set_seed(seed):
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def check_equal(A, B):
assert torch.allclose(A, B, rtol=1e-3, atol=1e-1) == True
def replace_parameter_add_grad(layer, weight=None, bias=None):
if weight is not None:
delattr(layer, 'weight')
setattr(layer, 'weight', weight)
layer.weight.requires_grad = True
if bias is not None:
delattr(layer, 'bias')
setattr(layer, 'bias', bias)
layer.bias.requires_grad = True
def broadcast_tensor_chunk(tensor, chunk_size=1, local_rank=0):
dist.broadcast(tensor, src=0)
tensor_chunk = torch.chunk(tensor, chunk_size, dim=-1)[local_rank]
return tensor_chunk.clone()
def tensor_equal(t_a: torch.Tensor, t_b: torch.Tensor, rtol: float = 1e-3, atol: float = 1e-1):
assert_close(t_a, t_b, rtol=rtol, atol=atol)
return True
def tensor_shard_equal(tensor: torch.Tensor,
shard: torch.Tensor,
rank: int,
world_size: int,
rtol: float = 1e-3,
atol: float = 1e-1):
assert tensor.ndim == shard.ndim
if tensor.shape == shard.shape:
return tensor_equal(tensor, shard, rtol, atol)
else:
dims_not_eq = torch.nonzero(torch.tensor(tensor.shape) != torch.tensor(shard.shape))
if dims_not_eq.numel() == 1:
# 1D shard
dim = dims_not_eq.item()
if world_size is None:
world_size = gpc.get_world_size(ParallelMode.PARALLEL_1D)
if rank is None:
rank = gpc.get_local_rank(ParallelMode.PARALLEL_1D)
return tensor_equal(tensor.chunk(world_size, dim)[rank], shard, rtol, atol)
else:
raise NotImplementedError
def split_param_single_dim_tp1d(dim, param, pg):
spec = (ShardSpec([dim], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D))
if param.process_group.tp_world_size() == 1:
param.set_process_group(pg)
param.set_tensor_spec(*spec)
def split_param_row_tp1d(param, pg):
split_param_single_dim_tp1d(0, param, pg)
def split_param_col_tp1d(param, pg):
split_param_single_dim_tp1d(-1, param, pg)
def debug_print(ranks, *args):
if dist.get_rank() in ranks:
print(*args)
dist.barrier()

64
tests/test_tensor/core/test_dist_spec_mgr.py
View File

@@ -1,64 +0,0 @@
import math
import pytest
import torch
import torch.distributed as dist
import colossalai
from colossalai.tensor import DistSpecManager, ProcessGroup, ReplicaSpec, ShardSpec
from colossalai.testing import rerun_if_address_is_in_use, spawn
def run():
group = ProcessGroup(tp_degree=dist.get_world_size())
rank = dist.get_rank()
size = dist.get_world_size()
depth = int(math.sqrt(size))
assert depth == math.sqrt(size)
x = torch.rand(8, 8).cuda()
old_dist_spec = ReplicaSpec()
row_spec = ShardSpec([0], [size])
col_spec = ShardSpec([-1], [size])
mat_spec = ShardSpec([0, 1], [depth, depth])
row_shard = DistSpecManager._shard_as(x, old_dist_spec, row_spec, group)
assert torch.equal(x.chunk(size, 0)[rank], row_shard)
assert torch.equal(x, DistSpecManager._gather(row_shard, row_spec, group))
col_shard = DistSpecManager._all_to_all(row_shard, row_spec, col_spec, group)
assert torch.equal(x.chunk(size, -1)[rank], col_shard)
assert torch.equal(x, DistSpecManager._gather(col_shard, col_spec, group))
mat_shard = DistSpecManager._shard_as(x, old_dist_spec, mat_spec, group)
assert torch.equal(x.chunk(depth, 0)[rank // depth].chunk(depth, 1)[rank % depth], mat_shard)
assert torch.equal(x, DistSpecManager._gather(mat_shard, mat_spec, group))
def check_mem():
pg = ProcessGroup(tp_degree=dist.get_world_size())
size = dist.get_world_size()
assert torch.cuda.memory_allocated() == 0
x = torch.rand(32, 32).cuda()
orig_mem = x.numel() * x.element_size()
assert torch.cuda.memory_allocated() == orig_mem
old_dist_spec = ReplicaSpec()
row_spec = ShardSpec([0], [size])
x.data = DistSpecManager._shard_as(x, old_dist_spec, row_spec, pg)
assert x.size(0) == 32 // size and x.size(1) == 32
assert torch.cuda.memory_allocated() == orig_mem // size
x.data = DistSpecManager._gather(x, row_spec, pg)
assert torch.cuda.memory_allocated() == orig_mem
def run_dist(rank, world_size, port):
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
check_mem()
run()
@pytest.mark.dist
@pytest.mark.parametrize('world_size', [1, 4])
@rerun_if_address_is_in_use()
def test_dist_spec_mgr(world_size):
spawn(run_dist, world_size)
if __name__ == '__main__':
test_dist_spec_mgr(4)

5
tests/test_tensor/test_comm_spec_apply.py
View File

@@ -1,7 +1,7 @@
import pytest
import torch
import torch.distributed as dist
from colossalai.core import global_context as gpc
from colossalai.device.device_mesh import DeviceMesh
from colossalai.initialize import launch
from colossalai.logging import disable_existing_loggers
@@ -184,7 +184,7 @@ def check_comm(rank, world_size, port):
launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
physical_mesh_id = torch.arange(0, 4)
assert rank == gpc.get_global_rank()
assert rank == dist.get_rank()
mesh_shape = (2, 2)
# [[0, 1,
@@ -205,7 +205,6 @@ def check_comm(rank, world_size, port):
# test all reduce in 1D flatten device mesh
check_all_reduce_in_flatten_device_mesh(device_mesh, rank)
gpc.destroy()
@pytest.mark.dist

6
tests/test_tensor/test_dtensor/test_comm_spec.py
View File

@@ -1,7 +1,7 @@
import pytest
import torch
import torch.distributed as dist
from colossalai.core import global_context as gpc
from colossalai.device.device_mesh import DeviceMesh
from colossalai.initialize import launch
from colossalai.logging import disable_existing_loggers
@@ -127,7 +127,7 @@ def check_comm(rank, world_size, port):
launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
physical_mesh_id = torch.arange(0, 4)
assert rank == gpc.get_global_rank()
assert rank == dist.get_rank()
mesh_shape = (2, 2)
# [[0, 1,
@@ -149,8 +149,6 @@ def check_comm(rank, world_size, port):
check_all_reduce_fwd(process_group_dict, rank)
check_all_reduce_bwd(process_group_dict, rank)
gpc.destroy()
@pytest.mark.dist
@rerun_if_address_is_in_use()

4
tests/test_tensor/test_mix_gather.py
View File

@@ -1,7 +1,7 @@
import pytest
import torch
import torch.distributed as dist
from colossalai.core import global_context as gpc
from colossalai.device.device_mesh import DeviceMesh
from colossalai.initialize import launch
from colossalai.logging import disable_existing_loggers
@@ -295,7 +295,7 @@ def check_comm(rank, world_size, port):
launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
physical_mesh_id = torch.arange(0, 8)
assert rank == gpc.get_global_rank()
assert rank == dist.get_rank()
mesh_shape = (2, 4)
# [[0, 1, 2, 3],

34
tests/test_tensor/test_parameter.py
View File

@@ -1,34 +0,0 @@
import pytest
import torch
from common_utils import tensor_equal
import colossalai
from colossalai.tensor import ColoParameter, ColoTensor, ColoTensorSpec, ProcessGroup
from colossalai.testing import free_port
@pytest.mark.skip
def test_multiinheritance():
colossalai.launch(config={}, rank=0, world_size=1, host='localhost', port=free_port(), backend='nccl')
colo_param = ColoParameter(None, requires_grad=True)
assert colo_param.dist_spec.placement.value == 'r'
assert isinstance(colo_param, ColoTensor)
assert isinstance(colo_param, torch.nn.Parameter)
# __deepcopy__ overload
import copy
colo_param2 = copy.deepcopy(colo_param)
assert isinstance(colo_param2, ColoParameter)
assert tensor_equal(colo_param.data, colo_param2.data)
assert colo_param.requires_grad == colo_param2.requires_grad
# __repr__ overload
assert 'ColoParameter' in str(colo_param)
# __torch_function__
clone_param = torch.clone(colo_param)
assert isinstance(clone_param, ColoTensor)
if __name__ == '__main__':
test_multiinheritance()