Resumen Técnico: Análisis de Rendimiento y Optimización de Modelos de Difusión Generativa 3D a través de Arquitecturas de GPU

Planteamiento del Problema

La síntesis de imágenes médicas 3D de alta fidelidad, particularmente para RM, depende fuertemente de los Modelos Probabilísticos de Difusión de Eliminación de Ruido (DDPM, por sus siglas en inglés), como Med-DDPM, debido a su superior fidelidad estructural y estabilidad de entrenamiento en comparación con las GAN. Sin embargo, el despliegue de estos modelos está severamente limitado por costos computacionales inmensos. Un solo muestreo de difusión 3D requiere cientos de evaluaciones secuenciales de la U-Net, cada una de las cuales involucra convoluciones 3D complejas, capas de normalización y operaciones elemento a elemento. Esto resulta en un patrón de ejecución de GPU altamente heterogéneo caracterizado por lanzamientos frecuentes de kernels y fases alternantes de dependencia de cómputo (compute-bound) y de memoria (memory-bound), además de demandas significativas de recursos.

La investigación actual en imágenes médicas evalúa típicamente los modelos de difusión únicamente mediante métricas de calidad de imagen, careciendo de una perspectiva a nivel de sistemas. En consecuencia, los cuellos de botella arquitectónicos específicos —que van desde los patrones de acceso a memoria y conversiones de disposición de tensores hasta la utilización subóptima de los pipelines de Tensor Core— siguen sin comprenderse del todo. Este artículo aborda esta brecha investigando por qué el entrenamiento de modelos de difusión 3D exige cientos de horas de GPU incluso para conjuntos de datos modestos, y cómo las arquitecturas de hardware actuales (Volta, Ampere, Hopper) gestionan estas cargas de trabajo específicas.

Metodología

Los autores realizaron un análisis de rendimiento integral de "full-stack" del marco de trabajo Med-DDPM de última generación a través de tres generaciones de GPUs NVIDIA:

• Hardware: NVIDIA V100 (Volta), A100 (Ampere) y H100 (Hopper).
• Stack de Software: PyTorch 2.0.1, CUDA 12.9, cuDNN 8.9.2 y el profiler Nsight Compute (ncu).
• Configuración Experimental: El entrenamiento se ejecutó con hiperparámetros idénticos (25 épocas, tamaño de lote 1, tamaño de entrada 128³) en todas las plataformas. Para asegurar cargas de trabajo bit-idénticas, los autores utilizaron entornos contenedorizados (Docker en Indiana Jetstream2/TACC Stampede3; Singularity/Apptainx en SDSC Expanse) para eliminar la variabilidad de las librerías del sistema.
• Estrategia de Perfilado: El estudio utilizó Nsight Compute para extraer métricas microarquitectónicas, incluyendo:
  • Desgloses de tiempo de ejecución a nivel de kernel y mezclas de instrucciones.
  • Perfiles de instrucciones dinámicas (FP, INT, carga/almacenamiento, control).
  • Estimaciones de puntuación de prioridad (combinando la duración del kernel y la mejora de velocidad estimada).
  • Utilización a nivel de SM, ocupación de warp y desgloses de tipos de estancamiento (dependencias de larga latencia, estrangulamiento del pipeline matemático, ramas divergentes).
  • Comportamiento de la jerarquía de memoria (tasas de acierto de L1/L2, utilización del ancho de banda de DRAM).

Contribuciones Clave

1. Caracterización a Nivel de Sistema: El artículo presenta la primera caracterización microarquitectónica detallada de Med-DDPM a través de las arquitecturas Volta, Ampere y Hopper, revelando que el entrenamiento está dominado abrumadoramente por kernels de convolución 3D de cuDNN e implícitos-GEMM.
2. Identificación de Cuellos de Botella: A través de los desgloses de mezcla de kernels y el análisis de la puntuación de prioridad, los autores identifican que las ineficiencias surgen no solo de las irregularidades en el acceso a la memoria, sino también de la fragmentación de la programación de warps, los efectos de residencia en caché y la limitada utilización de Tensor Cores en las arquitecturas Ampere/Hopper.
3. Evaluación de Optimización: El estudio evalúa dos optimizaciones conscientes de la arquitectura:
   • Activación de TF32 Tensor Core: Habilitar la precisión TF32 para la multiplicación de matrices y convoluciones 3D en Ampere y Hopper.
   • Disposición de Canales al Final (Channels-Last Layout): Reestructurar la disposición de la memoria de NCHW a NHWC para mejorar la localidad de los datos para los tensor cores.
