Quiconque travaille avec le code modèle PyTorch commence à poser les mêmes questions : 
 
 Pourquoi est-ce que cela prend si longtemps?Comment puis-je faire mon cycle d'entraînement plus rapide? Pourquoi est-ce que cela prend si longtemps?Comment puis-je faire mon cycle d'entraînement plus rapide? Que vous soyez un ingénieur ML, un chercheur ou juste décidé de jouer avec un dépôt ML aléatoire le week-end, vous allez finalement essayer de comprendre comment accélérer votre code. Cependant, avant de pouvoir le faire, nous devons apprendre à mesurer correctement les performances. Et puis tirer les bonnes conclusions à partir de ces mesures. Cet article est sur exactement cela, sur le bon benchmarking du code CUDA ou PyTorch. Exécution du flux Utilisons la multiplication de matrice sur un H100 comme exemple en cours. Cependant, avant de penser aux mesures, prenons la terminologie droite. Par « noyau » nous entendons une fonction (ou vraiment n’importe quelle opération) qui s’exécute sur une GPU. Votre code pourrait en avoir beaucoup, de dizaines à des milliers d’opérations au cours de chaque itération du cycle d’entraînement, par exemple: 
 
 
 
 
 
 torch.zeros((16, 16), appareil=”cuda) a + b, où a et b sont des tensors dans la mémoire GPU A @ B flash_attention (q, k, v) Et ainsi on Bien que le noyau soit exécuté sur une GPU, son est contrôlée par la CPU. Si nous imaginons un fichier Python avec le code de formation, il alternera entre les travaux « préparatoires », tels que les for-loops ou les if-statements, et les lancements réels du noyau. À cause de cela, le CPU avance et chaque fois qu'il frappe un noyau, il planifie ce noyau pour l'exécution sur la GPU en l'ajoutant à une file d'attente. Lorsque le prochain noyau est ajouté à la file d'attente, il n'y a aucune garantie que le précédent est déjà terminé. lancement Cela permet à la GPU d'éviter de s'asseoir dans le vide : elle peut toujours trouver du travail dans cette queue de noyau (sauf lorsque la synchronisation est explicitement appelée ou lorsque l'exécution est liée à la CPU). L’approche naïve et pourquoi elle est fausse Lorsque nous pensons pour la première fois à la mesure du temps en Python, l’instinct naturel est d’atteindre le module du temps et d’exécuter quelque chose comme ceci : 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
 Si nous exécutons ce code sur deux ensembles de formes différents : et En fait, dans ma course, il est en fait plus rapide de multiplier des matrices qui sont 256 fois plus grandes. Ce résultat est très suspect et devrait nous faire penser à ce que ce code mesure réellement. 16x32 4096x8192 En fait, nous n’attendons jamais que le matmul GPU soit terminé. Puisque tout le code « préparatoire » s’exécute sur la CPU, est en fait mesurant le temps qu'il nous faut pour planifier le noyau matmul dans la file d'attente de la GPU et aller de l'avant avec nos vies. time.perf_counter_ns() Ce que nous pouvons faire est d’ajouter après le matmul pour forcer le processeur à attendre la GPU. Mais ce n'est toujours pas idéal, car nous mesurons le temps d'horloge murale écoulé sur la CPU, ce qui inclut la planification au-dessus du cap. Nous n'isolons pas l'exécution GPU réelle de cette façon. torch.cuda.synchronize() Mesurer avec les événements CUDA La bonne façon de mesurer le temps de la GPU au lieu du temps de la CPU est avec les événements CUDA. Ce sont des marqueurs que nous pouvons insérer directement dans le flux d'exécution de la GPU. La GPU enregistre des timestamps chaque fois qu'elle atteint chaque événement CUDA, nous donnant le temps d'exécution réel sur l'appareil. 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
 Les événements sont insérés dans le même flux que notre noyau matmul. Il s’agit d’enregistrer les horaires.Le La méthode nous donne la différence entre les timestamps de début et de fin. enable_timing=True .elapsed_time Maintenant, nous obtenons des chiffres réels, pas des résultats absurdement rapides de l'approche naïve. À partir de cette référence, le matmul avec de grandes formes prend au moins 15 fois plus de temps. Il convient également de noter qu'il est préférable d'inclure des itérations de réchauffement au début, car il y a souvent des coûts uniques inclus dans le premier lancement, tels que la compilation JIT. Le cache L2 Lorsque nous exécutons le même noyau à plusieurs reprises avec les mêmes données (et c’est exactement ce que nous faisons), ces données restent dans le cache L2 de la GPU. Dans les GPU NVIDIA, il y a une hiérarchie de la mémoire: HBM est la plus grande et la plus lente, suivie par des unités de mémoire progressivement plus petites et plus rapides (cache L2 → cache L1 ou mémoire partagée → mémoire de registre). 
 
 
 
 
 
 
 
 
 
 
 
 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 de retour 6MB - 32GB A100 Ampère - 40 MB - 80GB H100 Hopper 50 MB - 80GB Le B200 Blackwell 126 Mo - 192 GB En raison de ce cache, nos mesures pourraient montrer de meilleurs résultats que ce que nous verrions dans la production, où les données changent entre les itérations. Pour obtenir des mesures plus réalistes, nous aurions besoin de rincer le cache. Une façon de le faire est d'allouer un grand tampon et de le mettre à jour au début de chaque itération. 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
 En fait, il y avait une différence significative dans le résultat entre cette exécution et la précédente, mais en réalité, la réexécution du code plusieurs fois modifie les résultats un peu, et notre différence est bien dans la variance attendue - donc pour ce problème particulier le cache n'a pas eu d'impact significatif. Solutions intégrées Est-ce que nous devons nous souvenir de tout cela juste pour comparer quelque chose?La réponse est non.En fait, il y a une bibliothèque appelée , qui a été développé à l'origine par OpenAI en 2021 pour écrire les noyaux GPU de manière pythonique. Aujourd'hui, il a un module de test intégré avec des utilitaires de benchmarking appropriés, qui sont capables de faire tout ce qui précède par eux-mêmes. 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)
 En plus de ce que nous avons couvert, permet d'obtenir plusieurs statistiques de la référence, pas seulement la moyenne. Par exemple, on pourrait demander la médiane ou le 99e percentile. Il y a d'autres arguments utiles, qui sont couverts dans le docstring. do_bench Pour utiliser cette fonction, nous pouvons simplement appeler: import triton

