En un artículo anterior, se presentó el Tensor Escaso Universal (UST), que permite a los desarrolladores desacoplar la escasez de un tensor de su disposición en memoria, mejorando la flexibilidad y el rendimiento. Ahora, se anuncia la integración del UST en nvmath-python v0.9.0, lo que acelerará aplicaciones científicas y de aprendizaje profundo con matrices dispersas.
Este artículo ofrece un recorrido por las características clave del UST, detalles de implementación y un resumen del rendimiento, incluyendo:
- Interoperabilidad sin costo: Conversión sin movimiento de datos con PyTorch, SciPy y CuPy.
- Formatos personalizados: Definición de nuevos esquemas de escasez.
- Operaciones polimórficas: Funciones independientes de la escasez que utilizan automáticamente núcleos optimizados o generan código disperso personalizado.
- Inyección en PyTorch: Integración fácil de los beneficios del UST en modelos existentes de PyTorch.
- Cacheo transparente: Evita recompilaciones y replanteamientos, amortizando la sobrecarga en ejecuciones repetidas de la misma operación.
Formato de tensor DSL
El UST describe formatos de almacenamiento de tensor escaso comunes y menos comunes (p. ej., COO, CSR, CSC y BSR) mediante un lenguaje específico de dominio (DSL) simple. Por ejemplo, un formato CSC se describe de la siguiente manera.
(i, j) -> (j : denso, i : comprimido) # CSC
Integrar este DSL en un lenguaje de programación implica decisiones de diseño con diversos compromisos.
En la biblioteca C++ MatX, por ejemplo, el DSL se implementa puramente a través de tipos y plantillas, lo que permite optimizaciones interesantes en tiempo de compilación para formatos de almacenamiento específicos. En el siguiente método de despacho foo(), se pueden evaluar pruebas específicas de formato en tiempo de compilación.
template foo(TensorType a) { if constexpr (TensorType::Format::isCOO()) { ... despacho al código COO ... } else if constexpr (TensorType::Format::isCSR()) { ... despacho al código CSR ... } else ... }
En nvmath-python, adoptamos una forma más sencilla de presentar el DSL, sacrificando algo de rendimiento en la inspección de objetos de formato. El formato CSC mencionado anteriormente se expresa así.
i, j = ( Dimension(dimensi), Dimension(dimensi) ) CSC = TensorFormat( [i, j], {j: LevelFormat.DENSE, i: LevelFormat.COMPRESSED} )
Una gran ventaja de estos objetos es que todo se puede construir dinámicamente en tiempo de ejecución. Sin embargo, realizar tareas específicas de formato requiere inspeccionar los contenidos reales en ese momento. Dado que tales decisiones generalmente ocurren fuera de rutas críticas de rendimiento, este compromiso parece una opción aceptable.
Interoperabilidad con PyTorch, SciPy, CuPy
La implementación del UST en nvmath-python proporciona interoperabilidad con tensores de PyTorch, SciPy, CuPy y NumPy. La conversión es mayormente sin costo, lo que significa que convertir formatos densos o dispersos como COO, CSR, CSC, BSR, BSC y DIA a un objeto UST o de vuelta se realiza sin movimiento o copia de datos. En su lugar, el objeto UST hace referencia a los búferes de almacenamiento de la estructura de datos original.
Por ejemplo, la siguiente conversión de una matriz dispersa de SciPy en formato COO a un objeto UST de nvmath-python utiliza los arreglos de fila, columna y valores para formar las posiciones, coordenadas y arreglo de valores del UST.
row = np.array([0, 0, 1, 1, 2, 3], dtype=np.int32) col = np.array([0, 1, 1, 3, 2, 3], dtype=np.int32) val = np.array([1, 2, 3, 4, 5, 6], dtype=np.float32) coo = sps.coo_array((val, (row, col)), shape=(4, 8)) # SciPy ust = Tensor.from_package(coo)
Esta conversión genera automáticamente el DSL para COO, ilustrado a continuación. La conversión de otros formatos resultará en el DSL correspondiente para esos formatos.
TensorFormat( [i, j] -> {i: (, ), j: } )
El objeto UST se convierte nuevamente a un tensor disperso en el paquete original de la siguiente manera. Esta también es una conversión sin costo al pasar simplemente referencias a los búferes constituyentes.
coo = ust.to_package() # esto genera nuevamente un formato COO de SciPy
Formatos menos comunes
Los desarrolladores también pueden definir su propio formato de escasez novedoso usando el DSL. Por ejemplo, el siguiente código convierte un tensor denso de PyTorch en un formato comprimido de delta de 2 bits (similar a CSR, pero utilizando diferencias en tiempo de ejecución para las coordenadas), que no está soportado nativamente en PyTorch.
A = torch.tensor([[1, 0, 0, 0, 0, 0, 0, 2], [0, 0, 3, 4, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 5]], dtype=torch.float64) U = Tensor.from_package(A) # denso UST delta_format = TensorFormat([i, j], {i: LevelFormat.DENSE, j: (LevelFormat.DELTA, 2)}) # usa 2 bits S = U.convert(tensor_format=delta_format) # UST disperso
Esto produce el almacenamiento UST disperso mostrado en la Figura 1, donde cada coordenada utiliza una distancia al siguiente elemento almacenado en lugar de un índice de columna directo. Es importante mencionar que fue necesario un relleno (los ceros mostrados en rojo) para los casos donde esa distancia superaría 3 (máximo utilizando 2 bits).
Utilidades de impresión y dibujo
Se proporcionan diversas utilidades para ayudar a los usuarios a depurar, probar y familiarizarse con la implementación del UST en nvmath-python. Primero, el método dunder __repr__ proporciona una representación de cadena inequívoca que puede usarse para imprimir los contenidos de almacenamiento del tensor.
A = torch.arange(4 * 5).reshape(4, 5).cuda().to_sparse_csr() U = Tensor.from_package(A) print(U)
Este fragmento produce la siguiente salida, que muestra los tipos utilizados para los valores, posiciones y arreglos de coordenadas en el almacenamiento UST, así como el dispositivo, las dimensiones y extents de nivel, los datos almacenados y un resumen del número de bytes utilizados para el almacenamiento.
---- Tensor disperso formato : [i, j] -> (i: , j: ) dispositivo : cuda dim : [4, 5] lvl : [4, 5] nse : 19 pos[1] : [0, 4, 9, 14, 19] #5 crd[1] : [1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4] #19 valores : [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19] #19 datos : 344 bytes sparsity : 5.00% ----
Los métodos draw_storage() y draw() generan una imagen del almacenamiento del tensor o del contenido del tensor, respectivamente. El primero se utilizó para generar la Figura 1. Llamar al segundo método en el ejemplo anterior de UST genera la imagen mostrada en la Figura 2, donde el color gris denota elementos no cero almacenados.
Ambos métodos muestran el UST para ejemplos relativamente pequeños, ya que no escalan a tamaños muy grandes. Para tensores más grandes, se puede utilizar el método draw_raw() para obtener una visualización de solo la estructura no cero del tensor, como se ilustra en la Figura 3 para una matriz tomada de la biblioteca SuiteSparse.
Operaciones polimórficas
El DSL del formato de tensor del UST permite a los desarrolladores enfocarse en la escasez del tensor. Sin embargo, definir el almacenamiento de los tensores es solo parte del requisito. Los desarrolladores también pueden especificar las operaciones a realizar.
La implementación del UST en nvmath-python permite a los desarrolladores definir cálculos utilizando operaciones independientes de la escasez como matmul() y solve(). Este enfoque es poderoso porque el sistema inspecciona los formatos de los operandos para determinar la ruta de ejecución, despachando a una biblioteca optimizada o núcleo para formatos comunes o confiando en la generación automática de código disperso cuando no existe una solución optimizada.
Este diseño permite el uso de formatos dispersos novedosos sin codificación manual explícita, mientras se logra un alto rendimiento utilizando soluciones de bibliotecas existentes para formatos estándar. La sintaxis para planificar y realizar una operación matmul() se muestra a continuación.
# c := a @ b + c con Matmul(a, b, c, beta=1.0) como mm: mm.plan() mm.execute()
Este enfoque amortiza el costo inicial de planificación en muchas ejecuciones de la misma invocación de matmul(), lo que realmente resulta beneficioso en, por ejemplo, solvers dispersos iterativos o capas lineales (recortadas) en aprendizaje profundo.
Inyección de UST en PyTorch
La interoperabilidad entre tensores dispersos de PyTorch y nvmath-python permite a los desarrolladores experimentar con esquemas de almacenamiento disperso novedosos dentro de modelos de PyTorch. Sin embargo, los investigadores en IA son reacios a reescribir su código de modelo para reemplazar las operaciones GEMM realizadas dentro de capas lineales por las operaciones polimórficas del UST, a pesar de que esto ofrece un rendimiento mucho mejor.
La implementación del UST en nvmath-python proporciona una forma de “inyectar” el UST en modelos existentes sin reescribir el código original del modelo. Esto se logra envolviendo el UST en una subclase de torch.Tensor y proporcionando envolturas a las operaciones polimórficas del UST. Al ofrecer envolturas para operaciones comunes de tensores de PyTorch, los investigadores pueden usar el UST sin tener que aprender mucho sobre el funcionamiento interno.
Estas envolturas también añaden otra capa de cacheo para reutilizar instancias de MatMul planificadas para operaciones similares, reduciendo el tiempo de ejecución posterior a solo la fase de ejecución.
También está disponible un método reformat_model() que itera a través de los pesos de todas las capas lineales, permitiendo que una función definida por el usuario inspeccione, potencialmente recorte y haga escaso (o salte) cada matriz de pesos.
weights = torchvision.models.get_model_weights(model_name).DEFAULT model = torchvision.models.get_model(model_name, weights=weights) model.to(device) model.eval() ... reformat_model(model, func=reformat) ... con torch.inference_mode(): prediction = model(batch)
La forma general del método de reformateo proporcionado por el usuario se muestra a continuación. Retornar None indica que se deja la matriz de pesos sin cambios. Después de esto, la inferencia utilizará el UST para todas las aplicaciones de capas lineales que han sido modificadas.
def reformat(weight): ... inspeccionar (posiblemente incluso recortar peso), luego decidir sobre el almacenamiento ... if (...) return TorchUST.from_torch(weight.to_sparse_csr()) return None }
Ejemplos de rendimiento
Primero, comparamos el rendimiento de SpMV utilizando la operación @ de CuPy entre formatos soportados nativamente (COO, CSR y DIA), junto con mv() de PyTorch para COO y CSR frente al matmul() del UST en la matriz 1,489,752 x 1,48,9752 atmosmodl de Matrix Market. Para el último, medimos solo el tiempo de ejecución de execute(), lo que es una comparación justa para aplicaciones con multiplicación repetida. Los tiempos resultantes se muestran en la Figura 4.
El UST logró aceleraciones que van desde 1.1 hasta 444 en este escenario. Se observa un bajo rendimiento de la implementación COO de CuPy y PyTorch. Para el formato CSR, todas las versiones utilizan la implementación cuSPARSE, beneficiándose el UST de la reutilización de la fase de planificación. La aceleración observada con el formato DIA se debe al uso de un núcleo especializado por parte del UST, en contraste con el método de CuPy que primero convierte a CSR.
El segundo experimento se basa en la publicación MACKO: Sparse Matrix-Vector Multiplication for Low Sparsity, que proporciona una implementación eficiente de la operación SpMV, un paso importante para la inferencia de un solo token en aprendizaje profundo. Los autores introducen el formato MACKO, que es esencialmente el formato comprimido de delta de la Figura 1, y ofrecen una impresionante implementación CUDA que ya entrega aceleraciones sobre implementaciones densas para un 50% de esparcidad no estructurada y más.
Como el UST de nvmath-python proporciona una forma fácil de incorporar nuevos núcleos escritos a mano en el mecanismo de búsqueda, experimentamos con usar su implementación como núcleo de backend para el formato comprimido delta. La Figura 5 compara el tiempo de ejecución de una implementación densa (GEMV), la implementación cuSPARSE de SPMV, el MACKO UST y la implementación original de MACKO para matrices dispersas uniformemente aleatorias de 8192×8192, variando desde 0% de escasez (denso) hasta 100% de escasez (cero).
Aprende más
Para una comprensión más profunda de la implementación del UST en nvmath-python, consulte la documentación en línea de nvmath-python.
