- TensorRT 11.0 introduce inferencia multi-device nativa, apoyada en NVIDIA NCCL para colectivas distribuidas de alto throughput, escalando pipelines generativos entre múltiples GPUs, incluyendo escenarios de edge.
- El paralelismo de contexto, soportado vía las primitivas
IDistCollectiveLayeren TensorRT 11.0, particiona las secuencias de entrada entre GPUs y se optimiza con estrategias como AllGather KV, Ring Attention y DeepSpeed Ulysses, cada una con su propio balance de cómputo, memoria y overhead de comunicación.
- Los benchmarks sobre NVIDIA Cosmos 3 y FLUX.1 indican que DeepSpeed Ulysses entrega la menor latencia para generación de medios basada en diffusion con contextos extremadamente largos, mientras que Ring Attention también escala bien hasta 4 GPUs.
Las cargas de IA generativa están superando rápidamente el presupuesto de memoria y cómputo de una sola GPU. Para desarrolladores de inferencia que construyen pipelines de generación de medios, el desafío es escalar entre varios dispositivos sin sacrificar las optimizaciones críticas (fusiones de kernel, planificación de memoria, cuantización) que NVIDIA TensorRT entrega para despliegues de producción.
El soporte multi-device, una función nueva en TensorRT 11.0, lleva inferencia multi-GPU nativa de alto rendimiento al runtime de TensorRT, habilitando despliegues de producción multi-device dirigidos también a dispositivos edge.
Combinado con Torch-TensorRT, los desarrolladores pueden convertir y desplegar modelos masivos de PyTorch fuera del framework, rompiendo los límites de memoria y cómputo de un solo dispositivo.
Se puede descargar TensorRT 11.0 con soporte multi-device desde el NVIDIA Developer Portal.
NCCL como capa de transporte distribuida
La NVIDIA Collective Communications Library (NCCL) entrega operaciones colectivas multi-GPU y multi-nodo de alto rendimiento que potencian el entrenamiento de modelos a gran escala sobre miles de GPUs. NCCL selecciona automáticamente el transporte óptimo según la topología, abstrayendo NVIDIA NVLink, NVIDIA NVSwitch, PCIe e InfiniBand detrás de una interfaz uniforme. Al integrarse directamente con NCCL, TensorRT hereda esta optimización de transporte para cargas de inferencia. Más información en https://developer.nvidia.com/nccl.
La nueva función multi-device cubre el set completo de colectivas distribuidas de NCCL: AllReduce, Broadcast, Reduce, AllGather, ReduceScatter, AlltoAll, Gather y Scatter.
¿Qué estrategias de paralelismo soporta TensorRT 11.0?
La inferencia distribuida puede expresarse con varias estrategias de paralelismo, cada una con distintos trade-offs entre ahorro de memoria, escalamiento de cómputo y overhead de comunicación. Las dos más comunes son paralelismo de tensor y paralelismo de contexto.
Paralelismo de tensor
En el paralelismo de tensor, los pesos de una sola capa se particionan entre GPUs. Cada GPU computa un shard de la multiplicación matricial de la capa y luego combina los resultados parciales mediante una colectiva para producir la salida completa. Esto reduce la memoria por dispositivo destinada a pesos, lo que lo vuelve la elección natural (y muchas veces la única) cuando los pesos de una capa exceden la memoria de una GPU, independiente del largo de secuencia o del batch.
En un bloque transformer, las proyecciones column-parallel (por ejemplo, QKV y la up-projection del MLP) se aparean con proyecciones row-parallel (la salida de atención y la down-projection del MLP), de modo que cada bloque requiere solo un AllReduce y el overhead de comunicación queda acotado.
Paralelismo de contexto
En el paralelismo de contexto, la secuencia de entrada se particiona entre GPUs por la dimensión de secuencia. Cada GPU procesa solo un slice de la secuencia, mientras que las operaciones colectivas hacen disponible la secuencia global donde se necesita, como durante la atención. Es especialmente efectivo para cargas de secuencia larga, donde el escalamiento cuadrático de la atención con el largo de secuencia la convierte en el consumidor dominante de cómputo y memoria.
También es un ajuste especialmente natural para modelos de diffusion y DiT, cuya atención bidireccional evita los problemas de balance de carga que surgen con causal masks.
Pueden leer el paper Context Parallelism for Scalable Million-Token Inference para detalles adicionales sobre paralelismo de contexto.
NVIDIA TensorRT 11.0 introduce soporte para las primitivas IDistCollectiveLayer requeridas por las distintas estrategias de paralelización. El resto del artículo se enfoca en paralelismo de contexto, que aborda directamente el costo dominante en pipelines generativos modernos de medios: atención sobre secuencias largas.
Paralelismo de contexto para generación de medios
Los pipelines de diffusion para imagen y video gastan una fracción grande de su presupuesto de cómputo y memoria dentro de bloques de atención operando sobre secuencias largas de tokens. Un latente de imagen de alta resolución o un clip de video multi-frame puede producir secuencias de decenas de miles de tokens por bloque, y la atención escala cuadráticamente con el largo de secuencia.
AllGather KV
El paralelismo de contexto particiona la secuencia entre GPUs. Cada rank procesa un slice de los queries (Q) correspondiente a su partición de secuencia. Una forma directa de implementarlo es el enfoque AllGather KV, donde los ranks intercambian sus shards de key (K) y value (V) mediante una colectiva AllGather antes de computar la atención local, habilitando a cada rank a atender sobre la secuencia completa. El resultado es una salida de atención por-rank cubriendo la secuencia completa al costo de una colectiva adicional por bloque de atención, mientras que la multiplicación local Q × Kᵀ se reduce proporcionalmente al número de ranks.
Para diffusion de video e imagen de alta resolución, este trade-off compone favorablemente entre pasos de denoising. El overhead de comunicación por paso permanece acotado por el AllGather en la dimensión de secuencia, mientras que los ahorros de cómputo y memoria aplican a cada capa de atención en cada paso.
Ring Attention
El paralelismo de contexto puede implementarse de varias maneras, cada una con trade-offs distintos. Una mejora potencial sobre AllGather KV es Ring Attention, donde comunicación y cómputo se solapan. Esto habilita a cada GPU a procesar su Q local mientras K y V circulan continuamente en una topología de anillo. Ring Attention también reduce la huella de memoria: usando un softmax online, los tensores K y V de tamaño completo no necesitan materializarse en ninguna GPU. Pueden leer el paper Ring Attention with Blockwise Transformers for Near-Infinite Context para más detalle.
DeepSpeed Ulysses
Para contexto largo (decenas de miles de tokens), una alternativa de implementación de paralelismo de contexto es DeepSpeed Ulysses. Inicialmente particiona las muestras individuales por la dimensión de secuencia entre las GPUs participantes. Antes del cómputo de atención, emplea una colectiva all-to-all sobre los Q, K y V particionados.
Esto asegura que cada GPU reciba el largo de secuencia completo, pero solo para un subset no solapado de las attention heads, habilitándolas a computar atención en paralelo. Finalmente, una segunda all-to-all reúne los resultados a través de las attention heads mientras los reparticiona por la dimensión de secuencia. Más sobre paralelismo de contexto para secuencias largas en el paper DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models.
¿Qué muestran los benchmarks en C++?
Los siguientes benchmarks evalúan inferencia multi-device con TensorRT para cargas de generación de medios pensadas para deployment de producción en C++. Se usan dos pipelines representativos: uno de generación de video basado en NVIDIA Cosmos 3 y otro de generación de imagen basado en FLUX.1.
Estos pipelines se autoraron primero en PyTorch y luego se convirtieron fuera del framework usando Torch-TensorRT para producir engines de TensorRT adecuados para aplicaciones de inferencia en C++. Este flujo permite a los desarrolladores mantener PyTorch como entorno de desarrollo de modelos mientras despliegan engines TensorRT optimizados en sistemas de producción.
Los benchmarks comparan latencia end-to-end entre distintas estrategias de paralelismo de contexto: AllGather KV, Ring Attention y Ulysses. Todos los resultados se obtuvieron sobre un único nodo con 8 GPUs.
Generación de video con NVIDIA Cosmos 3
La plataforma de modelos NVIDIA Cosmos es una plataforma de world foundation models, y el modelo Cosmos3-Nano puede generar imágenes, video, audio y otros formatos a partir de entradas multimodales: texto, imágenes y video. Para los benchmarks se usó el archivo de prompt de ejemplo. De acuerdo con esos resultados, Ulysses es el ganador claro cuando un modelo de diffusion trabaja con contextos excesivamente largos (del orden de decenas de miles de tokens).
Generación de imagen con FLUX.1
El detalle completo de los benchmarks de generación de imagen con FLUX.1, incluyendo comparativas adicionales de latencia y speedup, está disponible en el blog técnico de NVIDIA Developer.
