Matrices dispersas por bloques para modelos de lenguaje más pequeños y rápidos

'Matrices dispersas por bloques para modelos de lenguaje más pequeños y rápidos' can be condensed to 'Matrices dispersas por bloques para modelos de lenguaje más eficientes'.

Ahorrando espacio y tiempo, un cero a la vez

En publicaciones anteriores, presentamos las matrices dispersas y lo que pueden hacer para mejorar las redes neuronales.

La suposición básica es que las capas densas completas a menudo son excesivas y se pueden podar sin una pérdida significativa en la precisión. En algunos casos, las capas lineales dispersas incluso pueden mejorar la precisión y/o generalización.

El problema principal es que el código actualmente disponible que admite el cálculo algebraico disperso carece gravemente de eficiencia. También estamos esperando el soporte oficial de PyTorch.

Por eso nos quedamos sin paciencia y nos tomamos un tiempo este verano para abordar esta “laguna”. Hoy, nos emociona lanzar la extensión pytorch_block_sparse.

Por sí mismo, o incluso mejor combinado con otros métodos como la destilación y la cuantización, esta biblioteca permite redes que son tanto más pequeñas como más rápidas, algo que Hugging Face considera crucial para permitir que cualquier persona use redes neuronales en producción a bajo costo, y para mejorar la experiencia para el usuario final.

Uso

El módulo proporcionado BlockSparseLinear es un reemplazo directo de torch.nn.Linear, y es trivial usarlo en tus modelos:

# from torch.nn import Linear
from pytorch_block_sparse import BlockSparseLinear

...

# self.fc = nn.Linear(1024, 256)
self.fc = BlockSparseLinear(1024, 256, densidad=0.1)

La extensión también proporciona un BlockSparseModelPatcher que permite modificar un modelo existente “sobre la marcha”, como se muestra en este cuaderno de ejemplo. Luego, este modelo se puede entrenar como de costumbre, sin cambios en el código fuente del modelo.

NVIDIA CUTLASS

Esta extensión se basa en el concepto de prueba de azulejos dispersos por Yulhwa Kim.

Está utilizando plantillas C++ CUDA para la multiplicación de matrices dispersas por bloques basada en CUTLASS.

CUTLASS es una colección de plantillas de CUDA C++ para implementar núcleos CUDA de alto rendimiento. Con CUTLASS, es posible acercarse al rendimiento de cuBLAS en núcleos personalizados sin recurrir a código de lenguaje ensamblador.

Las últimas versiones incluyen todas las primitivas de Ampere Tensor Core, lo que proporciona aceleraciones de x10 o más con una pérdida limitada de precisión. Las próximas versiones de pytorch_block_sparse aprovecharán estas primitivas, ya que la dispersión por bloques es 100% compatible con los requisitos de Tensor Cores.

Rendimiento

En la etapa actual de la biblioteca, el rendimiento de las matrices dispersas es aproximadamente dos veces más lento que su contraparte densa optimizada por cuBLAS, y confiamos en que podemos mejorar esto en el futuro.

Esto es una gran mejora en las matrices dispersas de PyTorch: su implementación actual es una orden de magnitud más lenta que la densa.

Pero el punto más importante es que la ganancia de rendimiento al usar matrices dispersas aumenta con la dispersión, por lo que una matriz 75% dispersa es aproximadamente 2 veces más rápida que la equivalente densa.

Los ahorros de memoria son aún más significativos: para una dispersión del 75%, el consumo de memoria se reduce en un 4 veces, como cabría esperar.

Trabajo futuro

Poder entrenar de manera eficiente capas lineales dispersas por bloques fue solo el primer paso. El patrón de dispersión actualmente está fijo en la inicialización, y, por supuesto, optimizarlo durante el aprendizaje proporcionará grandes mejoras.

Por lo tanto, en futuras versiones, se pueden esperar herramientas para medir la “utilidad” de los parámetros para poder optimizar el patrón de dispersión. El patrón disperso del 50% de NVIDIA Ampere dentro de los bloques probablemente proporcionará otra mejora significativa de rendimiento, al igual que la actualización a versiones más recientes de CUTLASS.

¡Así que mantente atento a más beneficios de dispersión en un futuro cercano!