Implementar un modelo de aprendizaje profundo en producción siempre ha implicado una brecha dolorosa entre el modelo que entrena un investigador y el modelo que realmente se ejecuta de manera eficiente a escala. TensorRT existe, Torch-TensorRT existe, TorchAO existe, pero conectarlos, decidir qué backend usar para cada capa y validar que el modelo ajustado aún produce resultados correctos ha significado históricamente un trabajo de ingeniería personalizado sustancial. El equipo de IA de NVIDIA ahora está abriendo un conjunto de herramientas diseñado para agrupar ese esfuerzo en una única API de Python.
NVIDIA AITune es un conjunto de herramientas de inferencia diseñado para ajustar e implementar modelos de aprendizaje profundo con un enfoque en las GPU NVIDIA. Disponible bajo la licencia Apache 2.0 e instalable a través de PyPI, el proyecto está dirigido a equipos que desean una optimización de inferencia automatizada sin tener que reescribir sus canalizaciones de PyTorch existentes desde cero. Cubre TensorRT, Torch Inductor, TorchAO y más, los compara todos en su modelo y hardware, y elige al ganador, sin adivinanzas ni ajustes manuales.
Lo que realmente hace AITune
En esencia, AITune opera a nivel de nn.Module. Proporciona capacidades de ajuste de modelos a través de rutas de compilación y conversión que pueden mejorar significativamente la velocidad de inferencia y la eficiencia en diversas cargas de trabajo de IA, incluidas la visión por computadora, el procesamiento del lenguaje natural, el reconocimiento de voz y la IA generativa.
En lugar de obligar a los desarrolladores a configurar manualmente cada backend, el kit de herramientas permite un ajuste perfecto de los modelos y canalizaciones de PyTorch utilizando varios backends como TensorRT, Torch-TensorRT, TorchAO y Torch Inductor a través de una única API de Python, con los modelos optimizados resultantes listos para su implementación en entornos de producción.
También ayuda comprender qué son realmente estos backends. TensorRT es el motor de optimización de inferencia de NVIDIA que compila capas de redes neuronales en núcleos de GPU altamente eficientes. Torch-TensorRT integra TensorRT directamente en el sistema de compilación de PyTorch. TorchAO es el marco de optimización acelerada de PyTorch y Torch Inductor es el backend del compilador propio de PyTorch. Cada uno tiene diferentes fortalezas y limitaciones e históricamente, elegir entre ellos requería compararlos de forma independiente. AITune está diseñado para automatizar esa decisión por completo.
Dos modos de sintonización: anticipado y justo a tiempo
AITune admite dos modos: ajuste anticipado (AOT), en el que proporciona un modelo o una canalización y un conjunto de datos o cargador de datos, y confía en inspeccionar para detectar módulos prometedores para ajustarlos o seleccionarlos manualmente, y ajuste justo a tiempo (JIT), en el que establece una variable de entorno especial, ejecuta su script sin cambios y AITune, sobre la marcha, detectará los módulos y los ajustará uno por uno.
La ruta AOT es la ruta de producción y la más poderosa de las dos. AITune perfila todos los backends, valida la corrección automáticamente y serializa el mejor como un artefacto .ait: compila una vez, sin preparación en cada redistribución. Esto es algo que torch.compile por sí solo no te ofrece. Los pipelines también son totalmente compatibles: cada submódulo se ajusta de forma independiente, lo que significa que diferentes componentes de un único pipeline pueden terminar en diferentes backends dependiendo de qué pruebas comparativas son más rápidas para cada uno. El ajuste de AOT detecta el eje por lotes y los ejes dinámicos (ejes que cambian de forma independientemente del tamaño del lote, como la longitud de la secuencia en LLM), permite seleccionar módulos para ajustarlos, admite la mezcla de diferentes backends en el mismo modelo o canalización y le permite elegir una estrategia de ajuste, como el mejor rendimiento para todo el proceso o por módulo. AOT también admite el almacenamiento en caché, lo que significa que no es necesario reconstruir un artefacto previamente ajustado en ejecuciones posteriores, solo cargarlo desde el disco.
La ruta JIT es la ruta rápida y es más adecuada para una exploración rápida antes de comprometerse con AOT. Establezca una variable de entorno, ejecute su script sin cambios y AITune descubre automáticamente los módulos y los optimiza sobre la marcha. Sin cambios de código, sin configuración. Una restricción práctica importante: import aitune.torch.jit.enable debe ser la primera importación en su script al habilitar JIT mediante código, en lugar de mediante la variable de entorno. A partir de la versión 0.3.0, el ajuste JIT requiere solo una muestra y ajustes en la primera llamada del modelo, una mejora con respecto a las versiones anteriores que requerían múltiples pases de inferencia para establecer la jerarquía del modelo. Cuando un módulo no se puede ajustar, por ejemplo, porque se detecta una ruptura en el gráfico, lo que significa que torch.nn.Module contiene lógica condicional en las entradas, por lo que no hay garantía de un gráfico estático y correcto de los cálculos, AITune deja ese módulo sin cambios e intenta ajustar sus hijos en su lugar. El backend alternativo predeterminado en el modo JIT es Torch Inductor. Las desventajas de JIT en relación con AOT son reales: no puede extrapolar tamaños de lotes, no puede realizar comparaciones entre backends, no admite el almacenamiento de artefactos y no admite el almacenamiento en caché: cada nueva sesión del intérprete de Python se vuelve a sintonizar desde cero.
Tres estrategias para la selección de backend
Una decisión de diseño significativa en AITune es su abstracción estratégica. No todos los backends pueden ajustar todos los modelos; cada uno se basa en una tecnología de compilación diferente con sus propias limitaciones, como la exportación ONNX para TensorRT, rupturas de gráficos en Torch Inductor y capas no compatibles en TorchAO. Las estrategias controlan cómo AITune maneja esto.
Se proporcionan tres estrategias. FirstWinsStrategy prueba los backends en orden de prioridad y devuelve el primero que tiene éxito, lo que resulta útil cuando desea una cadena alternativa sin intervención manual. OneBackendStrategy utiliza exactamente un backend especificado y muestra la excepción original inmediatamente si falla, lo cual es apropiado cuando ya ha validado que un backend funciona y desea un comportamiento determinista. HighestThroughputStrategy perfila todos los backends compatibles, incluido TorchEagerBackend como base junto con TensorRT y Torch Inductor, y selecciona el más rápido, a costa de un tiempo de ajuste inicial más largo.
Inspeccionar, ajustar, guardar, cargar
La superficie API es deliberadamente mínima. ait.inspect() analiza la estructura de un modelo o canalización e identifica qué subcomponentes nn.Module son buenos candidatos para el ajuste. ait.wrap() anota los módulos seleccionados para su ajuste. ait.tune() ejecuta la optimización real. ait.save() conserva el resultado en un archivo de punto de control .ait, que agrupa los pesos de los módulos originales y sintonizados junto con un archivo hash SHA-256 para verificar la integridad. ait.load() lo vuelve a leer. En la primera carga, se descomprime el punto de control y se cargan las pesas; las cargas posteriores utilizan los pesos ya descomprimidos de la misma carpeta, lo que agiliza la redistribución.
El backend de TensorRT proporciona inferencia altamente optimizada utilizando el motor TensorRT de NVIDIA e integra TensorRT Model Optimizer en un flujo perfecto. También es compatible con ONNX AutoCast para inferencia de precisión mixta a través de TensorRT ModelOpt y CUDA Graphs para reducir la sobrecarga de la CPU y mejorar el rendimiento de la inferencia: CUDA Graphs captura y reproduce automáticamente las operaciones de la GPU, lo que elimina la sobrecarga del lanzamiento del kernel para llamadas de inferencia repetidas. Esta característica está deshabilitada de forma predeterminada. Para los desarrolladores que trabajan con modelos instrumentados, AITune también admite ganchos hacia adelante en los modos de ajuste AOT y JIT. Además, la versión 0.2.0 introdujo soporte para caché KV para LLM, ampliando el alcance de AITune a canales de modelos de lenguaje basados en transformadores que aún no cuentan con un marco de servicio dedicado.
Conclusiones clave
NVIDIA AITune es un conjunto de herramientas de Python de código abierto que compara automáticamente múltiples backends de inferencia (TensorRT, Torch-TensorRT, TorchAO y Torch Inductor) en su modelo y hardware específicos, y selecciona el de mejor rendimiento, eliminando la necesidad de una evaluación manual del backend. AITune ofrece dos modos de ajuste: anticipado (AOT), la ruta de producción que perfila todos los backends, valida la corrección y guarda el resultado como un artefacto .ait reutilizable para una redistribución sin preparación; y justo a tiempo (JIT), una ruta de exploración sin código que sintoniza la primera llamada al modelo simplemente configurando una variable de entorno. Tres estrategias de ajuste (FirstWinsStrategy, OneBackendStrategy y HighestThroughputStrategy) brindan a los desarrolladores de IA un control preciso sobre cómo AITune selecciona un backend, desde cadenas de respaldo rápidas hasta perfiles de rendimiento exhaustivos en todos los backends compatibles. AITune no reemplaza a vLLM, TensorRT-LLM o SGLang, que están diseñados específicamente para servir modelos de lenguaje grandes con características como procesamiento por lotes continuo y decodificación especulativa. En cambio, se dirige al panorama más amplio de los modelos y canalizaciones de PyTorch (visión por computadora, difusión, habla e incrustaciones) donde tales marcos especializados no existen.
Consulte el repositorio. Además, no dude en seguirnos en Twitter y no olvide unirse a nuestro SubReddit de más de 120.000 ML y suscribirse a nuestro boletín. ¡Esperar! estas en telegrama? Ahora también puedes unirte a nosotros en Telegram.
¿Necesita asociarse con nosotros para promocionar su repositorio de GitHub O su página principal de Hugging O su lanzamiento de producto O seminario web, etc.? Conéctate con nosotros