Los que han invadido casi todas las facetas de nuestra vida diaria son los modelos decodificadores autorregresivos. Estos modelos aplican operaciones de kernel con gran cantidad de computación para producir tokens uno por uno de una manera que, a primera vista, parece extremadamente ineficiente. Dada la enorme demanda de IA generativa, no sorprende que se esté invirtiendo un extraordinario esfuerzo de ingeniería en su optimización. Ya sea a través de kernels CUDA personalizados, CUDA Graphs, aceleradores de IA dedicados o muestreo especulativo, cualquier técnica que reduzca la latencia y/o el costo incluso en una fracción de un porcentaje es una victoria.
En esta publicación, demostramos una técnica para optimizar la generación de tokens en PyTorch utilizando el entrelazado de flujos CUDA. Si bien es fácil de implementar, el método aborda un cuello de botella específico, que a menudo se pasa por alto, y puede generar mejoras significativas en el rendimiento. Si bien la ejecución de modelos de canalización mediante flujos CUDA es común en la ingeniería de sistemas de IA, no encontramos ningún tutorial que documente la aplicación específica de nivel PyTorch que describimos aquí. Si encuentra útil la técnica, tenga la amabilidad de consultar esta publicación.
Para facilitar nuestra discusión, utilizaremos un modelo de decodificador PyTorch GPT-2 simple de la biblioteca de transformadores de HuggingFace (v5.1.0). Realizaremos nuestros experimentos en una GPU NVIDIA L40S y PyTorch (2.10.0).
Descargo de responsabilidad: el código que compartiremos tiene fines demostrativos. No confíe en su precisión u optimidad. No interprete nuestras menciones a ninguna biblioteca, plataforma o servicio como una aprobación de su uso.
Es importante destacar que el valor del método basado en flujo CUDA que analizaremos puede variar mucho según los detalles de su modelo y entorno de ejecución. Asegúrese de ejecutar sus propios puntos de referencia antes de integrar su uso.
Nuestro enfoque en esta publicación está en las cargas de trabajo de inferencia nativas de PyTorch que siguen siendo extremadamente frecuentes en entornos de desarrollo y prueba. Sin embargo, es importante tener en cuenta que, para entornos de producción, las bibliotecas de inferencia LLM dedicadas, como vLLM o NVIDIA TensorRT-LLM, tienden a ofrecer un mayor rendimiento y deben usarse siempre que sea relevante.
Un modelo de juguete GPT-2
Para simplificar nuestra discusión, usaremos un modelo de decodificador GPT-2 de la biblioteca de transformadores HuggingFace y lo ejecutaremos de forma autorregresiva en un lote de mensajes vacíos.
En el siguiente bloque de código, inicializamos el modelo y definimos una función ingenua de generación de tokens que crea un lote de flujos aleatorios hasta una longitud determinada.
importar antorcha desde transformadores importar GPT2LMHeadModel, GPT2Config torch.set_float32_matmul_precision('high') DEVICE = "cuda" # definir el modelo del decodificador config = GPT2Config.from_pretrained("gpt2") model = GPT2LMHeadModel(config).to(DEVICE).eval() @torch.inference_mode() def generate_sequence(model, max_seqlen, tamaño_por lotes): # Inicializar mensajes con el token BOS all_tokens = torch.full( (tamaño_por lotes, 1), config.bos_token_id, dispositivo=DEVICE, dtype=torch.long ) terminado = torch.zeros(batch_size, dispositivo=DEVICE, dtype=torch.bool) for i in range(max_seqlen): salidas = modelo(all_tokens) # extraer nuevo token logits = salidas.logits[:, -1, :] new_tokens = torch.argmax(logits, dim=-1) # agregar nuevo token a la secuencia all_tokens = torch.cat( [all_tokens, new_tokens.unsqueeze(-1)], dim=-1 ) terminado |= (new_tokens == config.eos_token_id) stop_gpu = torch.all(terminado) # comprobando la condición de parada if stop_gpu.item(): print(f"Todas las secuencias terminaron en el paso {i+1}") break return all_tokens
A continuación, definimos una función de evaluación comparativa simple que utilizamos para medir el rendimiento del tiempo de ejecución y la utilización de la memoria de nuestro generador de tokens en diferentes escenarios.
tiempo de importación, estadísticas def benchmark(func, num_runs=10): # Función de calentamiento() torch.cuda.synchronize() tiempos de ejecución =[]for _ in range(num_runs): # restablecer las estadísticas de memoria antes de cada ejecución torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() start = time.perf_counter() _ = func() torch.cuda.synchronize() end = time.perf_counter() runtimes.append(end – start) # Obtener estadísticas del asignador de memoria de la última ejecución mem_stats = torch.cuda.memory_stats() asignado_pico = mem_stats.get('allocated_bytes.all.peak', 0) reservado_pico = mem_stats.get('reserved_bytes.all.peak', 0) f_peak = reservado_pico – asignado_pico f_pct = ( 100 * f_peak / reservado_pico si reservado_pico > 0 más 0 ) print(f"n{'='*60}") print(f"Resultados de tiempo de ejecución:") print(f" Media: {statistics.mean(times de ejecución):.4f}s") print(f" Std: {statistics.stdev(times de ejecución):.4f}s") print(f" Min: {min(times de ejecución):.4f}s") print(f" Max: {max(times de ejecución):.4f}s") print(f"nEstadísticas de memoria:") print(f" Bytes asignados (pico): {allocated_peak / 1e9:.3f} GB") print(f" Bytes reservados (pico): {reserved_peak / 1e9:.3f} GB") print(f" Fragmentación (pico): {f_peak / 1e9:.3f} GB ({f_pct:.1f}%)") print(f"{'='*60}n") tamaño_lote = 32 para max_seqlen en [100, 200, 400]: print( f"Generación de evaluación comparativa con tamaño de lote {batch_size} " f"y longitud máxima de secuencia {max_seqlen}…" ) benchmark( lambda: generate_sequence( model, max_seqlen=max_seqlen, lote_size=batch_size ) )
En la siguiente tabla capturamos los resultados para un tamaño de lote de 32 y varias longitudes de secuencia diferentes:
A medida que la longitud de la secuencia se duplica, el tiempo de ejecución se cuadruplica, pareciendo seguir un patrón de escala clásico O(N²). Además, la alta fragmentación de la memoria indica una tensión severa en el asignador de memoria CUDA, lo que puede resultar en fallas frecuentes de la memoria y degradar el rendimiento del tiempo de ejecución. La fragmentación resulta de que cada paso solicita asignaciones de tensor ligeramente más grandes, un patrón que termina dejando múltiples bolsas de memoria inutilizable.
Nuestra primera optimización, el almacenamiento en caché de KV, aborda la complejidad del tiempo de ejecución de nuestro modelo de decodificador.
Almacenamiento en caché de KV
Nuestro ingenuo generador es extremadamente ineficiente: en lugar de almacenar y reutilizar los tensores intermedios de tokens anteriores, recalcula la secuencia completa en cada paso.
Abordamos la ineficiencia informática mediante el uso de almacenamiento en caché KV: almacenamos y reutilizamos los tensores intermedios de clave y valor para tokens anteriores. El almacenamiento en caché de KV reduce la complejidad del tiempo de ejecución de la generación de tokens de O(N²) a O(N).
En el siguiente bloque de código, utilizamos el soporte integrado de la biblioteca de transformadores para el almacenamiento en caché KV para reprogramar nuestra función de generación de tokens para calcular un único lote de tokens en cada paso.
@torch.inference_mode() def generate_sequence(model, max_seqlen, batch_size, use_cache=False): # Inicializa las indicaciones con el token BOS all_tokens = torch.full( (batch_size, 1), config.bos_token_id, dispositivo=DEVICE, dtype=torch.long ) terminado = torch.zeros(batch_size, dispositivo=DEVICE, dtype=torch.bool) # past_key_values se utiliza para almacenar las claves/valores almacenados en caché para cada capa past_key_values = Ninguno para i en el rango(max_seqlen): current_input = ( all_tokens si past_key_values es Ninguno más all_tokens[:, -1:] ) salidas = model( current_input, past_key_values=past_key_values, use_cache=use_cache ) # actualizar caché para el siguiente paso past_key_values = salidas.past_key_values logits = salidas.logits[:, -1, :] new_tokens = torch.argmax(logits, dim=-1) # agregar nuevo token a la secuencia all_tokens = torch.cat( [all_tokens, new_tokens.unsqueeze(-1)], dim=-1 ) terminado |= (new_tokens == config.eos_token_id) stop_gpu = torch.all(finalizado) # comprobando la condición de parada if stop_gpu.item(): print(f"Todas las secuencias terminaron en el paso {i+1}") break return all_tokens
Las cifras de rendimiento resultantes se capturan en la siguiente tabla:
La mejora del rendimiento es profunda y, como era de esperar, aumenta en función de la duración de la secuencia.
Aunque algo mejor que en nuestro experimento inicial, el grado de fragmentación de la memoria sigue siendo una preocupación. Para abordar esto, exploramos dos métodos: asignaciones de memoria expandibles y almacenamiento en caché KV estático.
Asignaciones de memoria CUDA ampliables
Para reducir la fragmentación de la memoria CUDA, programamos PyTorch para que utilice segmentos de memoria expandibles. Al momento de escribir este artículo, esta optimización de la memoria es una característica experimental y debe usarse con precaución. Consulte la documentación de PyTorch para obtener más detalles. Para utilizar la función configuramos la siguiente variable de entorno:
exportar PYTORCH_ALLOC_CONF="segmentos_expandibles: Verdadero"
Volviendo a ejecutar nuestros resultados de referencia en la siguiente tabla:
No sólo vemos una marcada mejora en la fragmentación, sino que también obtenemos una mejora adicional (marginal) en el rendimiento del tiempo de ejecución.
Almacenamiento en caché de KV con StaticCache
El caché predeterminado en HuggingFace es dinámico: crece a medida que aumenta la cantidad de claves y valores durante la generación. HuggingFace admite una caché de tamaño fijo, StaticCache, que preasigna un tamaño de caché máximo para los pares KV y reduce la tensión en el asignador de memoria CUDA. La desventaja de utilizar StaticCache es que la longitud total de la caché participa en el cálculo de la atención en cada paso de generación de tokens, donde los tokens irrelevantes se enmascaran. Esto da como resultado un desperdicio de cálculo que crece con la longitud de la secuencia. Por ejemplo, al generar una secuencia de 400 tokens, el cálculo de atención para cada token se ejecutará en tensores completos de tamaño 400X400.
En el bloque de código siguiente mejoramos nuestro generador de secuencias para admitir el uso de StaticCache:
che: de transformadores import StaticCache @torch.inference_mode() def generate_sequence( model, max_seqlen, batch_size, use_cache=False, use_static_cache=False ): # Inicializa las indicaciones con el token BOS all_tokens = torch.full( (batch_size, 1), config.bos_token_id, dispositivo=DEVICE, dtype=torch.long ) terminado = torch.zeros(batch_size, dispositivo=DEVICE, dtype=torch.bool) # Inicializa el caché estático si se solicita si use_cache y use_static_cache: past_key_values = StaticCache( config=config, max_batch_size=batch_size, max_cache_len=max_seqlen, dispositivo=DEVICE, dtype=model.dtype ) else: past_key_values = Ninguno # Inicializa la posición del caché seguimiento de caché estático cache_positions = torch.arange(max_seqlen, dispositivo=DEVICE) para i in range(max_seqlen): current_input = ( all_tokens if past_key_values es Ninguno más all_tokens[:, -1:] ) cache_position = ( cache_positions[i:i+1] si use_static_cache else Ninguno ) salidas = model( current_input, past_key_values=past_key_values, cache_position=cache_position, use_cache=use_cache ) # actualizar caché para el siguiente paso past_key_values = salidas.past_key_values logits = salidas.logits[:, -1, :] new_tokens = torch.argmax(logits, dim=-1) # agregar nuevo token a la secuencia all_tokens = torch.cat( [all_tokens, new_tokens.unsqueeze(-1)], dim=-1 ) terminado |= (new_tokens == config.eos_token_id) stop_gpu = torch.all(finished) # verificando la condición de parada if stop_gpu.item(): print(f"Todas las secuencias terminaron en el paso {i+1}") break return all_tokens
Los resultados actualizados se capturan a continuación:
El uso de una caché de tamaño fijo mejora enormemente la utilización de la memoria, como lo indica la disminución de la fragmentación de la memoria. Sin embargo, su impacto en el rendimiento en tiempo de ejecución es mixto: para 100 tokens reduce el rendimiento en comparación con un caché dinámico, mientras que para 200 y 400 tokens aumenta el rendimiento en un 9% y un 10%, respectivamente.
Existen métodos más avanzados para implementar la atención que optimizan la utilización de la memoria sin el costo del cálculo desperdiciado. En una publicación anterior, Optimización de modelos de transformadores para secuencias de entrada de longitud variable, cubrimos algunas técnicas de PyTorch para calcular la atención de forma dispersa para reducir el desperdicio de cálculo. Para la configuración de producción, bibliotecas como vLLM utilizan PagedAttention para maximizar la utilización de la memoria. Estos métodos están fuera del alcance de esta publicación.
Para obtener más detalles sobre el almacenamiento en caché en HuggingFace, consulte la descripción general de las estrategias de almacenamiento en caché.
Compilación de modelos
Una de las ventajas documentadas de utilizar una caché de tamaño fijo es que permite aprovechar muchas optimizaciones justo a tiempo (JIT).
En el siguiente bloque de código aplicamos nuestro punto de referencia a una versión compilada con PyTorch de nuestro modelo de decodificador:
tamaño_lote = 32 max_seqlen = 100 modelo = torch.compile(modelo) benchmark( lambda: generate_sequence( modelo, max_seqlen=max_seqlen, tamaño_lote=tamaño_lote, use_cache=True, use_static_cache=True ) )
La compilación del modelo da como resultado un aumento adicional en el rendimiento del tiempo de ejecución, como se muestra en la siguiente tabla:
Tenga en cuenta que también podemos aplicar la compilación de modelos cuando utilizamos el almacenamiento en caché dinámico. Sin embargo, torch.compile proporciona los mejores resultados cuando el gráfico de cálculo está compuesto por tensores de tamaño fijo (por ejemplo, consulte aquí para obtener más detalles).
La penalización por desempeño por detenerse antes de tiempo
Una parte integral de los generadores de tokens comunes es verificar el final de secuencia (EOS) al final de cada paso. Sin esta prueba, los generadores de tokens siempre se ejecutarían para max_seqlen, incluso si todas las secuencias del lote hubieran finalizado. Esto podría resultar en un desperdicio de cálculo considerable y una latencia innecesaria, especialmente cuando las longitudes de secuencias comunes son mucho más cortas que la longitud máxima. En el caso de nuestro experimento con juguetes, esperamos a que finalicen todas las secuencias del lote y dejemos de generar tokens. Las implementaciones de nivel de producción normalmente realizarán procesamiento por lotes continuo, reemplazando secuencias completadas con nuevas indicaciones en la cola de entrada.
terminado |= (new_tokens == config.eos_token_id) stop_gpu = torch.all(finalizado) # verificando la condición de parada if stop_gpu.item(): print(f"Todas las secuencias terminaron en el paso {i+1}") break
Es importante destacar que la llamada .item() al tensor stop_gpu desencadena un evento de bloqueo de sincronización del dispositivo host. Más específicamente, para evaluar la declaración condicional if, la CPU debe esperar a que la GPU complete su cálculo y copie el contenido del tensor en la memoria del host. Mientras la CPU espera, no puede ejecutar el siguiente paso del ciclo de generación de tokens o, más exactamente, no puede cargar los siguientes núcleos de cálculo en la GPU.
Para medir el impacto de la condición de detención en el rendimiento en tiempo de ejecución, agregamos instrumentación para crear perfiles de rendimiento con NVIDIA Nsight™ Systems (nsys) utilizando las API torch.cuda.profiler y nvtx (v0.2.14). (Consulte nuestra publicación reciente para obtener más detalles sobre la creación de perfiles de rendimiento con nsys).
Más detalles sobre la creación de perfiles de rendimiento con nsys). importar nvtx desde torch.cuda importar perfilador @torch.inference_mode() def generate_sequence( model, max_seqlen, batch_size, use_cache=False, use_static_cache=False ): # Inicializar mensajes con el token BOS all_tokens = torch.full( (batch_size, 1), config.bos_token_id, dispositivo=DEVICE, dtype=torch.long ) terminado = torch.zeros(batch_size, dispositivo=DEVICE, dtype=torch.bool) # Inicialice el caché estático si se solicita si use_cache y use_static_cache: past_key_values = StaticCache( config=config, max_batch_size=batch_size, max_cache_len=max_seqlen, dispositivo=DEVICE, dtype=model.dtype ) else: past_key_values = Ninguno # Inicialice el seguimiento de la posición de la caché para caché estática cache_positions = torch.arange(max_seqlen, dispositivo=DEVICE) para i in range(max_seqlen): if i == 30: # iniciar nsys perfilador torch.cuda.synchronize() perfilador.start() elif i == 50: # detener nsys perfilador torch.cuda.synchronize() perfilador.stop() con nvtx.annotate(f"Step {i+1}", color="blue"): with nvtx.annotate("Model Forward", color="green"): current_input = ( all_tokens si past_key_values es Ninguno más all_tokens[:, -1:] ) cache_position = ( cache_positions[i:i+1] si use_static_cache else Ninguno ) salidas = modelo( current_input, valores_clave_pasados=valores_clave_pasados, posición_cache=posición_caché, use_cache=use_cache ) valores_clave_pasados = salidas.valores_clave_pasados logits = salidas.logits[:, -1, :] nuevos_tokens = torch.argmax(logits, dim=-1) all_tokens = torch.cat( [all_tokens, new_tokens.unsqueeze(-1)], dim=-1 ) terminado |= (new_tokens == config.eos_token_id) stop_gpu = torch.all(finalizado) con nvtx.annotate("Comprobar condición de parada", color="red"): # comprobando la condición de parada if stop_gpu.item(): print(f"Todas las secuencias terminaron en el paso {i+1}") break return all_tokens
Ejecutamos nuestro script usando la opción cudaProfilerApi para iniciar y detener el generador de perfiles mediante programación. Consulte la documentación oficial para obtener detalles completos sobre la creación de perfiles desde la CLI de nsys.
perfil nsys –capture-range=cudaProfilerApi –trace=cuda,nvtx,osrt –output=baseline python train.py
El siguiente seguimiento, capturado para un tamaño de lote de 16 y una longitud de secuencia de 100, muestra la GPU inactiva durante aproximadamente 110 microsegundos entre pasos, una eternidad en el contexto de cargas de trabajo de GPU de alto rendimiento. Este es un resultado directo del evento de sincronización desencadenado por la prueba EOS.
En implementaciones de nivel de producción, estos problemas de sincronización se evitan mediante alguna combinación de 1) uso de código de nivel inferior (por ejemplo, C/C++) que evita la limitación del intérprete de Python, 2) uso de gráficos CUDA para reducir la sobrecarga de la carga del kernel, 3) movimiento de comprobaciones condicionales a la GPU utilizando nodos condicionales, y 4) preparación continua y asincrónica de solicitudes posteriores mientras la verificación EOS está en progreso.
En la siguiente sección, demostramos una técnica para ocultar la sobrecarga de la sincronización del dispositivo host en PyTorch utilizando transmisiones CUDA.
Una optimización de flujo CUDA
Una secuencia CUDA es una secuencia lineal de operaciones (kernels, copias de memoria, etc.) que se ejecutan en orden en la GPU. Si bien se garantiza que las operaciones dentro de una única secuencia se ejecutarán secuencialmente, las operaciones en diferentes secuencias pueden ejecutarse simultáneamente o superponerse.
En publicaciones anteriores (por ejemplo, aquí y aquí) demostramos el uso de flujos CUDA para canalizar cargas de trabajo comunes de IA/ML, por ejemplo, ejecutando un modelo en el lote N mientras se prepara el lote N+1. En esta publicación usaremos flujos CUDA para permitir que la CPU cargue los núcleos de GPU del paso N+1 antes de verificar los criterios de detención del paso N. Al contrario de nuestras demostraciones anteriores de flujos CUDA, nuestro ejemplo actual no necesariamente implicará la ejecución simultánea del núcleo de GPU.
Implementamos una función de generación de token alternativa que entrelaza dos flujos CUDA y ejecuta las siguientes operaciones de forma iterativa:
Programe el flujo i%2 para: (A) esperar a que el flujo (i-1)%2 complete su generación del token i-1, (B) usar los tensores actualizados para calcular el token i, (C) ejecutar la prueba EOS para el token i en la GPU y (D) realizar una copia (sin bloqueo) del resultado de la prueba EOS en la memoria fijada en la CPU.
En la secuencia CUDA predeterminada, espere a que la secuencia (i-1)%2 complete la generación del token i-1.
En la transmisión CUDA predeterminada, verifique si se cumplieron los criterios de detención para el token i-1. Si es así, detenga el generador y regrese. De lo contrario, incremente i y regrese al paso 1.
Mientras que anteriormente, la inicialización de la generación del token i estaba bloqueada por la prueba EOS en el token i-1, el uso de flujos CUDA nos permite programar la generación del token i antes de verificar el resultado de la prueba EOS en el token i-1. En la práctica, la prueba EOS para el token i-1 en la CPU se ejecuta mientras la GPU calcula el token i.
@torch.inference_mode() def generate_sequence_pipelined( model, max_seqlen, batch_size, use_cache=False, use_static_cache=False ): # Inicializa las indicaciones con el token BOS all_tokens = torch.full( (batch_size, 1), config.bos_token_id, dispositivo=DEVICE, dtype=torch.long ) terminado = torch.zeros(batch_size, device=DEVICE, dtype=torch.bool) past_key_values = Ninguno # Inicializa el caché estático si se solicita si use_cache y use_static_cache: past_key_values = StaticCache( config=config, max_batch_size=batch_size, max_cache_len=max_seqlen, device=DEVICE, dtype=model.dtype ) # Inicializa el seguimiento de la posición del caché para caché estático cache_positions = torch.arange(max_seqlen, dispositivo=DEVICE) # Flujos duales para canalizar flujos = [torch.cuda.Stream(), torch.cuda.Stream()] stop_host = [ torch.tensor(False, pin_memory=True), torch.tensor(False, pin_memory=True) ] para i in range(max_seqlen): curr_idx, prev_idx = i % 2, (i+1) % 2 curr_s, prev_s = streams[curr_idx], streams[prev_idx] # Lanzar la iteración i en el flujo actual con torch.cuda.stream(curr_s): # programar flujo para esperar a que se complete el flujo anterior curr_s.wait_stream(prev_s) current_input = ( all_tokens si past_key_values es Ninguno más all_tokens[:, -1:] ) cache_position = ( cache_positions[i:i+1] if use_static_cache else Ninguno ) salidas = model( current_input, past_key_values=pasado_key_values, cache_position=cache_position, use_cache=use_cache ) past_key_values = salidas.pasado_key_values logits = salidas.logits[:, -1, :] new_tokens = torch.argmax(logits, dim=-1) all_tokens = torch.cat( [all_tokens, new_tokens.unsqueeze(-1)], dim=-1 ) terminado |= (new_tokens == config.eos_token_id) stop_gpu = torch.all(terminado) stop_host[curr_idx].copy_(stop_gpu, non_blocking=True) # Verificar la señal de parada de la iteración anterior torch.cuda.current_stream().wait_stream(prev_s) if stop_host[prev_idx].item(): print(f"Todas las secuencias terminaron en el paso {i}") break return all_tokens
La siguiente imagen captura el seguimiento de nsys para nuestro nuevo generador de tokens:
En la sección CUDA del seguimiento podemos ver el uso de dos flujos CUDA, con la generación de tokens pasando de un lado a otro en una especie de efecto ping-pong: un flujo genera todos los tokens impares y el segundo todos los tokens pares. La CPU está aproximadamente medio paso por delante de la GPU, lo que le permite programar el paso i mientras la GPU calcula el paso i-1. La verificación de parada de EOS del lado de la CPU del paso i-1 (en rojo) ocurre después de que el paso i esté completamente programado (y haya comenzado a ejecutarse). Lo más importante es que ahora encontramos que la utilización de la GPU es consistente: el ralentí que veíamos antes desapareció.
El entrelazado de flujo CUDA da como resultado un aumento adicional del rendimiento, como se muestra en la siguiente tabla:
Es de esperar que el beneficio de la solución de ping-pong que hemos implementado se vea afectado por la relación entre el tiempo de inactividad de la GPU (es decir, la sobrecarga de la carga del kernel) y el tiempo de cálculo del kernel. Para probar esto, fijamos la longitud de la secuencia en 100 y volvemos a ejecutar el punto de referencia para varios tamaños de lote:
Como se esperaba, la mayor ganancia de rendimiento, 11,6%, se produce cuando el tamaño del lote es más pequeño y la carga de cálculo del kernel es la más baja. A medida que aumenta el cálculo del núcleo, la relación entre la carga del núcleo y el tiempo de cálculo del núcleo disminuye, al igual que el impacto del entrelazado de flujos CUDA.
Tenga en cuenta que el uso de transmisiones CUDA conlleva cierta sobrecarga. Esto se puede demostrar comparando nuestra solución de entrelazado con un generador de tokens que omite por completo la prueba EOS:
Los posibles problemas de rendimiento del uso de CUDA Streams
Las transmisiones CUDA deben usarse con extrema precaución. Cuando usamos la transmisión predeterminada, podemos confiar en PyTorch para realizar cualquier sincronización necesaria cuando se mueven datos. Sin embargo, cuando utilizamos transmisiones CUDA, debemos garantizar explícitamente una sincronización adecuada. En particular, debemos garantizar una transferencia de datos adecuada entre las transmisiones. De lo contrario, podemos experimentar errores CUDA (por ejemplo, "afirmación del lado del dispositivo activada"), si tenemos suerte. Si tenemos menos suerte, podemos experimentar corrupción de datos sin siquiera saberlo. Consulte la documentación de la transmisión PyTorch CUDA para obtener más detalles sobre el uso apropiado.
Para cargas de trabajo de IA/ML con una gran utilización de memoria CUDA, como los LLM, otra consideración es la utilización de la memoria. El asignador de almacenamiento en caché de PyTorch administra la memoria por flujo; El uso de múltiples flujos puede provocar una mayor reserva de memoria y fragmentación. Esto podría dar como resultado un aumento de fallas de memoria que podrían eclipsar las ganancias potenciales del uso de transmisiones.
Resultados
En la siguiente tabla, resumimos los resultados en tiempo de ejecución de la aplicación de almacenamiento en caché estático, compilación y canalización en un lote de 32 secuencias y una longitud máxima de secuencia de 100. Los resultados están ordenados en orden creciente de rendimiento:
En el caso de nuestro modelo de juguete GPT-2, los mejores resultados (casi 5 veces el rendimiento básico) se logran cuando se emplea la compilación de PyTorch y el método de entrelazado de flujo CUDA que se analiza en esta publicación. Sin embargo, como hemos visto, el impacto del entrelazado CUDA podría variar mucho según las propiedades de la carga de trabajo y el entorno de ejecución, particularmente en la relación entre el tiempo de carga del kernel y el tiempo de cómputo del kernel. Asegúrese de ejecutar sus propios puntos de referencia antes de adoptar este método.
Resumen
En la ingeniería de IA de alto rendimiento, cualquier indicio de subutilización de la GPU presenta una oportunidad de optimización. Una de las principales herramientas de optimización de las GPU NVIDIA son las transmisiones CUDA. En esta publicación, demostramos su uso para resolver el tiempo de inactividad de la GPU que resulta de la sincronización del dispositivo host asociada con la detención anticipada en la generación de tokens autorregresivos nativos de PyTorch. Al entrelazar transmisiones CUDA en un patrón de "ping-pong", logramos ocultar con éxito la latencia impuesta por la verificación EOS, lo que resultó en un aumento significativo del rendimiento de la carga de trabajo. Al combinar esta técnica con los métodos conocidos de compilación de modelos y almacenamiento en caché estático, podemos maximizar el rendimiento de la inferencia nativa de PyTorch.