wehub-resource-sync
220 lines
8.0 KiB
Python
220 lines
8.0 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import torch
|
|
import numpy as np
|
|
import tempfile
|
|
import torch.distributed as dist
|
|
from pathlib import Path
|
|
from typing import Callable, Optional, Union
|
|
|
|
|
|
def flush_l2_cache(enabled: bool = True):
|
|
"""
|
|
Flush the GPU L2 cache by writing a large zero-initialized tensor.
|
|
|
|
Arguments:
|
|
enabled: if `False`, does nothing.
|
|
"""
|
|
l2_flush_cache_size = 256e6
|
|
if enabled:
|
|
torch.empty(int(l2_flush_cache_size // 4), dtype=torch.int, device='cuda').zero_()
|
|
|
|
|
|
def bench(fn, num_warmups: int = 50, num_tests: int = 50,
|
|
post_fn: Optional[Callable] = None, flush_l2: bool = True):
|
|
"""
|
|
Benchmark a function using CUDA events.
|
|
|
|
Arguments:
|
|
fn: the function to benchmark.
|
|
num_warmups: the number of warmup iterations.
|
|
num_tests: the number of measurement iterations.
|
|
post_fn: an optional function to call after each test iteration.
|
|
flush_l2: whether to flush the L2 cache before each iteration.
|
|
|
|
Returns:
|
|
avg: the average execution time in seconds.
|
|
min: the minimum execution time in seconds.
|
|
max: the maximum execution time in seconds.
|
|
"""
|
|
torch.cuda.synchronize()
|
|
|
|
# Warmup
|
|
for _ in range(num_warmups):
|
|
fn()
|
|
|
|
# Testing
|
|
start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
|
|
end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
|
|
for i in range(num_tests):
|
|
flush_l2_cache(flush_l2)
|
|
start_events[i].record()
|
|
fn()
|
|
end_events[i].record()
|
|
if post_fn is not None:
|
|
post_fn()
|
|
torch.cuda.synchronize()
|
|
|
|
times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)])[1:]
|
|
return np.average(times), np.min(times), np.max(times)
|
|
|
|
|
|
class empty_suppress:
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_):
|
|
pass
|
|
|
|
|
|
class suppress_stdout_stderr:
|
|
"""
|
|
Context manager to suppress stdout and stderr output.
|
|
"""
|
|
|
|
def __enter__(self):
|
|
self.outnull_file = open(os.devnull, 'w')
|
|
self.errnull_file = open(os.devnull, 'w')
|
|
|
|
self.old_stdout_fileno_undup = sys.stdout.fileno()
|
|
self.old_stderr_fileno_undup = sys.stderr.fileno()
|
|
|
|
self.old_stdout_fileno = os.dup(sys.stdout.fileno())
|
|
self.old_stderr_fileno = os.dup(sys.stderr.fileno())
|
|
|
|
self.old_stdout = sys.stdout
|
|
self.old_stderr = sys.stderr
|
|
|
|
os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup)
|
|
os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
|
|
|
|
sys.stdout = self.outnull_file
|
|
sys.stderr = self.errnull_file
|
|
return self
|
|
|
|
def __exit__(self, *_):
|
|
sys.stdout = self.old_stdout
|
|
sys.stderr = self.old_stderr
|
|
|
|
os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup)
|
|
os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)
|
|
|
|
os.close(self.old_stdout_fileno)
|
|
os.close(self.old_stderr_fileno)
|
|
|
|
self.outnull_file.close()
|
|
self.errnull_file.close()
|
|
|
|
|
|
def bench_kineto(fn,
|
|
kernel_names: Union[str, tuple],
|
|
num_tests: int = 30,
|
|
suppress_kineto_output: bool = False,
|
|
trace_path: Optional[str] = None,
|
|
flush_l2: bool = True,
|
|
barrier_comm_profiling: bool = False,
|
|
num_kernels_per_period: int = 1,
|
|
barrier: Optional[Callable] = None):
|
|
"""
|
|
Benchmark a function using the PyTorch profiler (kineto) to get per-kernel timing.
|
|
|
|
Arguments:
|
|
fn: the function to benchmark.
|
|
kernel_names: the CUDA kernel name(s) to profile.
|
|
num_tests: the number of test iterations.
|
|
suppress_kineto_output: whether to suppress profiler output.
|
|
trace_path: the path to save the Chrome trace (`None` to skip).
|
|
flush_l2: whether to flush the L2 cache before each iteration.
|
|
barrier_comm_profiling: whether to insert a barrier before each iteration to reduce
|
|
unbalanced CPU launch overhead.
|
|
num_kernels_per_period: the number of kernels launched per test period.
|
|
barrier: a custom barrier function to use instead of `dist.all_reduce`.
|
|
|
|
Returns:
|
|
durations: the average kernel duration(s) in seconds.
|
|
"""
|
|
assert isinstance(kernel_names, (str, tuple))
|
|
is_tuple = isinstance(kernel_names, tuple)
|
|
|
|
# Skip profiling
|
|
# Conflict with Nsight Systems, Nsight Compute and Compute Sanitizer
|
|
if int(os.environ.get('EP_USE_NVIDIA_TOOLS', 0)):
|
|
return (1, ) * len(kernel_names) if is_tuple else 1
|
|
|
|
# For some auto-tuning kernels with prints
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
|
|
# Profile
|
|
suppress = suppress_stdout_stderr if suppress_kineto_output else empty_suppress
|
|
barrier_comm_profiling &= int(os.environ.get('EP_DISABLE_BARRIER_PROFILING', 0)) == 0
|
|
with suppress():
|
|
schedule = torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1)
|
|
profiler = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA], schedule=schedule, acc_events=True)
|
|
dummy = torch.ones(1, dtype=torch.float, device='cuda')
|
|
with profiler:
|
|
for i in range(2):
|
|
for _ in range(num_tests):
|
|
# Flush L2 cache
|
|
flush_l2_cache(flush_l2)
|
|
|
|
# NOTES: use a large kernel and a barrier to eliminate the unbalanced CPU launch overhead
|
|
if barrier_comm_profiling:
|
|
torch.cuda._sleep(int(2e7)) # ~10ms
|
|
|
|
# Some network may have ring-based implement, so be careful to use `all_reduce`
|
|
if barrier is None:
|
|
dist.all_reduce(dummy)
|
|
else:
|
|
barrier()
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
profiler.step()
|
|
|
|
# Parse the profiling table
|
|
prof_lines = profiler.key_averages().table(sort_by='cuda_time_total', max_name_column_width=100).split('\n')
|
|
kernel_names = (kernel_names, ) if isinstance(kernel_names, str) else kernel_names
|
|
assert all([isinstance(name, str) for name in kernel_names])
|
|
for name in kernel_names:
|
|
assert sum([name in line for line in prof_lines]) <= 1, f'Errors of the kernel {name} in the profiling table: {prof_lines}'
|
|
|
|
# Save chrome traces
|
|
if trace_path is not None:
|
|
profiler.export_chrome_trace(trace_path)
|
|
|
|
# Return average kernel durations
|
|
units = {'ms': 1e3, 'us': 1e6}
|
|
kernel_durations = []
|
|
for name in kernel_names:
|
|
total_time = 0
|
|
total_num = 0
|
|
for line in prof_lines:
|
|
if name in line:
|
|
time_str = line.split()[-2]
|
|
num_str = line.split()[-1]
|
|
for unit, scale in units.items():
|
|
if unit in time_str:
|
|
total_time += float(time_str.replace(unit, '')) / scale * int(num_str)
|
|
total_num += int(num_str)
|
|
break
|
|
kernel_durations.append(total_time / total_num if total_num > 0 else 0)
|
|
|
|
# Expand the kernels by periods
|
|
if num_kernels_per_period > 1:
|
|
with tempfile.NamedTemporaryFile(suffix='.json') as tmp:
|
|
profiler.export_chrome_trace(tmp.name)
|
|
profile_data = json.loads(Path(tmp.name).read_text())
|
|
|
|
for i, kernel_name in enumerate(kernel_names):
|
|
events = [event for event in profile_data['traceEvents'] if f'::{kernel_name}' in event['name']]
|
|
events = sorted(events, key=lambda event: event['ts'])
|
|
durations = [event['dur'] / 1e6 for event in events]
|
|
assert len(durations) % num_kernels_per_period == 0
|
|
num_kernel_patterns = len(durations) // num_kernels_per_period
|
|
kernel_durations[i] = [sum(durations[j::num_kernels_per_period]) / num_kernel_patterns for j in range(num_kernels_per_period)]
|
|
|
|
# Return execution durations
|
|
return kernel_durations if is_tuple else kernel_durations[0]
|