Jeder, der mit PyTorch-Modellcode arbeitet, beginnt mit den gleichen Fragen: 
 
 Warum dauert das so lange?Wie kann ich meine Trainingsschleife schneller machen? Warum dauert das so lange?Wie kann ich meine Trainingsschleife schneller machen? Egal, ob Sie ein ML-Ingenieur, ein Forscher sind oder gerade beschlossen haben, über das Wochenende mit einem zufälligen ML-Repository zu spielen, Sie werden schließlich versuchen, zu verstehen, wie Sie Ihren Code beschleunigen können. Bevor wir das jedoch tun können, müssen wir lernen, wie wir die Leistung richtig messen können. Und dann die richtigen Schlussfolgerungen aus diesen Messungen ziehen. Dieser Artikel geht genau darum, wie man den CUDA- oder PyTorch-Code richtig benchmarkt. Durchführung Flow Lassen Sie uns die Matrix-Multiplikation auf einem H100 als das laufende Beispiel verwenden. Unter „Kernel“ meinen wir eine Funktion (oder wirklich jede Operation), die auf einer GPU ausgeführt wird. 
 
 
 
 
 
 torch.zeros((16, 16), Gerät = „Cuda) a + b, wo a und b Tensoren im GPU-Speicher sind A @ B Flash_Aufmerksamkeit (q, k, v) Und so ist Obwohl der Kernel auf einer GPU ausgeführt wird, wird von der CPU gesteuert. Wenn wir uns eine Python-Datei mit dem Trainingskode vorstellen, schaltet sie sich zwischen „vorbereitenden“ Arbeiten, wie For-Loops oder If-Statements, und tatsächlichen Kernel-Launchs ab. Aus diesem Grund läuft die CPU vorwärts und jedes Mal, wenn sie einen Kernel trifft, schlägt sie diesen Kernel für die Ausführung auf der GPU ein, indem sie ihn zu einer Warteschlange hinzufügt. Startschuss Dadurch kann die GPU nicht still sitzen: Sie kann immer Arbeit in dieser Kernel-Warteschlange finden (außer wenn die Synchronisierung ausdrücklich aufgerufen wird oder wenn die Ausführung an die CPU gebunden ist). Der naive Ansatz und warum es falsch ist Wenn wir zuerst über die Zeitmessung in Python nachdenken, ist der natürliche Instinkt, das Zeitmodul zu erreichen und so etwas auszuführen: 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
 Wenn wir diesen Code auf zwei verschiedenen Formsätzen ausführen: und , wir würden sehen, dass es im Grunde keinen Unterschied in der matmul-Zeit gibt. In der Tat ist es in meinem Lauf tatsächlich schneller, Matrizen zu multiplizieren, die 256 mal größer sind. 16x32 4096x8192 In der Tat warten wir nie, bis die GPU matmul fertig ist. Da der gesamte "vorbereitende" Code auf der CPU läuft, Es misst tatsächlich die Zeit, die es braucht, um den matmul-Kernel in die GPU-Warteschlange zu schalten und mit unserem Leben fortzufahren. time.perf_counter_ns() Was wir tun können, ist hinzuzufügen nach dem matmul, um die CPU zu zwingen, auf die GPU zu warten. aber das ist immer noch nicht ideal, weil wir die vergangene Wand-Uhrzeit auf der CPU messen, was das Planen über den Kopf beinhaltet. torch.cuda.synchronize() Messung mit CUDA Ereignissen Der richtige Weg, um die GPU-Zeit anstelle der CPU-Zeit zu messen, ist mit CUDA-Ereignissen.Dies sind Marker, die wir direkt in den Ausführungsstrom der GPU einfügen können.Die GPU zeichnet Timestamps auf, wenn sie jedes CUDA-Ereignis erreicht, was uns die tatsächliche Ausführungszeit auf dem Gerät gibt. 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
 Die Ereignisse werden in den gleichen Strom wie unser matmul Kernel eingefügt. stellt sicher, dass Timestamps aufgezeichnet werden. Die Methode gibt uns den Unterschied zwischen den Start- und Endzeitstempeln. enable_timing=True .elapsed_time Nun bekommen wir echte Zahlen, nicht absurd schnelle Ergebnisse aus dem naiven Ansatz. Von diesem Benchmark aus dauert der matmul mit großen Formen mindestens 15x länger. Es ist auch erwähnenswert, dass es am besten ist, Warmup-Iterationen zu Beginn aufzunehmen, da es oft einmalige Kosten gibt, die beim ersten Start enthalten sind, wie z. B. die JIT-Kompilation. Der L2 Cache Wenn wir den gleichen Kernel wiederholt mit den gleichen Daten ausführen (und das ist genau das, was wir tun), bleiben diese Daten im L2-Cache der GPU. In NVIDIA GPUs gibt es eine Speicherhierarchie: HBM ist der größte und langsamste, gefolgt von progressiv kleineren und schnelleren Speicher-Einheiten (L2 Cache → L1 Cache oder Shared Memory → Register-Speicher). 
 
 
 
 
 
 
 
 
 
 
 
 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 zurück 6 MB 32 GB A100 Ampere 40 MB 80 GB H100 Hopper 50 MB 80 GB B200 Blackwell 126 MB 192 GB Aufgrund dieses Caches könnten unsere Messungen bessere Ergebnisse zeigen als die, die wir in der Produktion sehen würden, wo sich die Daten zwischen den Iterationen ändern. Um realistischere Messungen zu erhalten, müssten wir den Cache ausspülen. Eine Möglichkeit, dies zu tun, besteht darin, einen großen Puffer zuzuweisen und ihn zu Beginn jeder Iteration zu aktualisieren. 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
 Tatsächlich gab es einen signifikanten Unterschied im Ergebnis zwischen diesem Run und dem vorherigen, aber in Wirklichkeit ändert das Re-Run des Codes mehrmals die Ergebnisse ein wenig, und unser Unterschied liegt gut innerhalb der erwarteten Varianz - so dass für dieses spezielle Problem der Cache keinen signifikanten Einfluss hatte. eingebaute Lösungen Muss man sich an all das erinnern, nur um etwas zu benchmarken?Die Antwort ist nein. , die ursprünglich von OpenAI im Jahr 2021 entwickelt wurde, um GPU-Kernen auf Pythonische Weise zu schreiben.Heute hat es ein eingebautes Testmodul mit richtigen Benchmarking-Tools, die alle oben genannten Dinge selbst tun können. 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)
 Neben dem, was wir abgedeckt haben, Ermöglicht es, mehrere Statistiken aus dem Benchmark herauszuholen, nicht nur den Durchschnitt. Zum Beispiel könnte man den Median oder den 99. Prozentil anfordern. Es gibt andere nützliche Argumente, die in der Docstring abgedeckt sind. do_bench Um diese Funktion zu nutzen, können wir einfach anrufen: import triton

