Що дійсно визначає швидкість вашого PyTorch-коду?

Кожен, хто працює з модельним кодом PyTorch, починає задавати ті ж самі питання: Чому це займає так довго?Як я можу зробити свій тренувальний цикл швидше? Чому це займає так довго?Як я можу зробити свій тренувальний цикл швидше? Незалежно від того, чи є ви інженером ML, дослідником або просто вирішили пограти з випадковим сховищем ML протягом вихідних, ви в кінцевому підсумку спробуєте зрозуміти, як прискорити свій код. Однак, перш ніж ми зможемо зробити це, ми повинні навчитися правильно вимірювати продуктивність, а потім зробити правильні висновки з цих вимірювань. Виконання потоку Давайте використаємо множення матриці на H100 як приклад. Однак, перш ніж ми думаємо про вимірювання, давайте отримаємо термінологію прямо. Під «ядро» ми маємо на увазі функцію (або насправді будь-яку операцію), яка працює на GPU. Ваш код може мати багато з цих, від десятків до тисяч операцій під час кожної ітерації циклу навчання, наприклад: torch.zeros((16, 16), пристрій=”cuda) a + b, де a і b є тенсорами в пам'яті GPU А @ Б flash_attention (q, k, v) в якості І так на Незважаючи на те, що ядро виконується на GPU, його Якщо ми уявляємо собі файл Python з кодом навчання, він змінюється між «підготовими» роботами, такими як «for-loops» або «if-statements», і фактичними запусками ядра. Через це процесор йде вперед і кожен раз, коли він потрапляє в ядро, він планує те ядро для виконання на GPU, додавши його в чергу. запуску Це дозволяє GPU уникнути бездіяльності: він завжди може знайти роботу в цій черзі ядра (за винятком випадків, коли синхронізація явно викликається або коли виконання пов'язане з CPU). Наївний підхід і чому він помиляється Коли ми спочатку думаємо про вимірювання часу в Python, природний інстинкт полягає в тому, щоб досягти модуля часу і запустити щось подібне: import time import torch def matmul(a, b): return a @ b def get_data(m, n, k): a = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) return a, b def benchmark_naive(m, n, k, n_iters: int = 100): a, b = get_data(m, n, k) times = [] for _ in range(n_iters): start = time.perf_counter_ns() matmul(a, b) end = time.perf_counter_ns() times.append(end - start) # perf_counter_ns returns nanoseconds, convert to ms return sum(times) / n_iters / 1e6 small_shapes_time = benchmark_naive(16, 32, 16) large_shapes_time = benchmark_naive(4096, 8192, 4096) print(f"(16, 32) by (32, 16) matmul time: {small_shapes_time:.5f} ms") print(f"(4096, 8192) by (8192, 4096) matmul time: {large_shapes_time:.5f} ms") # (16, 32) by (32, 16) matmul time: 0.01415 ms # (4096, 8192) by (8192, 4096) matmul time: 0.01350 ms Якщо ми запустимо цей код на двох різних наборах форм: та Насправді, в моєму бігу, насправді швидше множити матриці, які в 256 разів більші. 16x32 4096x8192 Насправді, ми ніколи не чекаємо, поки GPU matmul закінчиться. Оскільки весь «підготувальний» код працює на CPU, Це насправді вимірювання часу, який потрібно нам, щоб запланувати ядро matmul в чергу GPU і рухатися далі з нашими життями. time.perf_counter_ns() Те, що ми можемо зробити, це додати після matmul змусити процесор чекати на GPU. Але це все ще не ідеально, тому що ми вимірюємо минулий час на CPU, який включає в себе планування над головою. ми не ізолюємо фактичне виконання GPU таким чином. torch.cuda.synchronize() Вимірювання подій CUDA Правильний спосіб вимірювання часу GPU замість часу CPU - це з подіями CUDA. Це маркери, які ми можемо вставити безпосередньо в потоки виконання GPU. GPU записує часові штампи, коли він досягає кожного події CUDA, даючи нам фактичний час виконання на пристрої. import torch def matmul(a, b): return a @ b def get_data(m, n, k): a = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) return a, b def benchmark_cuda_events(m, n, k, n_iters: int = 100): a, b = get_data(m, n, k) start_events = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] end_events = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] for it in range(n_iters): start_events[it].record() matmul(a, b) end_events[it].record() torch.cuda.synchronize() times = [start.elapsed_time(end) for start, end in zip(start_events, end_events)] return sum(times) / n_iters small_shapes_time = benchmark_cuda_events(16, 32, 16) large_shapes_time = benchmark_cuda_events(4096, 8192, 4096) print(f"(16, 32) by (32, 16) matmul time: {small_shapes_time:.5f} ms") print(f"(4096, 8192) by (8192, 4096) matmul time: {large_shapes_time:.5f} ms") # (16, 32) by (32, 16) matmul time: 0.02093 ms # (4096, 8192) by (8192, 4096) matmul time: 0.34025 ms Події вставляються в той же потік, що і наше ядро matmul. Підтверджується, що дані записуються в Метод дає нам різницю між початковим і кінцевим часовим штампами. enable_timing=True .elapsed_time Тепер ми отримуємо реальні цифри, не абсурдно швидкі результати з наївного підходу. З цього еталона, матуль з великими формами займає принаймні 15 разів довше. Варто також відзначити, що краще включити ітерації нагрівання на початку, тому що часто є одноразові витрати, включені в перший запуск, такі як збірка JIT. L2 кеш Інша проблема з нашим поточним вимірюванням - відсутність L2 кеш-флюша.Коли ми запускаємо те ж саме ядро неодноразово з тими ж даними (і це саме те, що ми робимо), ці дані залишаються в кеші L2 GPU. У GPU NVIDIA існує ієрархія пам'яті: HBM є найбільшою і найповільнішою, за якою слідують поступово менші і швидші одиниці пам'яті (Cache L2 → Cache L1 або Shared Memory → Registry memory). GPU Generation L2 per GPU HBM per GPU V100 Volta 6MB 32GB A100 Ampere 40MB 80GB H100 Hopper 50MB 80GB B200 Blackwell 126MB 192GB В100 Повернення 6 Мб 32 Гб А100 Ампер 40 Мб 80 Гб H100 Хоппер 50 Мб 80 Гб Б200 Блеквелл 126 Мб 192 Гб Через цей кеш, наші вимірювання можуть показати кращі результати, ніж те, що ми бачимо у виробництві, де дані змінюються між ітераціями. Щоб отримати більш реалістичні вимірювання, нам потрібно вимити кеш. Один спосіб зробити це - виділити великий буфер і оновити його на початку кожної ітерації. import torch def flush_l2_cache(): # On H100, the L2 cache is 50MB, so we allocate something a bit bigger cache_size = 60 * 1024 * 1024 buffer = torch.zeros(cache_size // 4, dtype=torch.float32, device="cuda") buffer += 1 del buffer def matmul(a, b): return a @ b def get_data(m, n, k): a = torch.randn(m, n, device="cuda", dtype=torch.bfloat16) b = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) return a, b def benchmark_with_cache_flush(m, n, k, n_iters: int = 100): a, b = get_data(m, n, k) for _ in range(10): matmul(a, b) # warmup torch.cuda.synchronize() start_events = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] end_events = [torch.cuda.Event(enable_timing=True) for _ in range(n_iters)] for it in range(n_iters): flush_l2_cache() # flush the L2 between iterations start_events[it].record() matmul(a, b) end_events[it].record() torch.cuda.synchronize() times = [start.elapsed_time(end) for start, end in zip(start_events, end_events)] return sum(times) / n_iters small_shapes_time = benchmark_with_cache_flush(16, 32, 16) large_shapes_time = benchmark_with_cache_flush(4096, 8192, 4096) print(f"(16, 32) by (32, 16) matmul time: {small_shapes_time:.5f} ms") print(f"(4096, 8192) by (8192, 4096) matmul time: {large_shapes_time:.5f} ms") # (16, 32) by (32, 16) matmul time: 0.00634 ms # (4096, 8192) by (8192, 4096) matmul time: 0.40575 ms Насправді, була значна різниця в результаті між цим і попереднім, але в реальності перезавантаження коду кілька разів змінює результати трохи, і наша різниця знаходиться в межах очікуваної відмінності - так що для цієї конкретної проблеми кеш не мав значного впливу. Вбудовані рішення Чи потрібно все це пам'ятати, щоб щось порівняти?Відповідь ні. Насправді, є бібліотека під назвою , який спочатку був розроблений OpenAI ще в 2021 році для написання ядер GPU Pythonic способом. Сьогодні він має вбудований модуль тестування з відповідними утилітами бенчмаркінг, які здатні робити все вищесказане самостійно. Трітон # Source: https://github.com/triton-lang/triton/blob/main/python/triton/testing.py#L127C1-L190C64 def do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode="mean"): """ Benchmark the runtime of the provided function. By default, return the median runtime of :code:`fn` along with the 20-th and 80-th performance percentile. :param fn: Function to benchmark :type fn: Callable :param warmup: Warmup time (in ms) :type warmup: int :param rep: Repetition time (in ms) :type rep: int :param grad_to_none: Reset the gradient of the provided tensor to None :type grad_to_none: torch.tensor, optional :param quantiles: Performance percentile to return in addition to the median. :type quantiles: list[float], optional :param return_mode: The statistical measure to return. Options are "min", "max", "mean", "median", or "all". Default is "mean". :type return_mode: str """ assert return_mode in ["min", "max", "mean", "median", "all"] di = runtime.driver.active.get_device_interface() fn() di.synchronize() cache = runtime.driver.active.get_empty_cache_for_benchmark() # Estimate the runtime of the function start_event = di.Event(enable_timing=True) end_event = di.Event(enable_timing=True) start_event.record() for _ in range(5): runtime.driver.active.clear_cache(cache) fn() end_event.record() di.synchronize() estimate_ms = start_event.elapsed_time(end_event) / 5 # compute number of warmup and repeat n_warmup = max(1, int(warmup / estimate_ms)) n_repeat = max(1, int(rep / estimate_ms)) start_event = [di.Event(enable_timing=True) for i in range(n_repeat)] end_event = [di.Event(enable_timing=True) for i in range(n_repeat)] # Warm-up for _ in range(n_warmup): fn() # Benchmark for i in range(n_repeat): # we don't want `fn` to accumulate gradient values # if it contains a backward pass. So we clear the # provided gradients if grad_to_none is not None: for x in grad_to_none: x.grad = None # we clear the L2 cache before each run runtime.driver.active.clear_cache(cache) # record time of `fn` start_event[i].record() fn() end_event[i].record() # Record clocks di.synchronize() times = [s.elapsed_time(e) for s, e in zip(start_event, end_event)] return _summarize_statistics(times, quantiles, return_mode) Крім того, що ми розкрили, Наприклад, можна запитати про медіану або 99-й процентиль. Є інші корисні аргументи, які охоплюються в доктрині. do_bench Для використання цієї функції ми можемо просто зателефонувати: import triton bench_time = triton.testing.do_bench(lambda: matmul(a, b)) Виконавчий процесор і графіки CUDA Для ядер з дуже коротким часом виконання, запуск навантаження може стати настільки великим, що він перевищує самий час ядра. Те ж саме відбувається, коли є значна робота CPU (зазвичай не оптимізована) між запусками ядра. Додамо looong for-loop до нашого ядра matmul, оскільки ми знаємо, що for-loops дійсно повільні на CPU в Python: def cpu_heavy_matmul(a, b): cnt = 0 for _ in range(100_000): cnt += 1 return a @ b Ми побачимо, що (4096, 8192) по (8192, 4096) час матмулу збільшився з 0,4 мс до 2,2 мс - більш ніж в 5 разів! Тепер, залежно від того, як виглядає ваш фактичний тренінг, ви можете або виключити CPU-overhead з результатів бенчмаркування, або ні. Наприклад, якщо ваше навчання вже пов'язане з CPU, то під час виробництва також виникнуть ті ж проблеми. Тоді, швидше за все, не має сенсу позбутися часу CPU для наших вимірювань. Однак, в звичайних тренінгах, процесор не повинен стояти на шляху до оптимальної продуктивності тренінгу - зазвичай процесор виходить перед GPU. У цьому випадку, навіть якщо це одне ядро має значний overhead, це не змінить час роботи ядра на GPU. Тоді ми хотіли б усунути overhead з наших вимірювань. Для цього ми можемо використовувати графіки CUDA. Графіки CUDA дозволяють записувати ряд ядер один раз і відтворювати їх з мінімальним залученням процесора. Використовуючи це для нашого нового matmul, ми отримуємо початковий час ядра (до того, як був введений CPU overhead). triton.testing.do_bench_cudagraph a, b = get_data(4096, 8192, 4096) triton.testing.do_bench_cudagraph(lambda: cpu_heavy_matmul(a, b)) # back to ~0.4ms! Звичайно, графіки CUDA мають свої обмеження (статичні форми, статичний контрольний потік і Вони найкраще підходять для ізоляції деяких фрагментів продуктивності GPU, а не e2e тренувальних ланцюгів. Так є Так є Не тільки системи не спрацьовують - дивіться дані Іноді навіть ідеальна методологія бенчмаркування не розповідає всю історію, тому що виробничі умови відрізняються від ізольованого середовища, в якому ми керуємо бенчмарком. Наприклад, припустимо, що ми хочемо порівняти груповане ядро Gemm всередині шару Mixture-of-Experts. У типовому шарі MoE, токени динамічно маршрутуються до різних експертів на основі ймовірності маршрутизатора. В результаті, фактичне навантаження, побачене кожним експертом (і груповим ядром Gemm), залежить від того, наскільки збалансовано маршрутизація. Наївно, ми генерували б випадкові ймовірності маршрутизатора і штучно маршрутували токени на їх основі. проте, протягом навчання наш маршрутизатор може бути незбалансованим, використовуючи декількох експертів набагато важче, ніж інші. Давайте подивимося, як основний розподіл токенів, спрямованих до кожного експерта, впливає на наші вимірювання. import numpy as np import torch import triton def sample_expert_assignments( seq_len: int, num_experts: int, top_k: int, use_beta: bool = True, alpha: float = 1.0, beta: float = 5.0, ) -> tuple[torch.Tensor, torch.Tensor]: if use_beta: expert_weights = np.random.beta(alpha, beta, num_experts) expert_weights = expert_weights / expert_weights.sum() else: expert_weights = np.ones(num_experts) / num_experts gumbel_noise = -np.log(-np.log(np.random.uniform(0, 1, (seq_len, num_experts)))) log_weights = np.log(expert_weights + 1e-10) scores = log_weights[None, :] + gumbel_noise expert_indices = torch.from_numpy(np.argsort(scores, axis=1)[:, -top_k:]) tokens_per_expert = torch.bincount( expert_indices.flatten(), minlength=num_experts ).to(torch.int32) return expert_indices, tokens_per_expert def get_tokens( seq_len: int, hidden: int, num_experts: int, top_k: int, expert_indices: torch.Tensor, tokens_per_expert: torch.Tensor, device: str = "cuda", ) -> torch.Tensor: x_original = torch.randn((seq_len, hidden), dtype=torch.bfloat16, device=device) token_indices = torch.arange(seq_len, device=device)[:, None].expand(-1, top_k) expert_flat = expert_indices.to(device).flatten() token_flat = token_indices.flatten() token_sorted = token_flat[torch.argsort(expert_flat)] return x_original[token_sorted] def get_tensors( seq_len: int = 1024, hidden: int = 4096, intermediate: int = 1536, num_experts: int = 128, top_k: int = 8, use_beta: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: expert_indices, tokens_per_expert = sample_expert_assignments( seq_len, num_experts, top_k, use_beta=use_beta, ) x = get_tokens( seq_len, hidden, num_experts, top_k, expert_indices, tokens_per_expert, device="cuda", ) w = torch.randn( (num_experts, hidden, intermediate), dtype=torch.bfloat16, device="cuda" ) offsets = torch.cumsum(tokens_per_expert, dim=0).to( dtype=torch.int32, device="cuda" ) return x, w, offsets def benchmark( seq_len: int = 1024, hidden: int = 4096, intermediate: int = 1536, num_experts: int = 128, top_k: int = 8, use_beta: bool = True, num_iters: int = 250, ): x, w, offsets = get_tensors( seq_len, hidden, intermediate, num_experts, top_k, use_beta, ) return triton.testing.do_bench_cudagraph( lambda: torch._grouped_mm(x, w, offs=offsets), rep=num_iters ) params = { "seq_len": 4096, "hidden": 4096, "intermediate": 1536, "num_experts": 128, "top_k": 8, } uniform = benchmark(**params, use_beta=False) beta = benchmark(**params, use_beta=True) print(f"Uniform: {uniform:.5f} ms") print(f"Beta: {beta:.5f} ms") # Uniform: 0.96811 ms # Beta: 1.02422 ms При однорідному присвоєнні токенів ядро GroupedGEMM на 6% швидше, ніж при вибірці з бета-розподілу (який ми використовуємо для моделювання незбалансованого навантаження)! Покласти все разом Щоб підсумувати: When writing the benchmarking code by hand, do not forget to: Flush the L2 cache Use CUDA events Use cuda.synchronize to wait for the completion of all GPU work Alternatively, use or for a built-in solution. triton.testing.do_bench triton.testing.do_bench_cudagraph Regardless of approach, do not forget about the underlying data you’re using. В цілому, ваш бенчмарк є тільки настільки корисним, як його актуальність до того, що ми насправді побачимо в реальних виробничих процесах.Якщо ми дбаємо про CPU, не позбавляйтеся від нього.Якщо ви очікуєте конкретного розподілу даних під час навчання, спробуйте взяти з нього зразок під час бенчмаркування.