common.py 4.89 KB
Newer Older
Jiarui Fang's avatar
Jiarui Fang committed
1
from functools import partial
2

Jiarui Fang's avatar
Jiarui Fang committed
3
import torch
4
import torch.distributed as dist
5

ver217's avatar
ver217 committed
6
from colossalai.logging import get_dist_logger
7
from colossalai.utils import checkpoint
8
from colossalai.zero.sharded_model import ShardedModelV2
Jiarui Fang's avatar
Jiarui Fang committed
9

10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
LOGGER = get_dist_logger('zero_test')

MP_PARALLEL_CONFIG = dict(fp16=dict(mode=None,), parallel=dict(pipeline=dict(size=1), tensor=dict(size=2, mode=None)))

_ZERO_MODEL_CONFIG = dict(reduce_scatter_bucket_size_mb=25,
                          fp32_reduce_scatter=False,
                          offload_config=None,
                          gradient_predivide_factor=1.0,
                          shard_param=True,
                          use_memory_tracer=False)

_ZERO_OPTIMIZER_CONFIG = dict(
    optimizer_class=torch.optim.Adam,
    cpu_offload=False,
    initial_scale=2**32,
    min_scale=1,
    growth_factor=2,
    backoff_factor=0.5,
    growth_interval=1000,
    hysteresis=2,
    max_scale=2**32,
)
32
33
34

ZERO_PARALLEL_CONFIG = dict(fp16=dict(mode=None,),
                            zero=dict(
35
36
                                model_config=_ZERO_MODEL_CONFIG,
                                optimizer_config=_ZERO_OPTIMIZER_CONFIG,
37
38
39
                            ),
                            parallel=dict(pipeline=dict(size=1), tensor=dict(size=1, mode=None)))

ver217's avatar
ver217 committed
40
41
42
43
44
45
46
47
48
49
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)))
Jiarui Fang's avatar
Jiarui Fang committed
50

51

52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def run_fwd_bwd(model, data, label, criterion, enable_autocast=False):
    model.train()
    with torch.cuda.amp.autocast(enabled=enable_autocast):
        if criterion:
            y = model(data)
            loss = criterion(y, label)
        else:
            loss = model(data, label)
        loss = loss.float()
    if isinstance(model, ShardedModelV2):
        model.backward(loss)
    else:
        loss.backward()


Jiarui Fang's avatar
Jiarui Fang committed
67
68
69
70
71
72
73
74
def checkpoint_wrapper(module, enable=True):
    if enable:
        module.forward = partial(checkpoint, module.forward)
    return module


def allclose(tensor_a: torch.Tensor, tensor_b: torch.Tensor, loose=False) -> bool:
    if loose:
75
        return torch.allclose(tensor_a, tensor_b, atol=1e-2, rtol=1e-3)
Jiarui Fang's avatar
Jiarui Fang committed
76
77
78
79
80
81
    return torch.allclose(tensor_a, tensor_b)


def check_grads(model, zero_model, loose=False):
    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
        zero_grad = zero_p.grad.clone().to(p.device)
82
83
84
        grad = p.grad.float()
        assert grad.dtype == zero_grad.dtype
        assert allclose(grad, zero_grad, loose=loose)
85

Jiarui Fang's avatar
Jiarui Fang committed
86
87
88
89

def check_params(model, zero_model, loose=False):
    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
        zero_p = zero_p.clone().to(p.device)
90
91
        # assert p.dtype == zero_p.dtype
        assert allclose(p.float(), zero_p.float(), loose=loose), f"diff {p.float() - zero_p.float()}"
Jiarui Fang's avatar
Jiarui Fang committed
92

93
94
95
96
97
98
99
100

def check_grads_padding(model, zero_model, loose=False):
    rank = dist.get_rank()
    for p, zero_p in zip(model.parameters(), zero_model.parameters()):
        zero_grad = zero_p.grad.clone().to(p.device)
        chunks = torch.flatten(p.grad).chunk(dist.get_world_size())
        if rank >= len(chunks):
            continue
101
        grad = chunks[rank].float()
102
103
104
        if zero_grad.size(0) > grad.size(0):
            zero_grad = zero_grad[:grad.size(0)]
        assert grad.dtype == zero_grad.dtype
105
        assert allclose(grad, zero_grad, loose=loose), f'diff: {grad - zero_grad}'
106
107
108
109
110
111
112
113
114
115
116
117
118
119


def check_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.clone().to(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)
ver217's avatar
ver217 committed
120
121
122
123
124


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()):
ver217's avatar
ver217 committed
125
        zero_p = zero_p.col_attr.data.payload.to(p.device).float()
ver217's avatar
ver217 committed
126
127
128
        chunks = torch.flatten(p).chunk(dist.get_world_size())
        if rank >= len(chunks):
            continue
ver217's avatar
ver217 committed
129
        p = chunks[rank].float()
ver217's avatar
ver217 committed
130
131
132
        if zero_p.size(0) > p.size(0):
            zero_p = zero_p[:p.size(0)]
        assert p.dtype == zero_p.dtype
ver217's avatar
ver217 committed
133
        assert allclose(p, zero_p, loose=loose), f'{p} vs {zero_p}'