Siapa pun yang bekerja dengan kode model PyTorch mulai mengajukan pertanyaan yang sama: 
 
 Mengapa ini memakan waktu begitu lama? bagaimana saya membuat lingkaran latihan saya lebih cepat? Mengapa ini memakan waktu begitu lama? bagaimana saya membuat lingkaran latihan saya lebih cepat? Apakah Anda seorang insinyur ML, seorang peneliti atau hanya memutuskan untuk bermain dengan repositori ML acak selama akhir pekan, Anda akhirnya akan mencoba untuk memahami bagaimana mempercepat kode Anda. Namun, sebelum kita bisa melakukannya, kita perlu mempelajari cara mengukur kinerja dengan benar. Dan kemudian menarik kesimpulan yang tepat dari pengukuran ini. Artikel ini tentang persis itu, tentang benchmarking CUDA atau kode PyTorch dengan benar. Eksekusi Flow Mari kita gunakan multiplikasi matriks pada H100 sebagai contoh berjalan sepanjang. Namun, sebelum kita berpikir tentang pengukuran, mari kita mendapatkan terminologi langsung. oleh "kernel" kita berarti fungsi (atau sebenarnya operasi apa pun) yang berjalan pada GPU. kode Anda mungkin memiliki banyak dari ini, dari puluhan hingga ribuan operasi selama setiap iterasi lingkaran pelatihan, misalnya: 
 
 
 
 
 
 torch.zeros((16, 16), perangkat=”cuda) a + b, di mana a dan b adalah tensor dalam memori GPU A @ B flash_attention (Q, K, V ) Dan demikianlah Meskipun kernel ini dijalankan pada GPU, adalah dikendalikan oleh CPU. Jika kita membayangkan file Python dengan kode pelatihan, itu bergantian antara pekerjaan "perawatan", seperti for-loops atau if-statements, dan peluncuran kernel yang sebenarnya. Karena ini, CPU berjalan ke depan dan setiap kali memukul kernel, itu jadwalkan kernel itu untuk dijalankan pada GPU dengan menambahkannya ke baris. peluncuran Ini memungkinkan GPU untuk menghindari duduk santai: ia selalu dapat menemukan pekerjaan di barisan kernel itu (kecuali ketika sinkronisasi secara eksplisit dipanggil atau ketika eksekusi terikat CPU). Pendekatan naif dan mengapa itu salah Ketika kita pertama kali berpikir tentang mengukur waktu di Python, naluri alami adalah untuk menjangkau modul waktu dan menjalankan sesuatu seperti ini: 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
 Jika kita menjalankan kode ini pada dua set bentuk yang berbeda: dan , kita akan melihat bahwa pada dasarnya tidak ada perbedaan dalam waktu matmul. pada kenyataannya, dalam perjalanan saya, sebenarnya lebih cepat untuk mengalikan matriks yang 256 kali lebih besar. hasil ini sangat mencurigakan dan harus membuat kita berpikir tentang apa yang sebenarnya diukur kode ini. 16x32 4096x8192 Sebenarnya, kita tidak pernah menunggu sampai GPU matmul selesai. Karena semua kode “perawatan” berjalan pada CPU, sebenarnya mengukur waktu yang dibutuhkan bagi kita untuk merencanakan kernel matmul ke baris GPU dan melanjutkan dengan kehidupan kita. time.perf_counter_ns() Apa yang bisa kita lakukan adalah menambahkan setelah matmul untuk memaksa CPU untuk menunggu GPU. tetapi itu masih tidak ideal, karena kami mengukur waktu jam dinding yang telah berlalu pada CPU, yang mencakup perencanaan overhead. torch.cuda.synchronize() Mengukur dengan acara CUDA Cara yang benar untuk mengukur waktu GPU alih-alih waktu CPU adalah dengan peristiwa CUDA. ini adalah penanda yang dapat kita masukkan langsung ke aliran eksekusi GPU. GPU merekam timestamps setiap kali mencapai setiap peristiwa CUDA, memberi kita waktu eksekusi aktual pada perangkat. 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
 Peristiwa-peristiwa tersebut disisipkan ke dalam aliran yang sama seperti kernel matmul kita. Hal ini bertujuan untuk memastikan bahwa timetables tercatat. Metode ini memberi kita perbedaan antara waktu awal dan akhir. enable_timing=True .elapsed_time Sekarang kita mendapatkan angka-angka nyata, bukan hasil yang sangat cepat dari pendekatan naif.Dari benchmark ini, matmul dengan bentuk besar membutuhkan waktu setidaknya 15x lebih lama. Perlu juga dicatat bahwa lebih baik untuk memasukkan iterasi pemanasan pada awalnya, karena seringkali ada biaya satu kali yang termasuk dalam peluncuran pertama, seperti kompilasi JIT. Menggunakan Cache L2 Masalah lain dengan pengukuran saat ini adalah kurangnya flush cache L2.Ketika kita menjalankan kernel yang sama berulang kali dengan data yang sama (dan ini adalah apa yang kita lakukan), data itu tetap berada di cache L2 GPU. Dalam GPU NVIDIA, ada hierarki memori: HBM adalah yang terbesar dan paling lambat, diikuti oleh unit memori yang semakin kecil dan lebih cepat (L2 cache → L1 cache atau 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 V100 Kembali 6 MB 32GB A100 Ampera 40 MB 80GB H100 Hopper yang 50 MB 80GB B200 yang Blackwell yang Jumlah 126 MB 192 GB Karena cache ini, pengukuran kami mungkin menunjukkan hasil yang lebih baik daripada yang akan kita lihat dalam produksi, di mana data berubah antara iterasi. Untuk mendapatkan pengukuran yang lebih realistis, kita perlu mencuci cache. Salah satu cara untuk melakukannya adalah mengalokasikan buffer besar dan memperbarui pada awal setiap iterasi. 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
 Sebenarnya, ada perbedaan signifikan dalam hasil antara berjalan ini dan sebelumnya, tetapi dalam kenyataannya menjalankan ulang kode beberapa kali mengubah hasilnya sedikit, dan perbedaan kami baik dalam varians yang diharapkan - jadi untuk masalah khusus ini cache tidak memiliki dampak yang signifikan. Solusi Built-in Apakah kita harus mengingat semua ini hanya untuk membandingkan sesuatu? jawabannya adalah tidak. , yang awalnya dikembangkan oleh OpenAI pada tahun 2021 untuk menulis kernel GPU dengan cara Pythonic. saat ini memiliki modul pengujian built-in dengan utilitas benchmarking yang tepat, yang mampu melakukan semua hal di atas sendiri. mari kita lihat: Triton # 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)
 Selain dari apa yang telah kami sampaikan, memungkinkan untuk mendapatkan beberapa statistik dari benchmark, bukan hanya rata-rata. misalnya, seseorang mungkin meminta median atau percentil 99 lainnya. do_bench Untuk menggunakan fungsi ini, kita dapat memanggil: import triton

bench_time = triton.testing.do_bench(lambda: matmul(a, b))
 Eksekusi CPU-bound dan graf CUDA Untuk kernel dengan waktu eksekusi yang sangat pendek, overhead peluncuran bisa menjadi begitu besar sehingga melebihi waktu kernel itu sendiri. Hal yang sama terjadi ketika ada pekerjaan CPU yang signifikan (biasanya tidak dioptimalkan) antara peluncuran kernel. Mari kita menambahkan looong for-loop ke kernel matmul kita, karena kita tahu bahwa for-loops benar-benar lambat pada CPU di Python: def cpu_heavy_matmul(a, b):
    cnt = 0
    for _ in range(100_000):
        cnt += 1
    return a @ b
 Kita akan melihat bahwa (4096, 8192) oleh (8192, 4096) waktu matmul meningkat dari 0,4ms menjadi 2,2ms - lebih dari 5x! Sekarang, tergantung pada apa yang terlihat dari pelatihan Anda sebenarnya, Anda mungkin ingin menghapus CPU-overhead dari hasil benchmarking atau tidak. Misalnya, jika pelatihan Anda sudah terikat dengan CPU, maka masalah yang sama akan muncul selama pelatihan produksi juga. Kemudian, kemungkinan besar, tidak masuk akal untuk menyingkirkan waktu CPU untuk pengukuran kami. Namun, dalam pelatihan normal, CPU tidak harus menghalangi kinerja pelatihan optimal - biasanya CPU berjalan di depan GPU. Dalam hal itu, bahkan jika satu kernel memiliki overhead yang signifikan, itu tidak akan mengubah waktu runtime kernel pada GPU. Kemudian kami ingin menghapus overhead dari pengukuran kami. Untuk melakukan ini, kita dapat menggunakan graf CUDA. graf CUDA memungkinkan kita merekam serangkaian kernel sekali dan memainkannya dengan keterlibatan CPU minimal. Dengan menggunakan ini untuk matmul baru kami, kita mendapatkan waktu kernel awal (sebelum overhead CPU diperkenalkan). 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!
 Tentu saja, graf CUDA datang dengan keterbatasan mereka sendiri (gaya statis, aliran kontrol statis dan Mereka paling cocok untuk mengisolasi beberapa potongan kinerja GPU daripada loop pelatihan e2e. Begitulah Begitulah Tidak hanya sistem yang gagal - melihat data Kadang-kadang bahkan metodologi benchmarking yang sempurna tidak menceritakan seluruh cerita, karena kondisi produksi berbeda dari lingkungan terisolasi di mana kita menjalankan benchmark. Misalnya, katakanlah kita ingin membandingkan kernel gemm yang dikelompokkan di dalam lapisan Mixture-of-Experts. Dalam lapisan MoE khas, token secara dinamis diarahkan ke spesialis yang berbeda berdasarkan probabilitas router. Hasilnya, beban kerja yang sebenarnya dilihat oleh setiap spesialis (dan oleh kernel gemm yang dikelompokkan) tergantung pada seberapa seimbang routing. Naif, kami akan menghasilkan probabilitas router acak dan secara buatan mengarahkan token berdasarkan mereka. namun, selama pelatihan, router kami mungkin tidak seimbang, menggunakan beberapa ahli jauh lebih berat daripada yang lain. Mari kita lihat bagaimana distribusi token yang mendasari yang diarahkan ke setiap ahli mempengaruhi pengukuran kami. 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
 Dengan atribut token yang seragam, kernel GroupedGEMM 6% lebih cepat dibandingkan ketika kita mengambil sampel dari distribusi beta (yang kita gunakan untuk memodelkan beban yang tidak seimbang)! Menempatkan Segalanya Bersama Untuk menyimpulkan : 
 
 
 
 
 
 
 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. Secara keseluruhan, benchmark Anda hanya berguna seperti relevansi dengan apa yang akan kita lihat dalam proses produksi nyata.Jika kita peduli tentang overhead CPU, jangan menyingkirkannya.Jika Anda mengharapkan distribusi data tertentu selama pelatihan, cobalah untuk mengambil sampel dari itu saat benchmarking.

Walkthroughs, tutorials, guides, and tips. This story will teach you how to do something new or how to do something better.

This story contains new, firsthand information uncovered by the writer.

Connect with me on LinkedIn

Read My Stories

Audio ini diproduksi dalam bahasa asli cerita!

Apa yang Sebenarnya Menentukan Kecepatan Kode PyTorch Anda?

About Author

KOMENTAR

HANG TAG

ARTIKEL INI DISAJIKAN PADA

Related Stories

Building an API-Powered Platform

The Magic Shop

And In The Open Air

Invest With Purpose: Interview with Arian Adeli, CEO at Rivo Trading

Building an API-Powered Platform

The Magic Shop

And In The Open Air

Invest With Purpose: Interview with Arian Adeli, CEO at Rivo Trading

Light-Mode

Classic

Newspaper

Minty

Dark-Mode

Neon Noir

Minty

HN StartUps