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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
|
import torch
from torch.profiler import profile, record_function, ProfilerActivity
from torch.utils.benchmark import Timer
import time
def profile_cuda_kernels(fn, args, string_id="Model time"):
print("################################################")
print(f"#### Profiling for {string_id} starts #########")
print("################################################")
warmup = 50
old_args = args[:]
n_repeats = 1
n_layers = 1
ref = fn(*old_args)
gO = torch.rand_like(ref)
for _ in range(0, warmup // n_layers):
args = list(old_args[:])
ref = fn(*args)
ref.backward(gO)
torch.cuda.synchronize()
# Forward profile
def fwd_run():
for _ in range(0, n_repeats // n_layers):
args = list(old_args[:])
for arg in args:
if isinstance(arg, torch.Tensor):
arg.grad = None
ref = fn(*args)
print(f"###### Forward profile for {string_id} starts #####")
with profile(activities=[ProfilerActivity.CUDA], record_shapes=True) as prof:
with record_function("baseline"):
fwd_run()
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=30))
print(f"###### Forward profile for {string_id} ends #####")
# Backward profile
def bwd_run():
for _ in range(0, n_repeats // n_layers):
args = list(old_args[:])
for arg in args:
if isinstance(arg, torch.Tensor):
arg.grad = None
ref = fn(*args)
print(f"###### Backward profile for {string_id} starts #####")
torch.cuda.synchronize()
with profile(
activities=[ProfilerActivity.CUDA], record_shapes=True
) as prof:
with record_function("baseline"):
ref.backward(gO)
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=30))
torch.cuda.synchronize()
print(f"###### Backward profile for {string_id} ends #####")
bwd_run()
print("################################################")
print(f"#### Profiling for {string_id} ends #########")
print("################################################\n\n\n\n")
def time_with_torch_timer(fn, args, string_id, kwargs=None):
if kwargs is None:
kwargs = {}
print("################################################")
print(f"#### Torch Timer for {string_id} starts #########")
print("################################################")
ref = fn(*args, **kwargs)
gO = torch.rand_like(ref)
env = {"args": args, "gO": gO, "kwargs": kwargs, "fn": fn}
grad_none = {"for x in args: x.grad=None"}
fn_call = "fn(*args, **kwargs)"
# Measure end-to-end fwd time
timer = Timer(stmt=f"{fn_call}", globals=env)
fwd_latency = round(timer.timeit(1000).mean * 10 ** 6, 3)
timer_blocked = timer.blocked_autorange()
print(f"Forward = {fwd_latency}")
# Measure end-to-end fwd bwd
timer = Timer(
stmt=f"{grad_none}; fwd = {fn_call}; fwd.backward(gO)",
globals=env,
)
fwd_bwd_latency = round(timer.timeit(1000).mean * 10 ** 6, 3)
timer_blocked = timer.blocked_autorange()
# print(f"Forward + sum + Backward = {fwd_sum_bwd_latency}")
bwd_latency = round(fwd_bwd_latency - fwd_latency, 3)
print(f"Backward = {bwd_latency}")
print("################################################")
print(f"#### Torch Timer for {string_id} ends ###############")
print("################################################\n\n\n\n")
def time_with_manual_timer(fn, args, string_id):
print("################################################")
print(f"#### Manual Timer for {string_id} starts #########")
print("################################################")
warmup = 50
repeats = 1000
old_args = args[:]
ref = fn(*old_args)
gO = torch.rand_like(ref)
for _ in range(0, warmup):
args = list(old_args[:])
for arg in args:
if isinstance(arg, torch.Tensor):
arg.grad = None
ref = fn(*args)
ref.backward(gO)
torch.cuda.synchronize()
fwd_times = []
bwd_times = []
for _ in range(0, repeats):
args = list(old_args[:])
for arg in args:
if isinstance(arg, torch.Tensor):
arg.grad = None
fwd_start = time.time()
ref = fn(*args)
torch.cuda.synchronize()
fwd_end = time.time()
bwd_start = time.time()
ref.backward(gO)
torch.cuda.synchronize()
bwd_end = time.time()
fwd_times.append(fwd_end - fwd_start)
bwd_times.append(bwd_end - bwd_start)
avg_fwd = round(sum(fwd_times) / repeats * 10 ** 6, 2)
avg_bwd = round(sum(bwd_times) / repeats * 10 ** 6, 2)
avg_total = round(avg_fwd + avg_bwd, 2)
print(f"Forward = {avg_fwd}")
print(f"Backward = {avg_bwd}")
print("################################################")
print(f"#### Manual Timer for {string_id} ends #########")
print("################################################\n\n\n")
|
