Introducción
exige modelos y datos a gran escala, lo que lleva el hardware informático al límite. Ya sea que esté entrenando modelos en imágenes complejas, procesando documentos de contexto extenso o ejecutando entornos de aprendizaje reforzado de alto rendimiento, maximizar la eficiencia de su GPU es fundamental. No es raro entrenar o ejecutar inferencias en modelos con miles de millones de parámetros en terabytes de datos. Una configuración no optimizada puede convertir un experimento rápido en horas o días de espera.
Cuando el entrenamiento o la inferencia son lentos, nuestro instinto suele ser culpar al tamaño del modelo o a la complejidad matemática. Las GPU modernas son calculadoras rápidas, pero dependen de la CPU para asignar el trabajo y de la ubicación de almacenamiento de datos en el dispositivo en la GPU. Por lo general, el cálculo en la GPU no es el cuello de botella. Si su CPU tiene dificultades para cargar, preprocesar y transferir sus lotes a través del puente PCIe, su GPU permanece inactiva, hambrienta de datos.
¿La buena noticia? No es necesario escribir kernels CUDA personalizados ni depurar código GPU de bajo nivel para solucionarlo. Si es un investigador, ingeniero o aficionado de ML interesado en optimizar los canales de GPU, ¡este blog es para usted! En esta publicación, exploramos la mecánica de este cuello de botella y analizamos decisiones de ingeniería viables para maximizar la utilización de la GPU. Cubriremos todo, desde ajustes fundamentales de la canalización de PyTorch hasta optimizaciones de hardware más avanzadas e integraciones de Hugging Face.
💡 Nota
Asumiremos un conocimiento práctico básico de Python y PyTorch DataLoaders en el futuro. No se requiere un conocimiento profundo de la arquitectura de la GPU, ya que proporcionaremos una descripción general de alto nivel de la GPU y cómo funciona. Todas las técnicas discutidas serán aplicables al entrenamiento y la inferencia a menos que se indique explícitamente.
Descripción general de la GPU
Para empezar, las unidades de procesamiento de gráficos (GPU) explotaron en popularidad junto con el aprendizaje profundo debido a su capacidad para entrenar y ejecutar modelos a la velocidad del rayo con operaciones paralelizables. Pero, ¿qué hace realmente una GPU? Antes de optimizar nuestro código, necesitamos un modelo mental compartido de lo que sucede bajo el capó de una GPU, las diferencias con una CPU y el flujo de datos entre las dos.
¿Qué es una GPU? ¿En qué se diferencia de una CPU?
Las GPU no superan universalmente a las CPU. Las CPU están diseñadas para resolver problemas altamente secuenciales con baja latencia y ramificaciones complejas (cualidades ideales para ejecutar un sistema operativo). Alternativamente, las GPU constan de miles de núcleos optimizados para completar operaciones básicas en paralelo.[1]. Mientras que una CPU necesitaría procesar mil multiplicaciones de matrices de forma secuencial (o hasta el límite de sus núcleos), una GPU puede ejecutar todas estas operaciones en paralelo (en una fracción de segundo). En el aprendizaje automático, rara vez nos ocupamos de problemas altamente secuenciales que requieren una CPU. La mayoría de las operaciones son multiplicación de matrices, una tarea altamente paralelizable.
En un nivel alto, una GPU consta de miles de pequeños núcleos de procesamiento agrupados en multiprocesadores de transmisión (SM) diseñados para computación paralela masiva. Los SM gestionan, programan y ejecutan cientos de subprocesos al mismo tiempo. Un grupo de memoria de alto ancho de banda, Video RAM (VRAM), rodea las unidades de cómputo junto con cachés ultrarrápidas que contienen datos temporalmente para un acceso rápido. VRAM es el almacén principal donde residen los pesos, gradientes y lotes de datos entrantes de su modelo. Las CPU y las GPU se comunican entre sí a través de un puente de interfaz, que analizaremos más en profundidad a continuación, ya que es el principal puente donde se producen los cuellos de botella.
💡 Nota
Si está utilizando GPU NVIDIA, hay otro componente dentro de la GPU llamado Tensor Cores, que acelera las matemáticas matriciales de precisión mixta utilizadas en el aprendizaje automático. Volverán a surgir cuando hablemos de Precisión mixta.
Puente PCIe
Como acabamos de mencionar, los datos viajan desde la CPU a la GPU a través de un puente de interfaz llamado Peripheral Component Interconnect Express (PCIe). Los datos se originan en su disco, se cargan en la RAM de la CPU y luego cruzan el bus PCIe para llegar a la VRAM de la GPU. Cada vez que envía un tensor de PyTorch al dispositivo usando .to('cuda'), está invocando una transferencia a través de este puente. Si su CPU envía constantemente pequeños tensores uno por uno en lugar de bloques grandes y contiguos, rápidamente obstruirá el puente con latencia y sobrecarga.
¿Qué es la utilización de GPU?
Ahora que hemos cubierto la anatomía de la GPU, debemos comprender las métricas que estamos rastreando. Cuando utilizamos nvidia-smi, Weights and Biases, PyTorch Profiler, NVIDIA Nsight Systems o cualquier otro método de seguimiento de GPU, generalmente analizamos dos porcentajes principales: uso de memoria y utilidad de GPU volátil.
Uso de memoria (VRAM): VRAM es la memoria física de la GPU. Su VRAM puede estar al 100% de su capacidad mientras su GPU no hace nada. Un uso elevado de VRAM solo significa que ha cargado correctamente los pesos, gradientes y un lote de datos de su modelo en la memoria física de la GPU. GPU-Util volátil (utilización informática): esta es la métrica crucial. Mide el porcentaje de tiempo durante el último período de muestra (generalmente 1 segundo) en el que los núcleos informáticos de la GPU estuvieron ejecutando instrucciones activamente. ¡El objetivo es maximizar constantemente este porcentaje!
Cuello de botella CPU-GPU
Ahora que hemos cubierto las CPU y las GPU, veamos cómo se produce el cuello de botella CPU-GPU y qué podemos hacer para solucionarlo. La GPU tiene miles de núcleos listos para paralelizar operaciones, pero necesita de la CPU para delegar tareas. Cuando entrenas un modelo, tu GPU no puede leer directamente desde tu SSD. La CPU debe cargar y decodificar los datos sin procesar, aplicar aumentos, agruparlos y entregarlos. Si su CPU tarda 50 milisegundos en preparar un lote y su GPU solo tarda 10 milisegundos en calcular los pasos hacia adelante y hacia atrás, su GPU pasa 40 milisegundos inactiva.
Modelo de línea de techo
Este problema se formaliza en el modelo de línea de techo. Mide el rendimiento (FLOP/segundo) frente a la intensidad aritmética (FLOP/byte), siendo los FLOP operaciones de punto flotante por segundo. Cuando la intensidad aritmética es baja (cargas una gran cantidad de datos pero haces muy pocos cálculos con ellos), te topas con el techo inclinado de la “memoria ligada”. Cuando la intensidad aritmética es alta (cargas una pequeña cantidad de datos pero haces una enorme cantidad de multiplicaciones de matrices con ellos), llegas al techo plano "compute-bound".
El paralelismo de GPU rara vez es un obstáculo para los experimentos de investigación. Normalmente se producen ralentizaciones en el régimen de memoria: análisis de datos de la CPU, obstrucción del bus PCle o límites de ancho de banda de VRAM. La clave para esto es casi siempre una mejor gestión del flujo de datos.
Optimización del canal de datos
Seguimiento del uso de GPU
Antes de que podamos optimizar la canalización de datos, debemos comprender cómo monitorear la utilización de GPU y VRAM. La forma más sencilla es utilizar nvidia-smi para obtener una tabla con todas las GPU disponibles, la VRAM actual y la utilización de GPU volátil.
Con watch -n 1 nvidia-smi, las métricas se pueden monitorear y actualizar cada segundo. Sin embargo, la mejor manera de obtener métricas de GPU más detalladas es utilizar PyTorch Profiler o Weights and Biases. NVIDIA Nsight Systems también es una gran herramienta de monitoreo que es relativamente rápida de configurar aquí. Weights and Biases ofrece la visualización más sencilla de los gráficos de utilización de GPU para nuestros propósitos, y estos gráficos son la forma más sencilla de diagnosticar una optimización deficiente de la GPU. Aquí se muestra un ejemplo de configuración para Weights and Biases (tomado directamente de la documentación de Weights and Biases).[2]):
import wandb # Proyecto en el que se registra la ejecución project = "my-awesome-project" # Diccionario con hiperparámetros config = {"epochs": 1337, "lr": 3e-4} # La sintaxis `with` marca la ejecución como finalizada al salir del bloque `with`, # y marca la ejecución como "fallida" si hay una excepción. # # En un cuaderno, puede ser más conveniente escribir `run = wandb.init()` # y llamar manualmente a `run.finish()` en lugar de usar un bloque `with`. con wandb.init(project=project, config=config) como ejecución: # Código de entrenamiento aquí # Registrar valores en W&B con run.log() run.log({"accuracy": 0.9, "loss": 0.1})
El indicador más sencillo de una canalización de GPU no optimizada es un gráfico de utilización de GPU en forma de diente de sierra. Aquí es donde la utilización de la GPU está inactiva al 0%, aumenta brevemente al 100% y luego vuelve a estar inactiva al 0%, lo que significa un problema de cuello de botella entre la CPU y la GPU. Alcanzar una utilización periódica del 100 % no es una señal de que la utilización de la GPU esté maximizada. La GPU analiza los datos disponibles en una fracción de segundo y los valles del 0% representan la espera a que la CPU prepare el siguiente lote. El objetivo es la utilización continua: una línea plana e ininterrumpida cercana al 100%, lo que significa que la GPU nunca tiene que esperar. A continuación se muestra un ejemplo de un gráfico de utilización de GPU de diente de sierra en el mismo formato que Pesos y sesgos:
Veamos por qué sucede esto en un PyTorch DataLoader básico. De forma predeterminada, un PyTorch DataLoader podría definirse de la siguiente manera:
Cargador de datos (conjunto de datos, tamaño de lote = 32, reproducción aleatoria = Verdadero, num_workers = 0, pin_memory = Falso)
Con num_workers=0 y pin_memory=False (los valores predeterminados en un DataLoader), el proceso principal de Python tiene que hacer todo de forma secuencial:
Obtenga los archivos del disco. Aplicar aumentos de imágenes o preprocesar texto. Mueva el lote a la GPU. La GPU calcula los pases hacia adelante y hacia atrás.
Este es el peor de los casos para la utilización de GPU. Para los pasos 1 a 3, la utilización de la GPU es del 0%. Cuando se completa el paso 3, la utilización de la GPU aumenta al 100% y luego los pasos 1 y 2 se repiten para el siguiente lote.
Las siguientes secciones analizan cómo optimizar la canalización de datos.
num_workers (Paralelizando la CPU)
La solución más impactante es paralelizar la preparación de datos. Al aumentar num_workers, le indica a PyTorch que genere subprocesos dedicados para la recuperación y preparación por lotes en segundo plano mientras la GPU calcula. Sin embargo, más trabajadores no siempre significa más velocidad. Si tiene una CPU de 8 núcleos y configura num_workers=16, ralentizará su entrenamiento debido a la sobrecarga del cambio de contexto y la comunicación entre procesos (IPC). Cada trabajador crea una copia del conjunto de datos en la memoria. Demasiados trabajadores pueden provocar daños en la memoria y bloquear el sistema. Una buena regla general es comenzar en num_workers=4 y crear perfiles desde allí.
💡 Nota
El número óptimo de num_workers no solucionará una implementación lenta del conjunto de datos. Para mantener __getitem__ eficiente, evite crear instancias de objetos, conexiones de base de datos por elemento o un preprocesamiento intenso en la llamada a la función. Su única tarea es recuperar bytes sin procesar, convertirlos en un tensor y regresar.
pin_memory=True (Transferencia de datos optimizada)
Incluso con trabajadores en segundo plano, es importante cómo se transfieren físicamente los datos a través del puente PCIe. Normalmente los tensores no van directamente a la GPU cuando los transfieres. Primero se lee del disco a la RAM del sistema paginado y la CPU lo copia en un área especial de RAM no paginada (o con página bloqueada) antes de cruzar el bus PCIe a la GPU VRAM.
Configurar pin_memory=True crea una vía rápida de datos. Le indica a su DataLoader que asigne lotes directamente en la memoria de página bloqueada. Esto permite que la GPU utilice el acceso directo a la memoria (DMA) para extraer los datos directamente a través del puente sin que la CPU tenga que actuar como intermediario para la transferencia final, lo que reduce significativamente la latencia.
pin_memory=True viene con una compensación de hardware. Con la memoria de página bloqueada, el sistema operativo no puede intercambiar esta RAM al disco duro si se queda sin espacio. Normalmente, los datos se pueden intercambiar al disco cuando no hay memoria, pero si te quedas sin RAM con página bloqueada, se generará un error de Sin memoria. Si obtiene un OOM en su secuencia de comandos, asegúrese de investigar primero el indicador pin_memory antes de realizar una depuración más compleja. Además, tenga cuidado al combinar pin_memory=True con un recuento alto de num_workers. Debido a que cada proceso de trabajo genera y mantiene activamente lotes en la memoria, esto puede inflar rápidamente la huella de RAM bloqueada de su sistema.
prefetch_factor (poner en cola datos)
A veces el cuello de botella no es la potencia de procesamiento de la CPU, sino el disco mismo. Al leer miles de archivos desde una unidad de red, pueden producirse picos repentinos en la latencia de E/S sin que sea culpa del usuario.
El argumento prefetch_factor dicta cuántos lotes debe preparar y mantener cada trabajador en una cola en la CPU por adelantado. Si tiene 4 trabajadores y un prefetch_factor=2, la CPU siempre intentará mantener en cola 8 lotes listos para usar. Si el disco de repente se cuelga durante medio segundo en un archivo corrupto, la GPU no morirá de hambre: simplemente sale de la cola de captación previa mientras el trabajador se pone al día.
💡 Nota
Asegúrese de no configurar prefetch_factor demasiado alto, ya que puede hacer que la GPU espere a que la CPU se ponga al día. Una buena regla general es establecer prefetch_factor en 2 o 3. También puede causar errores de CUDA sin memoria si lo configura demasiado alto.
Al ajustar estos parámetros, puede suavizar la utilización del diente de sierra en una curva de utilización alta y continua. Este es un ejemplo de a lo que debería aspirar:
La utilización de la GPU es ahora una línea alta y continua cercana al 100 %, lo que significa que la GPU nunca tiene que esperar. Solo con los parámetros de carga de datos, pudimos pasar de una utilización ineficiente de la GPU a una utilización efectiva y continua.
Computación y memoria en la GPU
Una vez que se optimiza el DataLoader, los datos vuelan a través del puente PCIe, lo que ya no crea el cuello de botella en la utilización de la GPU en forma de diente de sierra. Pero una vez que los datos pasan a la GPU VRAM, ¿cómo nos aseguramos de que se utilicen de manera eficiente?
Tamaño del lote
Revisemos el modelo de línea del techo. Para escapar del techo inclinado "limitado a la memoria" y alcanzar el máximo rendimiento plano "limitado a la computación" de su GPU, necesita una alta intensidad aritmética. La forma más sencilla de aumentar la intensidad aritmética es aumentar el tamaño del lote. Cargar una única matriz masiva de tamaño 1024×1024 y hacer los cálculos todos a la vez es más eficiente para los multiprocesadores de transmisión de la GPU que cargar 32 matrices más pequeñas secuencialmente.
En teoría, deberíamos cargar todos nuestros datos a la vez y tener un lote singular, pero en la práctica, esto termina en el temido error CUDA sin memoria. Esto significa que estás intentando cargar más datos en la VRAM, pero no hay más espacio para asignar en la GPU.
💡 Nota
Quizás se pregunte por qué el tamaño del lote, el número de trabajadores o cualquier métrica de aprendizaje profundo suele ser una potencia de 2. No es solo una regla de facto, sino el resultado del diseño del hardware de NVIDIA.
Dentro de un SM, una GPU no ejecuta subprocesos individualmente. Los agrupa en unidades llamadas Warps y, en las GPU de NVIDIA, un warp siempre contiene exactamente 32 subprocesos. Más allá del límite de deformación de 32 subprocesos, los controladores de memoria física de una GPU obtienen datos de la VRAM en fragmentos de potencia de 2 bytes. Si las dimensiones de su tensor no están alineadas con estos fragmentos, la GPU tiene que realizar múltiples búsquedas de memoria para capturar los datos desbordados.
Para obtener la máxima eficiencia, elija múltiplos de 32 o 64 (u 8 si tiene limitaciones de memoria) para el tamaño del lote y potencias de 2 para otras métricas de aprendizaje profundo.
Precisión mixta
De forma predeterminada, PyTorch inicializa todos los pesos y datos del modelo en FP32 (punto flotante de 32 bits). Para casi todas las tareas de aprendizaje profundo, esto es excesivo. La solución es la cuantificación del modelo, reduciendo los tensores a FP16 o BF16 (16 bits). ¿Por qué es importante esto para la utilización?
Reduce a la mitad el cuello de botella de la memoria: solo la mitad de bytes se mueven a través del puente PCIe y dentro de la VRAM interna de la GPU. Desbloquea el hardware: las GPU NVIDIA modernas poseen silicio especializado llamado Tensor Cores. Estos núcleos quedan completamente inactivos si les pasas matemáticas FP32. Están diseñados específicamente para ejecutar multiplicaciones de matrices de 16 bits a altas velocidades.
Antes de transmitir a FP16, asegúrese de que el rendimiento sea el mismo con un conjunto de datos submuestreado para confirmar que la tarea en cuestión no requiere FP32. Otra opción es utilizar torch.autocast de PyTorch. Esta función incorporada envuelve el paso hacia adelante en un administrador de contexto que determina automáticamente qué operaciones son seguras para convertir a 16 bits (como multiplicaciones de matrices) y cuáles deben permanecer en 32 bits para lograr estabilidad numérica (como Softmax o LayerNorm). Es esencialmente una aceleración 2x gratuita.
Sin embargo, FP16 o BF16 no siempre son el mejor método de cuantificación. En las arquitecturas NVIDIA modernas (A100 o H100), se debe usar BF16 en lugar de FP16 para evitar pérdidas de NaN o desbordamiento de gradiente. Otra buena opción es el formato TF32 (TensorFloat-32), propiedad de NVIDIA, que es un formato de punto flotante de 19 bits que mantiene la precisión de FP32 con una velocidad 10 veces mayor que FP32 en A100 y H100.[3].
Acumulación de gradiente (solo entrenamiento)
Al entrenar un modelo, en lugar de intentar forzar un lote de gran tamaño en la VRAM y fallar, utilice un "microlote" más pequeño. Cuando se utilizan estrategias de "microlotes", en lugar de actualizar los pesos del modelo inmediatamente (llamando a optimizador.step()), se acumulan los gradientes (loss.backward()) en varios pases consecutivos hacia adelante, creando un "tamaño de lote efectivo". Por ejemplo, un tamaño de lote de 8 con 8 pasos de acumulación de gradiente produce la misma actualización matemática que un tamaño de lote de 64 con una sola actualización. Las estrategias de microlotes pueden estabilizar el entrenamiento sin aumentar la huella de VRAM.
Eficiencia del núcleo
El concepto de eficiencia final se refiere a la eficiencia del núcleo en el entrenamiento y la inferencia. Esta es una exploración para comprender cómo funcionan los kernels CUDA dentro de una GPU para aquellos que desean trabajar con arquitecturas personalizadas. PyTorch 2.0+ (que casi siempre se usa), abstrae esto del usuario con un comando simple que se muestra a continuación. Comencemos con una inmersión más profunda en la GPU y cómo funcionan los núcleos:
Cada vez que ejecutas una operación en PyTorch (usaremos d=a+b+c como ejemplo simple), estás lanzando un "núcleo" en la GPU. Al abstraer las complejidades de una GPU, no puede hacer cálculos de una sola vez. La GPU debe:
Lea a y b de la VRAM en la memoria caché del SM. Calcular a+b. Escribe ese resultado intermedio en la VRAM. Lea ese resultado intermedio y c de VRAM. Calcula la suma final. Escribe d de nuevo en VRAM.
Esta es la sobrecarga del kernel. Al crear arquitecturas personalizadas desde cero, es fácil crear accidentalmente cientos de pequeñas lecturas y escrituras secuenciales. Los núcleos de GPU pasan todo el tiempo esperando en la VRAM interna en lugar de hacer cálculos. La fusión de núcleos CUDA personalizados puede reducir la sobrecarga, pero afortunadamente, PyTorch 2.0+ maneja esto implícitamente con torch.compile(). PyTorch analiza todo el gráfico computacional y utiliza Triton de OpenAI para escribir automáticamente núcleos fusionados altamente optimizados que pueden ahorrar horas de una larga ejecución de entrenamiento al acortar los viajes de ida y vuelta de la memoria.
Si bien torch.compile() es fenomenal para la fusión automática de operaciones de propósito general, a veces para obtener ganancias de rendimiento se requieren núcleos altamente especializados. Históricamente, integrar núcleos CUDA o Triton escritos a mano en su investigación significaba luchar con complejos sistemas de compilación C++ y versiones coincidentes del kit de herramientas CUDA. Afortunadamente, la biblioteca de núcleos de Hugging Face[4]trata las operaciones informáticas de bajo nivel como modelos previamente entrenados. En lugar de compilar desde el código fuente, puede recuperar archivos binarios precompilados y optimizados para hardware directamente desde el Hub con una sola llamada a una función de Python. La biblioteca detecta automáticamente su entorno exacto de PyTorch y GPU y descarga la combinación perfecta en segundos.
A continuación se muestra un ejemplo sencillo del uso de la biblioteca de kernels de Hugging Face (de la documentación de la biblioteca de kernels de Hugging Face[4]):
importar antorcha desde núcleos importar get_kernel # Descargar núcleos optimizados desde el centro Hugging Face activación = get_kernel("kernels-community/activation", versión=1) # Tensor aleatorio x = torch.randn((10, 10), dtype=torch.float16, dispositivo="cuda") # Ejecutar el núcleo y = torch.empty_like(x) activación.gelu_fast(y, x) print(y)
Transformadores de cara abrazada
Como comentario breve, la biblioteca de transformadores de Hugging Face[5]incluye toda la funcionalidad para optimizar su modelo para la GPU que hemos discutido a través de sus clases TrainingArguments y Trainer. Para lograr una mayor abstracción, simplemente puede utilizar el siguiente ejemplo para ejecutar un modelo a través de Hugging Face.
de transformadores import TrainingArguments, Trainer Training_args = TrainingArguments( # Especificar un directorio de salida output_dir="./results", # Data Pipeline dataloader_num_workers=4, dataloader_pin_memory=True, dataloader_prefetch_factor=2, # Computación y memoria per_device_train_batch_size=8, gradient_accumulation_steps=4, # Precisión y hardware bf16=Verdadero, tf32=Verdadero, # Eficiencia del kernel torch_compile=Verdadero, ) entrenador = Entrenador(modelo=modelo, args=training_args, train_dataset=train_dataset) trainer.train()
Conclusión
La optimización de su canal de GPU se reduce a dos principios básicos: mantener la GPU constantemente cargada con datos y hacer que cada operación cuente una vez que llegan los datos.
En el lado de la canalización de datos, ajustar el DataLoader aumentando num_workers, habilitando pin_memory y configurando un prefetch_factor produce una utilización de GPU más continua. En el lado de la computación, maximizar el tamaño del lote (o utilizar la acumulación de gradiente), bajar a precisión mixta (FP16/BF16 o TF32) y fusionar operaciones a través de torch.compile() o la biblioteca de kernels Hugging Face reduce drásticamente el tráfico de VRAM y la sobrecarga del kernel.
En conjunto, estos ajustes convierten horas de tiempo de inactividad desperdiciado y ligado a la memoria en un proceso de investigación de alta velocidad y totalmente utilizado.
Referencias
[1]Diseño de CPU frente a GPU: Guía de programación de NVIDIA Cuda
[2]Configuración de pesos y sesgos – Pesos y sesgos Github
[3]Formato de precisión TensorFloat-32 — Blog de NVIDIA
[4]Biblioteca de núcleos de Hugging Face – Hugging Face Github
[5]Biblioteca de transformadores Hugging Face – Hugging Face Github