O XNNPack é o mecanismo padrão de inferência de CPU do TensorFlow Lite para todos os modelos. Ele oferece melhorias de velocidade inovadoras em plataformas da Web, para dispositivos móveis e para computadores. Uma das otimizações empregadas no XNNPack é o reempacotamento de pesos estáticos dos operadores Convolution, Depthwise Convolution, Transposed Convolution e Fully Connected em um layout interno otimizado para cálculos de inferência. Durante a inferência, os pesos reempacotados são acessados em um padrão sequencial que é amigável para os pipelines dos processadores.
A redução da latência de inferência tem um custo: basicamente, o reempacotamento cria uma cópia extra dos pesos dentro do XNNPack. Já houve tentativas de reduzir esse custo com a adição de um cache em memória ao XNNPack. Esse cache permite compartilhar os pesos empacotados entre interpretadores independentes do TFLite que executariam o mesmo modelo de maneira autônoma.
A implementação do delegado XNNPack do TFLite foi aprimorada para lidar com algumas das deficiências do cache existente.
1. O cache reside na memória anônima, o que pode resultar em troca para o disco em caso de pressão da memória, levando ao baixo desempenho.
2. Necessidade de reempacotar os pesos iniciais sempre que um processo é iniciado.
3. Como o reempacotamento lê os pesos originais do TFLite e grava-os em um novo buffer, isso leva a um alto pico de uso da memória durante o empacotamento.
4. O cache requer a execução de etapas demoradas e o gerenciamento cuidadoso do ciclo de vida para permitir adequadamente o armazenamento em cache por meio do delegado do XNNPack.
5. O cache não permite compartilhar os pesos entre processos.
O XNNPack foi atualizado e fornece uma interface que permite implementar um provedor de cache de pesos. Um provedor de cache de pesos se comporta como um dicionário que o XNNPack preencherá e consultará para acessar buffers empacotados. Estas são as principais funções:
look_up: busca uma chave de buffer empacotada e retorna um identificador exclusivo (ou um identificador especial reservado para NotFound) que pode ser usado posteriormente para recuperar o endereço de buffer.reserve_space: reserva um buffer que pode ser usado para armazenar informações de um determinado tamanho. Esse buffer precisa ser confirmado usando a função look_up_or_insert.look_up_or_insert: verifica se um buffer correspondente à chave fornecida existe no provedor de cache. Se não existir, os dados fornecidos serão confirmados no provedor de cache. Essa função também retorna o identificador que pode ser usado para recuperar o endereço de buffer.offset_to_addr: retorna o endereço de buffer do identificador retornado por look_up e look_up_or_insert.As interações entre o XNNPack e o provedor de cache de pesos são ilustradas no diagrama a seguir.
O delegado do TFLite agora usa essa nova interface e tem seu próprio provedor de cache de pesos. Esse provedor é capaz de salvar e carregar os pesos empacotados diretamente no disco ou a partir dele. O TFLite já usa o flatbuffer e o mapeamento de memória baseado em arquivo há muito tempo. Estamos preenchendo a lacuna ao adotar a mesma técnica para obter as seguintes vantagens:
A persistência dos pesos empacotados no disco pula o dispendioso processo de reempacotamento a cada vez que um modelo é carregado. Isso resulta na redução significativa da latência de inicialização e do pico de uso da memória. Mesmo para a criação inicial, isso permite a eliminação de duplicação de dados empacotados e melhora ainda mais o desempenho do empacotamento ao evitar o reempacotamento dos mesmos dados.
O mmap usa o gerenciamento de memória virtual do sistema operacional para otimizar o desempenho e o uso geral da memória do sistema. Em nosso caso, isso é muito vantajoso para o acesso aleatório a arquivos somente leitura de grande volume, como os pesos constantes da operação de uma rede neural, por exemplo.
Com os dados empacotados armazenados no disco, o cache do XNNPack não depende mais da memória anônima, que pode estar sujeita a problemas de desempenho sob pressão da memória. Em vez disso, ele usa o gerenciamento de memória virtual do sistema operacional para uma operação mais fluida.
Ao eliminar a necessidade de copiar dados entre o sistema de arquivos e a memória, o mmap reduz significativamente o overhead e acelera os tempos de acesso.
Mais informações sobre mapeamentos de arquivos e uso da memória estão disponíveis diretamente na página principal do mmap e em outros estudos interessantes.
O carregamento de arquivos baseado no mmap possibilita o compartilhamento uniforme de pesos entre vários processos, já que o espaço de endereço virtual de cada processo é mapeado para as mesmas páginas da memória física. Isso reduz o tamanho geral da memória, uma vez que vários processos compartilham a mesma memória, e acelera o carregamento do modelo de maneira generalizada.
Em vez de exigir que o usuário configure e gerencie o objeto de cache durante toda a vida útil do aplicativo, ele pode simplesmente fornecer um caminho para o arquivo de cache.
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 garantir uma inferência precisa e eficiente, é essencial invalidar o cache do XNNPack sob condições específicas:
Evolução do modelo: se os pesos ou a estrutura do modelo mudarem, os dados armazenados em cache ficarão desatualizados e deverão ser invalidados. Isso significa remover o arquivo no caminho do cache fornecido.
Upgrades do XNNPack: atualizações no algoritmo de empacotamento interno do XNNPack podem resultar em pesos em cache incompatíveis, exigindo que o cache seja recalculado. Felizmente, o XNNPack é capaz de detectar isso e substitui o cache existente automaticamente.
Basicamente, qualquer modificação que possa afetar a maneira como os pesos são empacotados ou utilizados pelo XNNPack deve acionar uma invalidação de cache.
A inicialização de sessão é dominada pelo empacotamento de pesos. Para LLMs, vários subgrafos estão reutilizando os mesmos pesos. A criação do cache é mais rápida porque a funcionalidade de eliminação de duplicação evita empacotar esses mesmos pesos várias vezes. Para modelos mais comuns, como o Stable Diffusion, não há eliminação de duplicação, e o tempo de inicialização um pouco maior ocorre porque o cache é salvo no disco. Recarregar o cache (a partir da segunda execução) reduz a inicialização a uma fração do tempo anterior em todos os casos.
A melhoria da inicialização da sessão naturalmente afeta o tempo até o primeiro token para LLMs, dividindo-o aproximadamente por dois nos comparativos de mercado.
Também podem ser observados ganhos de memória provenientes da implementação do cache. O pico de tamanho do conjunto residente (RSS) é reduzido para LLMs graças à eliminação de duplicação. Para outros modelos que não se beneficiam da eliminação de duplicação, não há mudanças. Recarregar o cache reduz ainda mais o pico de RSS porque os modelos originais do TFLite não são mais lidos e, portanto, nunca são carregados na memória.
Atualmente, o cache está vinculado ao uso do sistema de arquivos. Nosso objetivo é conseguir aproveitar o mecanismo de eliminação de duplicação de dados de maneira independente para casos de uso que não desejam trocar a memória alocada tradicional por mapeamentos baseados em arquivo. Com o mmap, é possível fazer mapeamentos anônimos que permitem reutilizar a maior parte da implementação.
