colossalai

1
import pytest
2
import torch
3
import torch.distributed as dist
4
from torch.distributed.distributed_c10d import _get_default_group
5

6
import colossalai
7
from colossalai.accelerator import get_accelerator
8
from colossalai.tensor import ColoParameter
9
from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn
10
from colossalai.zero.gemini import TensorState
11
from colossalai.zero.gemini.chunk import Chunk
12

13

14
def dist_sum(x):
15
    temp = torch.tensor([x], device=get_accelerator().get_current_device())
16
    dist.all_reduce(temp)
17
    return temp.item()
18

19

20
def add_param(param_list, param_cp_list, *args, **kwargs):
21
    param = ColoParameter(torch.randn(*args, **kwargs))
22
    param_list.append(param)
23
    param_cp_list.append(param.clone())
24

25

26
def check_equal(param, param_cp):
27
    if param.device != param_cp.device:
28
        temp = param.data.to(param_cp.device)
29
    else:
30
        temp = param.data
31
    return torch.equal(temp, param_cp.data)
32

33

34
@parameterize("init_device", [None, torch.device("cpu")])
35
@parameterize("keep_gathered", [True, False])
36
@parameterize("pin_memory", [True, False])
37
def exam_chunk_basic(init_device, keep_gathered, pin_memory):
38
    world_size = torch.distributed.get_world_size()
39
    pg = _get_default_group()
40
    my_chunk = Chunk(
41
        chunk_size=1024,
42
        zero_group=pg,
43
        dtype=torch.float32,
44
        init_device=init_device,
45
        cpu_shard_init=True,
46
        keep_gathered=keep_gathered,
47
        pin_memory=pin_memory,
48
    )
49

50
    param_list = []
51
    param_cp_list = []
52

53
    add_param(param_list, param_cp_list, 8, 8, 8, device="cuda")
54
    add_param(param_list, param_cp_list, 4, 4)
55
    add_param(param_list, param_cp_list, 4, 8, 2, device="cuda")
56
    add_param(param_list, param_cp_list, 1, 1, 5)
57

58
    for param in param_list:
59
        my_chunk.append_tensor(param)
60
    assert my_chunk.utilized_size == 597
61
    for param, param_cp in zip(param_list, param_cp_list):
62
        check_equal(param, param_cp)
63
    my_chunk.close_chunk()
64

65
    if keep_gathered is False:
66
        assert my_chunk.cpu_shard.size(0) == 1024 // world_size
67
        assert my_chunk.device_type == "cpu"
68
        assert my_chunk.can_move
69
        my_chunk.shard_move(get_accelerator().get_current_device())
70
    else:
71
        assert my_chunk.cuda_global_chunk.size(0) == 1024
72
        assert my_chunk.device_type == "cuda"
73
        assert not my_chunk.can_move
74

75
    assert dist_sum(my_chunk.valid_end) == my_chunk.utilized_size
76
    flag = my_chunk.has_inf_or_nan
77
    assert not flag, "has_inf_or_nan is {}".format(flag)
78

79
    my_chunk.access_chunk()
80
    assert my_chunk.device_type == "cuda"
81
    for param, param_cp in zip(param_list, param_cp_list):
82
        check_equal(param, param_cp)
83

84
    assert my_chunk.tensor_state_cnter[TensorState.HOLD] == 4
85
    my_chunk.tensor_trans_state(param_list[0], TensorState.COMPUTE)
86
    assert my_chunk.tensor_state_cnter[TensorState.HOLD] == 3
87
    assert my_chunk.tensor_state_cnter[TensorState.COMPUTE] == 1
88
    assert not my_chunk.can_release
89

90
    for param in param_list:
91
        my_chunk.tensor_trans_state(param, TensorState.COMPUTE)
92
        my_chunk.tensor_trans_state(param, TensorState.HOLD_AFTER_BWD)
93
        my_chunk.tensor_trans_state(param, TensorState.READY_FOR_REDUCE)
94

95
    assert my_chunk.tensor_state_cnter[TensorState.READY_FOR_REDUCE] == 4
96
    assert my_chunk.can_reduce
97
    my_chunk.reduce()
98
    assert my_chunk.tensor_state_cnter[TensorState.HOLD] == 4
99

100
    if keep_gathered is False:
101
        assert my_chunk.cuda_shard.size(0) == 1024 // world_size
102
        assert my_chunk.device_type == "cuda"
103
        assert my_chunk.can_move
104
    else:
105
        assert my_chunk.cuda_global_chunk.size(0) == 1024
106
        assert my_chunk.device_type == "cuda"
107
        assert not my_chunk.can_move
108

109

110
def run_dist(rank, world_size, port):
111
    colossalai.launch(config={}, rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl")
112
    exam_chunk_basic()
113

114

115
@pytest.mark.dist
116
@pytest.mark.parametrize("world_size", [1, 2, 4])
117
@rerun_if_address_is_in_use()
118
def test_chunk_function(world_size):
119
    spawn(run_dist, world_size)
120

121

122
if __name__ == "__main__":
123
    test_chunk_function(4)
124