bench_time = triton.testing.do_bench(lambda: matmul(a, b))
 CPU-gebundene Ausführung und CUDA-Grafiken Bei Kerneln mit einer sehr kurzen Ausführungszeit kann der Launch-Overhead so groß werden, dass er die Kernelzeit selbst übersteigt. Lassen Sie uns einen looong for-loop zu unserem matmul-kernel hinzufügen, da wir wissen, dass for-loops auf der CPU in python wirklich langsam sind: def cpu_heavy_matmul(a, b):
    cnt = 0
    for _ in range(100_000):
        cnt += 1
    return a @ b
 Wir werden sehen, dass die (4096, 8192) von (8192, 4096) matmul Zeit von 0,4ms auf 2,2ms erhöht - mehr als 5x! Nun, je nachdem, wie Ihr tatsächlicher Trainingslauf aussieht, möchten Sie möglicherweise entweder den CPU-Overhead aus den Benchmarking-Ergebnissen beseitigen oder nicht. Zum Beispiel, wenn Ihr Training bereits CPU-gebunden ist, dann würden während des Produktionslaufs auch dieselben Probleme auftreten. Dann macht es höchstwahrscheinlich keinen Sinn, die CPU-Zeit für unsere Messungen loszuwerden. Allerdings sollte der CPU bei normalen Trainingslaufs nicht im Wege der optimalen Trainingsleistung stehen - normalerweise läuft der CPU vor der GPU. In diesem Fall, selbst wenn dieser Kernel einen signifikanten Overhead hat, wird er die Kernellaufzeit auf der GPU nicht ändern. Dann möchten wir den Overhead aus unseren Messungen beseitigen. Um dies zu tun, können wir CUDA-Diagramme verwenden. CUDA-Diagramme lassen uns eine Reihe von Kerneln einmal aufzeichnen und sie mit minimaler CPU-Beteiligung wiedergeben. Wenn wir dies für unser neues Matmul verwenden, erhalten wir die ursprüngliche Kernelzeit (vor der CPU-Overhead eingeführt wurde). 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!
 Natürlich haben CUDA-Diagramme ihre eigenen Einschränkungen (statische Formen, statische Steuerung und Sie eignen sich am besten zum Isolieren einiger Teile der GPU-Leistung anstelle von e2e-Trainingslöchern. So ist So ist Nicht nur Systeme scheitern – Blick auf Daten Manchmal erzählt auch die perfekte Benchmarking-Methode nicht die ganze Geschichte, denn die Produktionsbedingungen unterscheiden sich von der isolierten Umgebung, in der wir den Benchmark durchführen. Nehmen wir zum Beispiel an, wir wollen den gruppierten Gemm-Kernel innerhalb einer Mix-of-Experts-Schicht benchmarken. In einer typischen MoE-Schicht werden Token dynamisch an verschiedene Experten auf der Grundlage der Router-Wahrscheinlichkeiten verwiesen. Naiv, wir würden zufällige Router-Wahrscheinlichkeiten generieren und künstlich Routing-Token basierend auf ihnen. Lassen Sie uns sehen, wie die zugrunde liegende Verteilung von Token, die an jeden Experten weitergeleitet werden, unsere Messungen beeinflusst. 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
 Mit der einheitlichen Token-Zuweisung ist der GroupedGEMM-Kernel 6% schneller als bei der Probenahme aus der Beta-Verteilung (die wir verwenden, um eine ungleichgewichtige Belastung zu modellieren)! Alles zusammenbringen Um zusammenzufassen: 
 
 
 
 
 
 
 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. Alles in allem ist Ihr Benchmark nur so nützlich wie seine Relevanz für das, was wir tatsächlich in realen Produktionsläufen sehen werden.Wenn wir uns um CPU-Overhead kümmern, lassen Sie es nicht los.Wenn Sie während des Trainings eine bestimmte Datenverteilung erwarten, versuchen Sie es während des Benchmarking zu sammeln.

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

Dieses Audio ist in der Originalsprache der Geschichte produziert!

Was bestimmt wirklich die Geschwindigkeit Ihres PyTorch-Codes?

About Author

KOMMENTARE

Hängeetiketten

DIESER ARTIKEL WURDE VORGESTELLT IN

Related Stories

HackerNoon Decoded 2024: Celebrating Our Startups Community!

HackerNoon Decoded 2024: Celebrating Our Management Community!

HackerNoon Decoded 2024: Celebrating Our Writing Community!

HackerNoon Decoded: The Top 10 Countries Where HackerNoon Is the Most Active

HackerNoon Decoded 2024: Celebrating Our Startups Community!

HackerNoon Decoded 2024: Celebrating Our Management Community!

HackerNoon Decoded 2024: Celebrating Our Writing Community!

HackerNoon Decoded: The Top 10 Countries Where HackerNoon Is the Most Active

Light-Mode

Classic

Newspaper

Minty

Dark-Mode

Neon Noir

Minty

HN StartUps