4. Perspectivas de Rendimiento: El artículo demuestra que, si bien la activación de TF32 reduce significativamente los ciclos de SM y las instrucciones dinámicas, la disposición de canales al final cambia fundamentalmente la carga de trabajo a un régimen limitado por la memoria, dominado por pequeños kernels elemento a elemento, exponiendo nuevos cuellos de botella relacionados con la fusión de kernels y la sobrecarga de lanzamiento.

Resultados

La evaluación comparó cuatro configuraciones: Baseline, OPT1 (solo TF32), OPT2 (solo Channels-Last) y OPT12 (Combinado).

• Activación de TF32 (OPT1):

  • Eficiencia: En A100 y H100, habilitar TF32 redujo los ciclos de SM hasta en 5× y las instrucciones dinámicas hasta en 100× en comparación con el baseline.
  • Utilización: La utilización de Tensor Core aumentó significativamente de ~1.45× a 9.98× en A100. Por el contrario, la utilización de los núcleos CUDA disminuyó a medida que la carga de trabajo se desplazaba hacia instrucciones densas de MMA (Multiplicación-Acumulación de Matrices).
  • Memoria: Aunque las tasas de acierto de la caché L1 disminuyeron (debido a que los bloques de memoria global contiguos y más grandes fluyen a través de los SMs), la utilización del ancho de banda de DRAM aumentó, lo que indica que el camino de Tensor Core extrae datos agresivamente para alimentar los pipelines.
  • Estancamientos (Stalls): En A100, todas las categorías de estancamiento disminuyeron significativamente. En H100, el estrangulamiento del pipeline matemático disminuyó, pero los estancamientos por dependencias de larga latencia aumentaron, reflejando las cadenas de dependencia más profundas de la disponibilidad de operandos de Tensor Core.

• Disposición de Canales al Final (OPT2):

  • Cambio de Kernel: Esta optimización colapsó el dominio de los grandes kernels de conv/GEMM, desplazando la ejecución casi por completo a pequeños kernels elemento a elemento y de transformación de disposición (ej. padding, conversiones NHWC↔NCHW).
  • Rendimiento: Aunque esto redujo drásticamente los ciclos totales y las instrucciones (hasta un 0.01×–0.02× del baseline), resultó en motores de cómputo subutilizados. La utilización de Tensor Core se mantuvo cerca de los niveles del baseline (~1.5%) y el rendimiento de SM cayó aproximadamente un 50%.
  • Memoria: Las tasas de acierto de la caché L2 aumentaron drásticamente (hasta un 85%) y la utilización del ancho de banda de DRAM colapsó, confirmando que la carga de trabajo se volvió residente en caché pero carecía de la intensidad aritmética para saturar el hardware.

• Combinado (OPT12):

  • La combinación restauró parcialmente el uso de Tensor Core en A100, pero no logró alcanzar el rendimiento máximo en H100 o V100, donde la carga de trabajo siguió dominada por elementos a elemento y transformaciones de disposición. Los autores señalan que se requiere una mayor fusión de kernels para alcanzar el pleno potencial de esta disposición.

Significado y Reivindicaciones

El artículo sostiene que la principal barrera para escalar los modelos de difusión médica 3D no es meramente algorítmica, sino sistémica, arraigada en el desajuste entre las cargas de trabajo de difusión y las microarquitecturas de GPU.

• Realidad Arquitectónica: El estudio establece que el simple hecho de pasar a nuevas generaciones de GPU (ej. de V100 a H100) no resuelve automáticamente los cuellos de botella si la carga de trabajo sigue dominada por microoperaciones no computacionales, conversiones de disposición o flujo de control.
• Jerarquía de Optimización: Los autores concluyen que la activación de TF32 Tensor Core es una optimación base robusta e inmediata que comprime el trabajo aritmético sin degradar la calidad de la síntesis. En contraste, aunque la disposición de canales al final mejora la localidad de la caché, actualmente expone un nuevo régimen de micro-kernels limitados por la memoria que requieren mayor fusión y ajuste arquitectónico para ser efectivos.
• Guía Práctica: Los hallazgos proporcionan una hoja de ruta para optimizar los modelos de difusión 3D generales: alinear los modos de precisión y los formatos de memoria con las capacidades de Tensor Core de la GPU subyacente, y reconocer que reducir el conteo de ciclos no siempre equivale a mejoras proporcionales en la velocidad si la carga de trabajo resultante no logra saturar los motores de cómputo o de memoria.

El artículo afirma modestamente que estos conocimientos son esenciales para hacer que el despliegue práctico de generadores médicos 3D de alta fidelidad sea viable, yendo más allá de las métricas de calidad de imagen para abordar las realidades computacionales del entrenamiento.