".github/vscode:/vscode.git/clone" did not exist on "7bc5a8e338a1240070ca82f96721c451060f1a70"
test_3d.py 1.92 KB
Newer Older
zbian's avatar
zbian committed
1
2
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
Frank Lee's avatar
Frank Lee committed
3
4
5
import pytest
import torch
import torch.multiprocessing as mp
Frank Lee's avatar
Frank Lee committed
6
from colossalai.initialize import launch, get_default_parser
zbian's avatar
zbian committed
7

Frank Lee's avatar
Frank Lee committed
8
9
from checks_3d.check_layer_3d import *
from checks_3d.check_operation_3d import *
Frank Lee's avatar
Frank Lee committed
10
from colossalai.logging import get_dist_logger
Frank Lee's avatar
Frank Lee committed
11
from functools import partial
zbian's avatar
zbian committed
12
13
14
15
16

CONFIG = dict(parallel=dict(pipeline=1, tensor=dict(mode='3d', size=8)),
              seed=0)


Frank Lee's avatar
Frank Lee committed
17
18
19
20
21
22
23
# def check_operations():
#     check_AB()
#     check_ABT()
#     check_ATB()
#     check_add()
#     check_mul()
#     check_sum()
zbian's avatar
zbian committed
24
25
26


def check_layer():
Frank Lee's avatar
Frank Lee committed
27
    logger = get_dist_logger()
zbian's avatar
zbian committed
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
    liear_fwd_time, linear_bwd_time = check_linear()
    norm_fwd_time, norm_bwd_time = check_layernorm()
    attn_fwd_time, attn_bwd_time = check_attention()
    mlp_fwd_time, mlp_bwd_time = check_mlp()
    head_fwd_time, head_bwd_time = check_head()
    embed_fwd_time, embed_bwd_time = check_embed()
    loss_fwd_time, loss_bwd_time = check_loss()
    block_fwd_time = norm_fwd_time + attn_fwd_time + norm_fwd_time + mlp_fwd_time
    block_bwd_time = norm_bwd_time + attn_bwd_time + norm_bwd_time + mlp_bwd_time
    fwd_time = embed_fwd_time + NUM_BLOCKS * block_fwd_time + norm_fwd_time + head_fwd_time + loss_fwd_time
    bwd_time = embed_bwd_time + NUM_BLOCKS * block_bwd_time + norm_bwd_time + head_bwd_time + loss_bwd_time
    logger.info('ViT forward time: {:.3f} s | backward time: {:.3f} s'.format(
        fwd_time, bwd_time),
        ranks=[0])


Frank Lee's avatar
Frank Lee committed
44
def check_layer_and_operation(rank, world_size):
Frank Lee's avatar
Frank Lee committed
45
    launch(config=CONFIG,
Frank Lee's avatar
Frank Lee committed
46
47
48
49
50
           rank=rank,
           world_size=world_size,
           host='localhost',
           port=29923,
           backend='nccl')
zbian's avatar
zbian committed
51
52

    check_layer()
Frank Lee's avatar
Frank Lee committed
53
54
55
56
57
58
59
60
61
    gpc.destroy()
    torch.cuda.empty_cache()


@pytest.mark.dist
def test_3d():
    world_size = 8
    run_func = partial(check_layer_and_operation, world_size=world_size)
    mp.spawn(run_func, nprocs=world_size)
zbian's avatar
zbian committed
62
63
64


if __name__ == '__main__':
Frank Lee's avatar
Frank Lee committed
65
    test_3d()