test_naive_fp16.py 3.38 KB
Newer Older
1
import torch
2
import colossalai
3
4
5
import torch.multiprocessing as mp
from colossalai.amp import convert_to_naive_amp, convert_to_apex_amp
from tests.components_to_test.registry import non_distributed_component_funcs
6
from colossalai.testing import assert_close_loose, rerun_if_address_is_in_use
7
from colossalai.utils import free_port
8
9
10
11
12
13
from colossalai.amp import convert_to_naive_amp, convert_to_apex_amp

from tests.components_to_test.registry import non_distributed_component_funcs

import copy
import pytest
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
from functools import partial


def check_equal(a, b):
    """
    This function checks if two tensors are equal within tolerance
    """
    assert torch.allclose(a.float(), b.float(), rtol=1e-4, atol=1e-3), f'a = {a}, b = {b}'


def run_naive_amp():
    """
    In this test, we compare the naive fp16 optimizer implemented in colossalai 
    and fp32 torch optimizer
    """

30
31
32
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

33
    # create layer
34
    test_models = ['repeated_computed_layers', 'nested_model', 'resnet18']
35
36
    for test_name in test_models:
        get_component_func = non_distributed_component_funcs.get_callable(test_name)
37
        model_builder, train_dataloader, _, optim_class, _ = get_component_func()
38
39

        # create model
40
41
        naive_amp_model = model_builder(checkpoint=True).cuda()
        apex_amp_model = copy.deepcopy(naive_amp_model)
42
43

        # create optimizer
44
45
        naive_amp_optimizer = optim_class(naive_amp_model.parameters(), lr=1e-3)
        apex_amp_optimizer = optim_class(apex_amp_model.parameters(), lr=1e-3)
46

47
48
49
50
51
52
        # inject naive and apex amp
        naive_amp_config = dict(initial_scale=128)
        naive_amp_model, naive_amp_optimizer = convert_to_naive_amp(naive_amp_model, naive_amp_optimizer,
                                                                    naive_amp_config)
        apex_amp_config = dict(opt_level='O2', loss_scale=128, keep_batchnorm_fp32=False)
        apex_amp_model, apex_amp_optimizer = convert_to_apex_amp(apex_amp_model, apex_amp_optimizer, apex_amp_config)
53
54
55
56
57
58
59

        # create data
        data_iter = iter(train_dataloader)
        data, label = next(data_iter)
        data = data.cuda()

        # forward pass
60
61
62
        naive_amp_output = naive_amp_model(data)
        apex_amp_output = apex_amp_model(data)
        assert_close_loose(naive_amp_output, apex_amp_output)
63
64

        # backward
65
66
        naive_amp_optimizer.backward(naive_amp_output.mean())
        apex_amp_optimizer.backward(apex_amp_output.mean())
67
68

        # check grad
69
70
        for naive_amp_param, apex_amp_param in zip(naive_amp_model.parameters(), apex_amp_model.parameters()):
            assert_close_loose(naive_amp_param.grad, apex_amp_param.grad)
71
72

        # step
73
74
        naive_amp_optimizer.step()
        apex_amp_optimizer.step()
75
76

        # check updated param
77
78
        for naive_amp_param, apex_amp_param in zip(naive_amp_model.parameters(), apex_amp_model.parameters()):
            assert_close_loose(naive_amp_param, apex_amp_param)
79
80
81
82
83
84
85
86


def run_dist(rank, world_size, port):
    colossalai.launch(config=dict(), rank=rank, world_size=world_size, port=port, host='localhost')
    run_naive_amp()


@pytest.mark.dist
87
@rerun_if_address_is_in_use()
88
89
90
91
92
93
94
95
def test_naive_amp():
    world_size = 1
    run_func = partial(run_dist, world_size=world_size, port=free_port())
    mp.spawn(run_func, nprocs=world_size)


if __name__ == '__main__':
    test_naive_amp()