Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
layernorm_backward.py 3.73 KB
Newer Older
yanjl1's avatar
Initial  
yanjl1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114 
import hipdnn
import torch


def build_layernorm_backward_graph(
    hipdnn_handle,
    torch_tensor_x,
    torch_tensor_dy,
    torch_tensor_scale,
    torch_tensor_mean,
    torch_tensor_inv_var,
    hipdnn_data_type,
):
    # Create graph
    graph = hipdnn.pygraph(
        handle=hipdnn_handle,
        io_data_type=hipdnn_data_type,
        intermediate_data_type=hipdnn.data_type.FLOAT,
        compute_data_type=hipdnn.data_type.FLOAT,
        name="layernorm_backward",
    )

    # Create hipdnn tensors
    hipdnn_tensor_x = graph.tensor_like(torch_tensor_x)
    hipdnn_tensor_dy = graph.tensor_like(torch_tensor_dy)
    hipdnn_tensor_scale = graph.tensor_like(torch_tensor_scale)
    hipdnn_tensor_mean = graph.tensor_like(torch_tensor_mean)
    hipdnn_tensor_inv_var = graph.tensor_like(torch_tensor_inv_var)

    # Create layernorm op
    hipdnn_tensor_dx, hipdnn_tensor_dbias, hipdnn_tensor_dscale = graph.layernorm_backward(
        grad=hipdnn_tensor_dy,
        input=hipdnn_tensor_x,
        scale=hipdnn_tensor_scale,
        mean=hipdnn_tensor_mean,
        inv_variance=hipdnn_tensor_inv_var,
        name="layernorm_backward",
    )
    hipdnn_tensor_dx.set_output(True)
    hipdnn_tensor_dbias.set_output(True)
    hipdnn_tensor_dscale.set_output(True)
    graph.build(hipdnn_handle)

    return (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_dy,
        hipdnn_tensor_scale,
        hipdnn_tensor_mean,
        hipdnn_tensor_inv_var,
        hipdnn_tensor_dx,
        hipdnn_tensor_dbias,
        hipdnn_tensor_dscale,
    )


if __name__ == "__main__":
    # Input dimensions
    batch, seq_len, embedding_dim = 2, 1024, 768

    hipdnn_data_type = hipdnn.data_type.FLOAT
    torch_data_type = torch.float32

    torch_tensor_x = torch.rand(batch, seq_len, embedding_dim, dtype=torch_data_type, device="cuda")
    torch_tensor_dy = torch.rand(
        batch, seq_len, embedding_dim, dtype=torch_data_type, device="cuda"
    )
    torch_tensor_scale = torch.rand(embedding_dim, dtype=torch_data_type, device="cuda")
    torch_tensor_mean = torch.rand(batch * seq_len, dtype=torch_data_type, device="cuda")
    torch_tensor_inv_var = torch.rand(batch * seq_len, dtype=torch_data_type, device="cuda")

    hipdnn_handle = hipdnn.create_handle()

    (
        graph,
        hipdnn_tensor_x,
        hipdnn_tensor_dy,
        hipdnn_tensor_scale,
        hipdnn_tensor_mean,
        hipdnn_tensor_inv_var,
        hipdnn_tensor_dx,
        hipdnn_tensor_dbias,
        hipdnn_tensor_dscale,
    ) = build_layernorm_backward_graph(
        hipdnn_handle,
        torch_tensor_x,
        torch_tensor_dy,
        torch_tensor_scale,
        torch_tensor_mean,
        torch_tensor_inv_var,
        hipdnn_data_type,
    )

    torch_tensor_dx = torch.empty(hipdnn_tensor_dx.get_dim(), dtype=torch_data_type, device="cuda")
    torch_tensor_dbias = torch.empty(
        hipdnn_tensor_dbias.get_dim(), dtype=torch_data_type, device="cuda"
    )
    torch_tensor_dscale = torch.empty(
        hipdnn_tensor_dscale.get_dim(), dtype=torch_data_type, device="cuda"
    )
    variant_pack = {
        hipdnn_tensor_x: torch_tensor_x.data_ptr(),
        hipdnn_tensor_dy: torch_tensor_dy.data_ptr(),
        hipdnn_tensor_scale: torch_tensor_scale.data_ptr(),
        hipdnn_tensor_mean: torch_tensor_mean.data_ptr(),
        hipdnn_tensor_inv_var: torch_tensor_inv_var.data_ptr(),
        hipdnn_tensor_dx: torch_tensor_dx.data_ptr(),
        hipdnn_tensor_dbias: torch_tensor_dbias.data_ptr(),
        hipdnn_tensor_dscale: torch_tensor_dscale.data_ptr(),
    }
    workspace = torch.empty(graph.get_workspace_size(), dtype=torch.uint8, device="cuda")

    graph.exec(variant_pack=variant_pack, workspace=workspace.data_ptr())
    print("layernorm backward graph execution complete.")