Ejecuta matemáticas de alto rendimiento a escala con NVIDIA nvmath-python
- nvmath-python v1.0 proporciona una abstracción al estilo Python sobre las librerías matemáticas CUDA-X y NVPL (incluyendo cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR y cuBLASMp), ofreciendo operaciones matemáticas avanzadas aceleradas en CPU, GPU y sistemas distribuidos multi-nodo, con integración fluida para NumPy, CuPy y PyTorch.
- La librería introduce soporte para tensores dispersos universales (UST) mediante un lenguaje de dominio específico (DSL), permitiendo crear formatos dispersos óptimos personalizados. Además, incluye APIs genéricas para máxima compatibilidad y APIs especializadas para casos de uso de alto rendimiento, como operaciones compuestas fusionadas y fusión de kernels JIT con cuBLASLt.
- nvmath-python permite una instalación rápida y una selección flexible de backends, se integra con el módulo de logging de Python para el trazado computacional detallado, permite amortizar los costos de planificación y autotuning mediante APIs con estado, y facilita la fusión de kernels personalizada mediante numba-cuda y callbacks de FFT compilados con JIT para maximizar la productividad en cargas científicas.
El contenido generado por IA puede resumir información de forma incompleta. Verifique la información crítica. Aprenda más
NVIDIA nvmath-python es una librería diseñada para tender un puente entre la comunidad científica de Python y las librerías matemáticas NVIDIA CUDA-X. Permite a los usuarios de Python acceder al rendimiento de CUDA-X en operaciones matemáticas comunes sin interrumpir sus flujos de trabajo existentes. Dependiendo de la API, las operaciones pueden ejecutarse en una CPU, una GPU con soporte CUDA o en sistemas distribuidos multi-GPU y multi-nodo.
Lanzamiento de nvmath-python v1.0
Con la disponibilidad general de nvmath-python v1.0, exploramos el diseño y las capacidades únicas de la librería para acelerar operaciones matemáticas, desde una simple CPU o GPU hasta una escala de múltiples nodos. nvmath-python actúa como una capa de abstracción sobre librerías como cuFFT, cuBLASLt, cuDSS, cuSPARSE, cuTENSOR y cuBLASMp. Un enfoque innovador hacia la dispersión, el tensor disperso universal (UST), permite al usuario crear su propio formato disperso óptimo mediante un lenguaje de dominio específico sin necesidad de programarlo manualmente.
Instalación rápida y flexible
Instalar un paquete de Python con dependencias nativas complejas puede ser una experiencia frustrante. nvmath-python se instala rápidamente y puede personalizarse para diversos entornos.
- Elija un gestor de paquetes, como pip, conda, uv o pixi.
- Existe la opción de instalar todas las dependencias requeridas mediante el sistema de resolución del gestor de paquetes o realizar una instalación mínima, útil en escenarios como CI/CD o entornos que solo utilizan CPU.
- Seleccione el backend de CPU, el soporte de APIs de dispositivo o las APIs distribuidas.
- Elija una librería de arreglos complementaria para trabajar, como NumPy, CuPy o PyTorch (o todas ellas). Consulte la guía de instalación detallada para ver las opciones disponibles.
Un complemento útil para librerías de arreglos existentes
Al igual que otras librerías como NumPy, nvmath-python implementa operaciones numéricas centrales útiles en muchas aplicaciones de ingeniería y computación científica. Sin embargo, no pretende reemplazar a las librerías de arreglos de propósito general ni ofrecer funciones tradicionales como indexación, slicing o reducción.
En su lugar, nvmath-python se enfoca en exponer la funcionalidad completa y la potencia de las librerías matemáticas CUDA-X en Python, facilitando que las librerías y frameworks existentes utilicen rutinas aceleradas por GPU altamente optimizadas sin depender de interfaces C/C++ de bajo nivel.
En el siguiente ejemplo, nvmath-python consume arreglos de NumPy y el resultado es también un arreglo de NumPy.
import numpy as np
import nvmath
m, n, k = 10, 40, 100
a = np.random.randn(m, k) # a es un arreglo de NumPy
b = np.random.randn(k, n) # b es un arreglo de NumPy
c = nvmath.linalg.advanced.matmul(a, b) # c también es un arreglo de NumPyElección de espacios de memoria y ejecución
La flexibilidad de elegir una librería de arreglos se aplica tanto a librerías de GPU, como CuPy, como a librerías de CPU, como NumPy. Esto es posible porque nvmath-python está respaldado por lo siguiente:
- Librerías distribuidas como cuBLASMp, cuSOLVERMp o cuFFTMp.
Este soporte simplifica la migración de código entre CPU y GPU y permite flujos de trabajo híbridos y distribuidos que combinan la ejecución en ambos entornos.
El siguiente código ilustra cómo nvmath-python soporta múltiples espacios de memoria y ejecución:
import cupy as cp
import numpy as np
import nvmath
N = 2048
a_gpu = cp.random.randn(N) + 1j * cp.random.randn(N)
a_cpu = np.random.randn(N) + 1j * np.random.randn(N)
c_gpu = nvmath.fft.fft(a_gpu)
c_cpu = nvmath.fft.fft(a_cpu)El espacio de ejecución de fft para cada llamada se infiere a partir de su tensor de entrada, ya sea a_gpu o a_cpu, aunque se puede especificar un espacio de ejecución diferente. La función de registro de la librería muestra dónde se ejecutó cada operación.
APIs genéricas y especializadas
Las APIs dentro de nvmath-python se dividen en dos clases: APIs genéricas que actúan como herramientas versátiles (amplias pero superficiales), y APIs especializadas diseñadas como instrumentos precisos (estrechas y profundas).
Las APIs genéricas se enfocan en proporcionar una experiencia de usuario uniforme a través de diversos espacios de ejecución y memoria, así como tipos de operandos; sin embargo, limitan la configurabilidad a las características básicas compartidas en su amplio alcance. Mientras tanto, las APIs especializadas ofrecen un conjunto completo de características y configuraciones diseñadas específicamente para un rango operativo estrecho y pueden estar restringidas a hardware particular.
Para ilustrar, la multiplicación de matrices avanzada implementa la operación compuesta \(\scriptstyle \mathbf{D}=f(\mathbf{A}\mathbf{B}+\mathbf{C})\) específicamente para operandos densos en la GPU y proporciona cada configuración necesaria para extraer la mayor eficiencia de hardware posible. Por el contrario, la API de multiplicación de matrices genérica admite operandos densos y estructurados en espacios de ejecución de CPU y GPU, pero ofrece solo el subconjunto común de opciones aplicables a su alcance más amplio.
La elección óptima depende totalmente del caso de uso específico: las APIs especializadas son ideales cuando una operación se convierte en un cuello de botella computacional que exige optimizaciones específicas del hardware o acceso a características distintas. Mientras tanto, las APIs genéricas son más adecuadas para tareas que no son críticas en términos de rendimiento o cuando la personalización especializada es innecesaria. Todas las APIs especializadas residen dentro de los submódulos advanced para mantenerlas distintas de las APIs genéricas.
Registro con nvmath-python
La librería proporciona integración con el logger de la biblioteca estándar de Python desde el módulo logging para capturar detalles computacionales en varios niveles (debug, información, warning y error).
El siguiente ejemplo ilustra el flujo de datos entre los espacios de memoria y ejecución (usando la multiplicación de matrices avanzada).
Vía NVIDIA Developer.
