XNNPack es el motor de inferencia de CPU de TensorFlow Lite predeterminado para todos los modelos. Ofrece aceleraciones innovadoras en plataformas de dispositivos móviles, computadoras de escritorio y la Web. Una de las optimizaciones empleadas en XNNPack es el reempaquetamiento de los pesos estáticos de los operadores de convolución, convolución en profundidad, convolución transpuesta y completamente conectados en un diseño interno adaptado a cálculos de inferencia. Durante la inferencia, se accede a los pesos reempaquetados en un patrón secuencial que es compatible con las canalizaciones de los procesadores.
La reducción de la latencia de inferencia tiene un costo: en esencia, el reempaquetamiento crea una copia adicional de los pesos dentro de XNNPack. Ya se han dedicado esfuerzos a reducir ese costo agregando un almacenamiento en la memoria caché a XNNPack. Esta caché permite compartir los pesos empaquetados entre intérpretes de TFLite independientes que ejecutarían el mismo modelo de forma independiente.
Se mejoró la implementación del delegado de TFLite XNNPack para abordar algunas de las deficiencias de la caché existente.
1. El almacenamiento en caché se realiza en una memoria anónima, de modo que es necesario almacenar en disco en caso de que no haya suficiente memoria, lo que genera un bajo rendimiento.
2. Requiere volver a empaquetar los pesos iniciales cada vez que se inicia un proceso.
3. Debido a que, durante el reempaquetamiento, se leen los pesos de TFLite originales y se escribe en un nuevo búfer, se genera un alto uso máximo de memoria durante el empaquetado.
4. Es necesario seguir pasos tediosos y realizar una administración cuidadosa del ciclo de vida para habilitar adecuadamente el almacenamiento en caché a través del delegado de XNNPack.
5. No es posible compartir los pesos entre los procesos.
XNNPack se actualizó y proporciona una interfaz que te permite implementar un proveedor de almacenamiento en caché de pesos. Un proveedor de almacenamiento en caché de pesos se comporta como un diccionario que XNNPack llena y consulta para acceder a los búferes empaquetados. Estas son sus principales funciones.
look_up busca una clave de búfer empaquetada y muestra un identificador único (o un identificador especial reservado para NotFound) que se puede usar más tarde para recuperar la dirección del búfer.reserve_space reserva un búfer que se puede usar para almacenar información de un tamaño determinado. Luego, ese búfer se debe confirmar usando look_up_or_insert.look_up_or_insert comprueba si existe un búfer que coincida con la clave proporcionada en el proveedor de almacenamiento en caché. Si no es así, los datos proporcionados se envían al proveedor de almacenamiento en caché. Esta función también muestra el identificador que se puede usar para recuperar la dirección del búfer.offset_to_addr muestra la dirección del búfer del identificador que muestran look_up y look_up_or_insert.Las interacciones entre XNNPack y el proveedor de almacenamiento de pesos en caché se ilustran en el siguiente diagrama.
El delegado de TFLite ahora utiliza esta nueva interfaz y tiene su propio proveedor de almacenamiento en caché de pesos. Este proveedor es capaz de guardar y cargar los pesos empaquetados directamente en el disco y desde este. TFLite ha aprovechado la asignación de memoria flatbuffer y respaldada por archivos durante mucho tiempo. Estamos llenando el vacío utilizando la misma técnica, ya que se obtienen las siguientes ventajas.
Los pesos empaquetados persistentes en el disco evitan el costoso proceso de reempaquetamiento cada vez que se carga un modelo. Esto se traduce en una reducción significativa tanto en la latencia de inicio como en el uso máximo de memoria. Incluso para la compilación inicial, se ofrece deduplicación de datos empaquetados y se mejora aún más el rendimiento de empaquetamiento, ya que se evita reempaquetar los mismos datos.
mmap aprovecha la administración de la memoria virtual del sistema operativo, lo que le permite optimizar el uso y el rendimiento general de la memoria del sistema. En nuestro caso, esta característica es especialmente ventajosa para el acceso aleatorio a archivos voluminosos de solo lectura, como los pesos constantes de operación de una red neuronal.
Con los datos empaquetados almacenados en el disco, la memoria caché de XNNPack ya no depende de la memoria anónima, que puede tener propensión a sufrir problemas de rendimiento cuando aumenta su uso. En su lugar, aprovecha la administración de la memoria virtual del sistema operativo para brindar un funcionamiento más fluido.
Al eliminar la necesidad de copiar datos entre el sistema de archivos y la memoria, mmap reduce significativamente la sobrecarga y acelera los tiempos de acceso.
Puedes encontrar más información sobre las asignaciones de archivos y el uso de memoria directamente en la página del manual de mmap y otra información interesante.
La carga de archivos basada en mmap permite el intercambio de pesos sin problemas entre múltiples procesos, ya que el espacio de direcciones virtuales de cada proceso se asigna a las mismas páginas de memoria física. Esto no solo reduce la huella de memoria general, ya que varios procesos comparten la misma memoria, sino que también acelera la carga de modelos en todos los ámbitos.
En lugar de requerir que el usuario configure y administre el objeto de caché durante toda la vida útil de la app, simplemente puede proporcionar una ruta al archivo de caché.
std::unique_ptr<tflite::Interpreter> interpreter;
// Setup the options for the XNNPack delegate.
TfLiteXNNPackDelegateOptions xnnpack_options = TfLiteXNNPackDelegateOptionsDefault();
xnnpack_options.weight_cache_file_path = "/tmp/cache_file.xnn_cache";
// Create and apply the XNNPack delegate to a TFLite interpreter.
// Static weights will be packed and written into weights_cache on the first run.
// They will be automatically loaded for all other runs.
TfLiteDelegate* delegate = TfLiteXNNPackDelegateCreate(&xnnpack_options);
interpreter->ModifyGraphWithDelegate(delegate);
Para garantizar una inferencia precisa y eficiente, es fundamental invalidar la memoria caché de XNNPack en condiciones específicas:
Evolución de modelos: si cambian los pesos o la estructura de tu modelo, los datos almacenados en caché se vuelven obsoletos y se deben invalidar. Es decir, se debe eliminar el archivo en la ruta de caché proporcionada.
Actualizaciones de XNNPack: las actualizaciones del algoritmo de empaquetamiento interno de XNNPack pueden hacer que los pesos en caché sean incompatibles, lo que requiere que se vuelva a calcular la memoria caché. Afortunadamente, XNNPack es capaz de detectar esto y reemplazar la caché existente automáticamente.
En esencia, cualquier modificación que pueda afectar la forma en que XNNPack empaqueta o utiliza los pesos debería desencadenar una invalidación de la memoria caché.
La inicialización de la sesión está dominada por el empaquetamiento de pesos. Para los LLM, varios subgrafos reutilizan los mismos pesos. La creación de la caché es más rápida porque la funcionalidad de deduplicación evita empaquetar esos mismos pesos varias veces. Para modelos más estándar, como la difusión estable, no hay deduplicación y el tiempo de inicialización es ligeramente mayor debido a que se guarda la caché en el disco. Recargar la caché (a partir de la segunda ejecución) reduce la inicialización a una fracción del tiempo anterior en todos los casos.
Naturalmente, la mejora de la inicialización de la sesión afecta el tiempo hasta el primer token para los LLM: lo divide aproximadamente por 2 en los puntos de referencia.
También se pueden ver las ganancias de memoria aportadas por la implementación de la caché. El tamaño máximo del conjunto de residentes se reduce para los LLM gracias a la deduplicación. En otros modelos que no se benefician de la deduplicación, no hay cambios. Recargar la caché reduce aún más el RSS máximo porque los modelos originales de TFLite ya no se leen y, por lo tanto, nunca se guardan en la memoria.
Por el momento, la memoria caché está vinculada al uso del sistema de archivos. Nuestro objetivo es aprovechar el mecanismo de deduplicación de datos de forma independiente para los casos de uso en los que no se desea intercambiar la memoria asignada tradicional con asignaciones respaldadas por archivos. Con mmap, puedes realizar asignaciones anónimas que permitirán reutilizar la mayor parte de la implementación.
