Cualquiera que trabaje con el código modelo de PyTorch comienza preguntando las mismas preguntas: 
 
 ¿Por qué esto dura tanto tiempo? ¿Cómo puedo hacer que mi ciclo de entrenamiento sea más rápido? ¿Por qué esto dura tanto tiempo? ¿Cómo puedo hacer que mi ciclo de entrenamiento sea más rápido? Ya sea que usted sea un ingeniero de ML, un investigador o simplemente decidió jugar con un repositorio de ML aleatorio durante el fin de semana, eventualmente intentará entender cómo acelerar su código. Sin embargo, antes de que podamos hacer eso, necesitamos aprender a medir el rendimiento correctamente.Y luego sacar las conclusiones correctas de estas mediciones.Este artículo trata exactamente de eso, sobre el benchmarking adecuado del código CUDA o PyTorch. Ejecución de flujo Utilicemos la multiplicación de matriz en un H100 como ejemplo en curso. Sin embargo, antes de pensar en las mediciones, vamos a tener la terminología recta. Por "núcleo" nos referimos a una función (o realmente cualquier operación) que se ejecuta en una GPU. Su código podría tener muchos de estos, de decenas a miles de operaciones durante cada iteración del ciclo de entrenamiento, por ejemplo: 
 
 
 
 
 
 torch.zeros((16, 16), dispositivo=”cuda) a + b, donde a y b son tensores en la memoria GPU A @ B flash_atención(q, k, v) Y así en A pesar de que el kernel se ejecuta en una GPU, su es controlado por la CPU. Si imaginamos un archivo Python con el código de entrenamiento, se alternará entre el trabajo “preparatorio”, como for-loops o if-statements, y los lanzamientos reales del núcleo. Debido a esto, la CPU corre adelante y cada vez que toca un núcleo, programa ese núcleo para ejecutarse en la GPU añadiéndolo a una cola. Cuando el próximo núcleo se agrega a la cola, no hay garantía de que el anterior ya haya terminado. Lanzamiento Esto permite que la GPU evite estar sentada en el vacío: siempre puede encontrar trabajo en esa cola del núcleo (excepto cuando se llama explícitamente la sincronización o cuando la ejecución está vinculada a la CPU). El enfoque ingenuo y por qué está equivocado Cuando primero pensamos en medir el tiempo en Python, el instinto natural es alcanzar el módulo de tiempo y ejecutar algo así: 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 ejecutamos este código en dos conjuntos diferentes de formas: y , veríamos que básicamente no hay diferencia en el tiempo de matmul. De hecho, en mi curso, es realmente más rápido multiplicar matrices que son 256 veces más grandes. Este resultado es muy sospechoso y debería hacernos pensar en lo que este código realmente mide. 16x32 4096x8192 De hecho, nunca esperamos que la GPU matmul termine. puesto que todo el código “preparatorio” se ejecuta en la CPU, De hecho, está midiendo el tiempo que nos lleva a programar el núcleo matmul en la cola de la GPU y seguir adelante con nuestras vidas. time.perf_counter_ns() Lo que podemos hacer es añadir después de la matmul para forzar la CPU a esperar a la GPU. Pero eso todavía no es ideal, porque estamos midiendo el tiempo del reloj de pared transcurrido en la CPU, lo que incluye la planificación sobre la cabeza. torch.cuda.synchronize() Medición con eventos CUDA La manera correcta de medir el tiempo de la GPU en lugar del tiempo de la CPU es con eventos CUDA. Estos son marcadores que podemos insertar directamente en el flujo de ejecución de la GPU. La GPU registra timestamps cada vez que llega a cada evento CUDA, dándonos el tiempo de ejecución real en el 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
 Los eventos se insertan en el mismo flujo que nuestro núcleo matmul. garantiza que se registren los timestamps. El método nos da la diferencia entre los timestamps de inicio y final. enable_timing=True .elapsed_time Ahora obtenemos números reales, no resultados absurdamente rápidos del enfoque ingenuo. Desde este punto de referencia, el matmul con formas grandes toma al menos 15 veces más tiempo. También vale la pena señalar que es mejor incluir iteraciones de calentamiento al comienzo, porque a menudo hay costes únicos incluidos en el primer lanzamiento, como la compilación JIT. El cache L2 Cuando ejecutamos el mismo núcleo repetidamente con los mismos datos (y esto es exactamente lo que estamos haciendo), esos datos permanecen en la caché L2 de la GPU. En las GPUs de NVIDIA, existe una jerarquía de memoria: HBM es la mayor y más lenta, seguida de unidades de memoria progresivamente más pequeñas y más rápidas (cache L2 → cache L1 o memoria compartida → memoria de 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 Vuelta 6MB 32 GB A100 Ampera 40 MB 80GB H100 Hopper 50MB 80GB B200 Blackwell 612MB 192 GB Debido a esta caché, nuestras mediciones pueden mostrar mejores resultados que lo que veríamos en la producción, donde los datos cambian entre iteraciones. Para obtener mediciones más realistas, tendríamos que limpiar la caché. Una manera de hacer eso es asignar un buffer grande y actualizarlo al comienzo de cada iteración. Esto es, por supuesto, un heurístico en lugar de un flash de caché garantizado: 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
 De hecho, hubo una diferencia significativa en el resultado entre esta ejecución y la anterior, pero en realidad re-executar el código varias veces altera los resultados un poco, y nuestra diferencia está bien dentro de la variación esperada - por lo que para este problema en particular la caché no tuvo un impacto significativo. Soluciones construidas ¿Tenemos que recordar todo esto sólo para hacer referencia a algo?La respuesta es no. De hecho, hay una biblioteca llamada , que fue desarrollado originalmente por OpenAI en 2021 para escribir los núcleos de GPU de una manera Pythonic. Hoy en día tiene un módulo de prueba integrado con utilidades de benchmarking adecuadas, que son capaces de hacer todo lo anterior por sí solos. Tritón # 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)
 Además de lo que hemos cubierto, permite obtener varias estadísticas de la referencia, no sólo la media. Por ejemplo, uno podría solicitar el mediano o el 99o percentil. Hay otros argumentos útiles, que se cubren en la doctrina. do_bench Para utilizar esta función, podemos llamar simplemente: import triton

bench_time = triton.testing.do_bench(lambda: matmul(a, b))
 Ejecución vinculada a la CPU y gráficos CUDA Para los núcleos con un tiempo de ejecución muy corto, la carga de lanzamiento puede llegar a ser tan grande que excede el tiempo del núcleo mismo. Lo mismo ocurre cuando hay un trabajo significativo de la CPU (normalmente no optimizado) entre los lanzamientos del núcleo. Añadamos un looong for-loop a nuestro núcleo matmul, ya que sabemos que los for-loops son realmente lentos en la CPU en Python: def cpu_heavy_matmul(a, b):
    cnt = 0
    for _ in range(100_000):
        cnt += 1
    return a @ b
 Veremos que el tiempo de matmul (4096, 8192) por (8192, 4096) aumentó de 0.4ms a 2.2ms - más de 5x! Ahora, dependiendo de lo que su curso de entrenamiento real parezca, es posible que desee eliminar la CPU-overhead de los resultados de benchmarking o no. Por ejemplo, si su entrenamiento ya está ligado a la CPU, entonces surgirían los mismos problemas durante la producción. Entonces, más probablemente, no tiene sentido deshacerse del tiempo de CPU para nuestras mediciones. Sin embargo, en las operaciones de entrenamiento normales, la CPU no debería estar en el camino del rendimiento óptimo de entrenamiento - por lo general, la CPU corre antes de la GPU. En ese caso, incluso si ese núcleo tiene una overhead significativa, no alterará el tiempo de ejecución del núcleo en la GPU. Entonces querríamos eliminar la overhead de nuestras mediciones. Para ello, podemos utilizar gráficos CUDA. CUDA Graphs nos permite grabar una serie de núcleos una vez y reproducirlos con una participación mínima de la CPU. , que proporciona exactamente la misma funcionalidad. Usando esto para nuestro nuevo matmul, obtenemos el tiempo del núcleo inicial (antes de que se introdujera 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!
 Por supuesto, los gráficos CUDA vienen con sus propias restricciones (formas estáticas, flujo de control estático y Son más adecuados para aislar algunos pedazos del rendimiento de la GPU en lugar de los loop de entrenamiento e2e. so on Así es No sólo los sistemas fallan - mirando los datos A veces incluso la metodología de benchmarking perfecta no cuenta la historia entera, porque las condiciones de producción son diferentes del entorno aislado en el que ejecutamos la referencia. Por ejemplo, digamos que queremos comparar el núcleo gemm agrupado dentro de una capa Mixture-of-Experts. En una capa MoE típica, los tokens se redireccionan dinámicamente a diferentes expertos basados en las probabilidades del router. Como resultado, la carga de trabajo real vista por cada experto (y por el núcleo gemm agrupado) depende de cuán equilibrado es el enrutamiento. Naivamente, generaríamos probabilidades aleatorias de router y rotaríamos tokens artificialmente basados en ellos. sin embargo, durante el entrenamiento, nuestro router podría estar desequilibrado, utilizando varios expertos mucho más pesadamente que los otros. Veamos cómo la distribución subyacente de los tokens dirigidos a cada experto afecta a nuestras mediciones. 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 la asignación uniforme de token, el núcleo GroupedGEMM es un 6% más rápido que cuando tomamos muestras de la distribución beta (que usamos para modelar la carga desequilibrada)! Poniendo todo juntos Para resumir: 
 
 
 
 
 
 
 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. En resumen, su índice de referencia es tan útil como su relevancia para lo que realmente veremos en las producciones reales.Si nos preocupamos por la CPU, no lo desechemos.Si espera una distribución de datos específica durante el entrenamiento, intente tomar muestras de ella durante el 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

Este audio es producido en el idioma original de la historia!

¿Qué realmente determina la velocidad de tu código PyTorch?

About Author

COMENTARIOS

ETIQUETAS

ESTE ARTÍCULO FUE PRESENTADO EN

Related Stories

Crecimiento de las criptomonedas: creación de perfiles de usuarios eficaces

Navegando por las aguas: desarrollo de aplicaciones RAG de nivel de producción con lagos de datos

La fuga rápida del sistema Claude Sonnet 3.5: un análisis forense

Valhalla de Floki se une como patrocinador asociado de la gira de la India por Sri Lanka

Crecimiento de las criptomonedas: creación de perfiles de usuarios eficaces

Navegando por las aguas: desarrollo de aplicaciones RAG de nivel de producción con lagos de datos

La fuga rápida del sistema Claude Sonnet 3.5: un análisis forense

Valhalla de Floki se une como patrocinador asociado de la gira de la India por Sri Lanka

Light-Mode

Classic

Newspaper

Minty

Dark-Mode

Neon Noir

Minty

HN StartUps