Chiunque stia lavorando con il codice modello PyTorch inizia a porre le stesse domande: 
 
 Perché questo è così lungo? come faccio a rendere il mio ciclo di allenamento più veloce? Perché questo è così lungo? come faccio a rendere il mio ciclo di allenamento più veloce? Che tu sia un ingegnere ML, un ricercatore o appena deciso di giocare con un repository ML casuale durante il fine settimana, alla fine cercherai di capire come accelerare il tuo codice. Tuttavia, prima di poterlo fare, dobbiamo imparare a misurare correttamente le prestazioni. E poi trarre le giuste conclusioni da queste misurazioni. Questo articolo riguarda esattamente questo, circa il corretto benchmarking del codice CUDA o PyTorch. Flusso di esecuzione Utilizziamo la moltiplicazione di matrice su un H100 come esempio in esecuzione. Tuttavia, prima di pensare alle misurazioni, prendiamo la terminologia dritta.Con "core" intendiamo una funzione (o davvero qualsiasi operazione) che viene eseguita su una GPU. Il tuo codice potrebbe avere un sacco di questi, da decine a migliaia di operazioni durante ogni iterazione del ciclo di allenamento, ad esempio: 
 
 
 
 
 
 torch.zeros((16, 16), dispositivo=”cuda) a + b, dove a e b sono tensori nella memoria GPU A @ B flash_attenzione (q, k, v) E così su Nonostante il kernel sia eseguito su una GPU, la sua è controllato dalla CPU. Se immaginamo un file Python con il codice di formazione, si alternano tra il lavoro "preparatorio", come for-loops o if-statements, e i lanci del kernel effettivi. Per questo motivo, la CPU corre avanti e ogni volta che colpisce un kernel, pianifica che il kernel per l'esecuzione sulla GPU aggiungendolo a una coda. Quando il kernel successivo viene aggiunto alla coda, non c'è alcuna garanzia che il precedente sia già terminato. Lancio Ciò consente alla GPU di evitare di stare seduta: può sempre trovare lavoro in quella coda del kernel (tranne quando la sincronizzazione è esplicitamente chiamata o quando l'esecuzione è legata alla CPU). Il metodo ingenuo e perché è sbagliato Quando pensiamo per la prima volta a misurare il tempo in Python, l'istinto naturale è quello di raggiungere il modulo del tempo e eseguire qualcosa di simile a questo: 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
 Se eseguiamo questo codice su due diversi set di forme: e , vedremmo che non c'è fondamentalmente alcuna differenza nel tempo matmul. In realtà, nel mio corso, è in realtà più veloce moltiplicare matrici che sono 256 volte più grandi. Questo risultato è molto sospetto e dovrebbe farci pensare a cosa misura effettivamente questo codice. 16x32 4096x8192 In realtà, non aspettiamo mai che la GPU matmul finisca. Dal momento che tutto il codice "preparatorio" viene eseguito sulla CPU, In realtà sta misurando il tempo necessario per pianificare il kernel matmul nella coda della GPU e procedere con le nostre vite. time.perf_counter_ns() Ciò che possiamo fare è aggiungere Dopo il matmul per costringere la CPU ad aspettare la GPU. Ma questo non è ancora l'ideale, perché stiamo misurando il tempo trascorso sul CPU, che include la pianificazione overhead. torch.cuda.synchronize() Misurazione con eventi CUDA Il modo corretto per misurare il tempo della GPU invece del tempo della CPU è con gli eventi CUDA. Questi sono i marcatori che possiamo inserire direttamente nel flusso di esecuzione della GPU. La GPU registra i timestamp ogni volta che raggiunge ogni evento CUDA, dandoci il tempo di esecuzione effettivo sul dispositivo. 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
 Gli eventi vengono inseriti nello stesso flusso del nostro kernel matmul. garantisce che i timestamp siano registrati. Il metodo ci dà la differenza tra i timestamp di inizio e di fine. enable_timing=True .elapsed_time Ora otteniamo numeri reali, non risultati assurdamente veloci dall'approccio ingenuo. Da questo punto di riferimento, il matmul con forme grandi richiede almeno 15x di più. Vale anche la pena notare che è meglio includere iterazioni di riscaldamento all'inizio, perché ci sono spesso costi singoli inclusi nel primo lancio, come la compilazione JIT. Il cache L2 Un altro problema con la nostra misurazione attuale è la mancanza di un flush della cache L2.Quando eseguiamo lo stesso kernel ripetutamente con gli stessi dati (e questo è esattamente ciò che stiamo facendo), quei dati rimangono nella cache L2 della GPU. Nelle GPU NVIDIA, c'è una gerarchia di memoria: HBM è la più grande e più lenta, seguita da unità di memoria progressivamente più piccole e più veloci (cache L2 → cache L1 o memoria condivisa → memoria del registro). 
 
 
 
 
 
 
 
 
 
 
 
 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 Ritorno 6 MB 32GB di A100 Ampere 40 MB di 80GB H100 Hopper 50 MB di 80GB Il B200 Il Blackwell di 126 MB di 192GB A causa di questa cache, le nostre misurazioni potrebbero mostrare risultati migliori di quelli che vedremmo nella produzione, dove i dati cambiano tra iterazioni. Per ottenere misurazioni più realistiche, dovremmo spruzzare la cache. Uno dei modi per farlo è quello di assegnare un buffer grande e aggiornarlo all'inizio di ogni iterazione. Questo è, naturalmente, un flush heuristico piuttosto che un cache garantito: 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
 In realtà, c'è stata una differenza significativa nel risultato tra questa esecuzione e quella precedente, ma in realtà la re-esecuzione del codice più volte altera i risultati un po ', e la nostra differenza è ben all'interno della varianza prevista - quindi per questo particolare problema la cache non ha avuto un impatto significativo. Soluzioni built-in Dobbiamo ricordare tutto questo solo per benchmark qualcosa?La risposta è no. In realtà, c'è una libreria chiamata , che è stato originariamente sviluppato da OpenAI nel 2021 per scrivere i nuclei GPU in modo Pythonic. Oggi ha un modulo di test integrato con utili di benchmarking appropriati, che sono in grado di fare tutto quanto sopra da soli. Tritone # 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)
 Oltre a ciò che abbiamo coperto, Per esempio, si potrebbe richiedere la media o il 99 ° percentile. Ci sono altri argomenti utili, che sono coperti nella docstring. do_bench Per utilizzare questa funzione, possiamo semplicemente chiamare: import triton

bench_time = triton.testing.do_bench(lambda: matmul(a, b))
 Esecuzione legata alla CPU e grafici CUDA Per i kernel con un tempo di esecuzione molto breve, l'overhead di lancio potrebbe diventare così grande da superare il tempo del kernel stesso. Lo stesso accade quando c'è un lavoro significativo della CPU (di solito non ottimizzato) tra i lanci del kernel. Aggiungiamo un looong for-loop al nostro kernel matmul, poiché sappiamo che i for-loops sono davvero lenti sulla CPU in Python: def cpu_heavy_matmul(a, b):
    cnt = 0
    for _ in range(100_000):
        cnt += 1
    return a @ b
 Vedremo che il (4096, 8192) di (8192, 4096) tempo matmul aumentato da 0,4ms a 2,2ms - più di 5x! Ad esempio, se il tuo allenamento è già legato alla CPU, allora durante la produzione sorgerebbero gli stessi problemi. Allora, molto probabilmente, non ha senso sbarazzarsi del tempo di CPU per le nostre misurazioni. Tuttavia, nelle normali esecuzioni di allenamento, la CPU non dovrebbe ostacolare le prestazioni di allenamento ottimali - di solito la CPU corre prima della GPU. In questo caso, anche se quel kernel ha un overhead significativo, non cambierà il tempo di esecuzione del kernel sulla GPU. Poi vorremmo eliminare il overhead dalle nostre misurazioni. Per fare questo, possiamo usare i grafici CUDA. I grafici CUDA ci permettono di registrare una serie di nuclei una volta e riprodurli con un minimo coinvolgimento della CPU. Utilizzando questo per il nostro nuovo matmul, otteniamo il tempo del kernel iniziale (prima dell'introduzione dell'overhead della CPU). 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!
 Naturalmente, i grafici CUDA hanno le proprie restrizioni (forma statica, flusso di controllo statico e Sono meglio adatti per isolare alcuni pezzi della prestazione della GPU piuttosto che i loop di formazione e2e. so on Così è Non solo i sistemi falliscono - guardando i dati A volte anche la metodologia di benchmarking perfetta non racconta l'intera storia, perché le condizioni di produzione sono diverse dall'ambiente isolato in cui eseguiamo il benchmark. Ad esempio, diciamo che vogliamo benchmarkare il nucleo gemm raggruppato all'interno di un livello Mixture-of-Experts. In uno strato MoE tipico, i token vengono dinamicamente indirizzati a diversi esperti in base alle probabilità del router. Di conseguenza, il carico di lavoro effettivo visto da ciascun esperto (e dal nucleo gemm raggruppato) dipende da quanto sia equilibrato il routing. In modo ingenuo, genereremmo probabilità di router casuali e tracceremmo artificialmente token basati su di essi.Tuttavia, durante l'addestramento, il nostro router potrebbe essere squilibrato, utilizzando diversi esperti molto più pesantemente degli altri. Vediamo come la distribuzione sottostante dei token indirizzati a ciascun esperto influenza le nostre misurazioni. 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
 Con l'assegnazione uniforme dei token, il nucleo GroupedGEMM è del 6% più veloce rispetto a quando si prende il campione dalla distribuzione beta (che usiamo per modellare il carico squilibrato)! Mettendo tutto insieme per riassumere: 
 
 
 
 
 
 
 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. Tutto sommato, il tuo benchmark è utile solo quanto la sua rilevanza a ciò che vedremo effettivamente nelle esecuzioni di produzione reali.Se ci preoccupiamo dell'overhead della CPU, non sbarazzarci di esso.Se ti aspetti una distribuzione di dati specifica durante la formazione, prova a prenderne il campione durante il 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

Questo audio è prodotto nella lingua originale della storia!

Cosa determina davvero la velocità del tuo codice PyTorch?

About Author

COMMENTI

CARTELLINI

QUESTO ARTICOLO È STATO PRESENTATO IN

Related Stories

Arron Goolsbey Talks Future of Gaming with AI and Blockchain at Mythical Games

New Year, New Exciting Features: Mint Your Name, Watch Web 2.5 and Meet The Top Startups of The Year

Coding with Latenode's AI + JavaScript Assistant: A Low-Code Platform for Simplified Development

Meet Dealscard, Winner of Startups of The Year 2024 in Ghaziabad/Deal Aggregator

Arron Goolsbey Talks Future of Gaming with AI and Blockchain at Mythical Games

New Year, New Exciting Features: Mint Your Name, Watch Web 2.5 and Meet The Top Startups of The Year

Coding with Latenode's AI + JavaScript Assistant: A Low-Code Platform for Simplified Development

Meet Dealscard, Winner of Startups of The Year 2024 in Ghaziabad/Deal Aggregator

Light-Mode

Classic

Newspaper

Minty

Dark-Mode

Neon Noir

Minty

HN StartUps