bench_time = triton.testing.do_bench(lambda: matmul(a, b))
 Exécution liée à la CPU et graphiques CUDA Pour les noyaux avec un temps d'exécution très court, le débit de lancement peut devenir si grand qu'il dépasse le temps du noyau lui-même.La même chose se produit lorsqu'il y a un travail de CPU important (généralement non optimisé) entre les lancements du noyau. Ajoutez un looong for-loop à notre noyau matmul, car nous savons que les for-loops sont vraiment lents sur la CPU en Python : def cpu_heavy_matmul(a, b):
    cnt = 0
    for _ in range(100_000):
        cnt += 1
    return a @ b
 Nous verrons que le temps matmul (4096, 8192) par (8192, 4096) a augmenté de 0,4ms à 2,2ms - plus de 5x! Maintenant, en fonction de ce que votre course d'entraînement réel ressemble, vous voudrez peut-être éliminer la tête de la CPU des résultats de référencement ou non. Par exemple, si votre formation est déjà liée à la CPU, les mêmes problèmes se poseront également pendant la course de production. Ensuite, il n'a probablement pas de sens de se débarrasser du temps de CPU pour nos mesures. Cependant, dans les cours d'entraînement normaux, la CPU ne devrait pas s'opposer à la performance d'entraînement optimale - généralement la CPU s'avance sur la GPU. Dans ce cas, même si ce noyau a une tête d'exécution significative, il n'altérera pas le temps d'exécution du noyau sur la GPU. Pour ce faire, nous pouvons utiliser des graphiques CUDA. Les graphiques CUDA nous permettent d'enregistrer une série de noyaux une fois et de les reproduire avec un minimum d'implication de la CPU. En utilisant cela pour notre nouveau matmul, nous obtenons le temps du noyau initial (avant l'introduction de l'overhead de la 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!
 Bien sûr, les graphiques CUDA ont leurs propres contraintes (formes statiques, flux de contrôle statique et Ils sont mieux adaptés pour isoler certains morceaux de la performance de la GPU plutôt que les boucles d'entraînement e2e. Ainsi est Ainsi est Non seulement les systèmes échouent - regardez les données Parfois, même une méthodologie de benchmarking parfaite ne raconte pas toute l’histoire, car les conditions de production sont différentes de l’environnement isolé dans lequel nous exécutons le benchmark. Par exemple, supposons que nous voulions comparer le noyau gemm regroupé à l'intérieur d'une couche Mixture-of-Experts. Dans une couche MoE typique, les jetons sont dynamiquement redirigés vers différents experts en fonction des probabilités du routeur. En conséquence, la charge de travail réelle vue par chaque expert (et par le noyau gemm regroupé) dépend de l'équilibre du routage. Naïvement, nous générerions des probabilités de routeur aléatoire et routerions artificiellement des jetons en fonction d’eux. Cependant, tout au long de la formation, notre routeur pourrait être déséquilibré, en utilisant plusieurs experts beaucoup plus lourdement que les autres. Voyons comment la distribution sous-jacente des jetons dirigés vers chaque expert affecte nos mesures. 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
 Avec l'attribution uniforme de jetons, le noyau GroupedGEMM est 6% plus rapide que lorsque nous prenons l'échantillon de la distribution bêta (que nous utilisons pour modéliser la charge déséquilibrée)! Mettre tout ensemble Pour résumer : 
 
 
 
 
 
 
 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. Dans l'ensemble, votre référence n'est que aussi utile que sa pertinence pour ce que nous verrons réellement dans les opérations de production réelles.Si nous nous soucions de l'excès de CPU, ne vous en débarrassez pas.Si vous attendez une distribution de données spécifique pendant la formation, essayez d'en tirer l'échantillon pendant le référencement.

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

Cet audio est produit dans la langue originale de l'histoire !

Qu’est-ce qui détermine la vitesse de votre code PyTorch ?

About Author

COMMENTAIRES

ÉTIQUETTES

CET ARTICLE A ÉTÉ PARU DANS

Related Stories

Des forums aux flux : comment les algorithmes des réseaux sociaux façonnent l'interaction numérique

Télégramme : le pont de Crypto Island vers le continent

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

Vous voulez gagner un concours d’écriture HackerNoon ? Voici ce que recommandent les gagnants du concours #crypto-api

Des forums aux flux : comment les algorithmes des réseaux sociaux façonnent l'interaction numérique

Télégramme : le pont de Crypto Island vers le continent

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

Vous voulez gagner un concours d’écriture HackerNoon ? Voici ce que recommandent les gagnants du concours #crypto-api

Light-Mode

Classic

Newspaper

Minty

Dark-Mode

Neon Noir

Minty

HN StartUps