Melhorando a Performance de Quantização no TVM

A Quantização é um método usado em deep learning pra deixar os modelos menores e mais rápidos. Funciona mudando a forma como os números são armazenados no modelo. Normalmente, os modelos usam números de ponto flutuante de 32 bits, mas com a quantização, a gente troca por inteiros de 8 bits. Essa mudança ajuda a diminuir a quantidade de memória necessária e acelera os cálculos sem prejudicar muito a precisão.

#O que é o TVM?

TVM, ou Tensor Virtual Machine, é uma ferramenta de código aberto que ajuda a rodar modelos de deep learning em diferentes tipos de hardware. É feita pra ajudar os desenvolvedores de machine learning a rodar seus modelos de forma eficiente em vários dispositivos. O TVM tem duas partes principais que otimizam o Desempenho dos modelos. A primeira parte foca em como os dados fluem pelo modelo, enquanto a segunda parte melhora a maneira como o modelo usa a memória.

#O Desafio da Quantização no TVM

Mesmo com vários estudos mostrando os benefícios da quantização em reduzir o tempo e a memória usados pra rodar um modelo, a gente percebeu que a versão quantizada dos modelos no TVM não costuma performar tão bem quanto o esperado. Quando a galera usa quantização de 8 bits, geralmente espera que o modelo rode cerca de duas vezes mais rápido do que usando números de 32 bits. Mas, no TVM, essa versão quantizada às vezes rodava mais devagar, quase o dobro do tempo em comparação com a versão de precisão total.

Na nossa pesquisa, a gente investigou bem o que estava acontecendo e achou maneiras de fazer a quantização de 8 bits funcionar melhor no TVM. A gente focou em dois tipos diferentes de tarefas: aquelas onde os cálculos são o principal limite (computation-bound) e aquelas onde o uso de memória é o principal limite (memory-bound).

#Tipos de Tarefas

#Tarefas Computação-Bound

As tarefas computation-bound são aquelas em que o principal desafio vem da quantidade de potência de cálculo necessária. Esse tipo de tarefa geralmente envolve cálculos pesados, como multiplicação de matrizes ou execução de convoluções. Em machine learning, muitas tarefas se encaixam nessa categoria porque demandam um esforço computacional significativo pra fazer previsões ou treinar modelos. Por exemplo, se a gente tá rodando um modelo com um tamanho de lote de um, dá pra ver que a quantização ajuda. Isso porque usa operações aritméticas mais simples com menos demanda de recursos em comparação com os procedimentos padrão de 32 bits.

Alguns hardwares são melhores em realizar cálculos de menor precisão do que os de maior. Isso significa que cálculos usando números de 8 bits podem ser feitos mais rápido do que aqueles com números de 32 bits.

#Tarefas Memory-Bound

As tarefas memory-bound são diferentes porque são limitadas pela quantidade de dados que podem ser transferidos pra dentro e fora da memória rapidamente. Entradas grandes podem causar atrasos enquanto o sistema luta pra transferir dados entre a memória e o processador. Esse problema tende a aparecer mais com tamanhos de lote maiores (por exemplo, 8, 64 ou 256) já que essas requerem mais memória pra dados de entrada e cálculos.

Reduzindo o tamanho dos números de 32 bits pra 8 bits com a quantização, a gente pode economizar muita memória. Esse tamanho menor significa que menos dados precisam se mover de um lado pro outro entre a memória e o processador, o que pode ajudar a melhorar a velocidade.

#Configurando Experimentos

Pra ver como a quantização afeta o desempenho, a gente olhou pra um modelo chamado ResNet18 compilado pelo TVM. Rodamos esse modelo em um sistema poderoso com um CPU de 8 núcleos e uma quantidade decente de memória. Em nossos experimentos, testamos diferentes tamanhos de números e layouts pra ver como eles impactavam no desempenho. Cada teste envolveu rodar o modelo várias vezes, fazendo a média do tempo levado em cada execução.

#Corrigindo Problemas de Desempenho

Durante nossos testes, a gente percebeu que a quantização estava fazendo o modelo rodar mais devagar do que deveria. Depois de examinar a configuração, identificamos um bug que estava fazendo o modelo quantizado ter um desempenho ruim. Assim que consertamos esse bug, o modelo quantizado começou a mostrar um desempenho melhor.

A gente também descobriu que o TVM tem diferentes tipos de executores pra rodar modelos. Um executor é melhor pra modelos estáticos com operações fixas, enquanto o outro é mais adequado pra modelos dinâmicos que podem mudar. Pra nossos experimentos, a gente trocou pro executor de modelo estático, que nos permitiu otimizar melhor o modelo quantizado. Depois dessa mudança, a gente viu uma melhoria significativa no desempenho.

#Analisando o Desempenho Computação-Bound

Com o bug consertado e o executor certo no lugar, a gente olhou pra melhorar ainda mais o desempenho nas tarefas computation-bound. Focamos em otimizar as convoluções no nosso modelo, já que elas requerem muitos cálculos.

No entanto, aprendemos que nem todas as estratégias de otimização funcionam bem juntas. Diferentes configurações no TVM levam a diferentes resultados de desempenho porque algumas estratégias já estão ajustadas pra certas tarefas. As melhorias variam dependendo de como a configuração específica e os cronogramas se encaixam nas tarefas que estão sendo executadas.

Por exemplo, o packing espacial é uma técnica que ajuda a acelerar o acesso à memória pros trabalhos, mudando a forma como os dados são armazenados. O objetivo é facilitar pra que o hardware acesse os dados, o que ajuda a melhorar o desempenho. Essa mudança pode levar a um aumento significativo na velocidade dos cálculos.

#Analisando o Desempenho Memory-Bound

Além dos benefícios de desempenho dos cálculos melhorados, a quantização também ajuda no uso da memória. Usando inteiros de 8 bits ao invés de números de ponto flutuante de 32 bits, a gente pode reduzir quanto de memória o modelo precisa e com que frequência ele precisa buscar dados.

A gente notou que com tamanhos de lote maiores, as vantagens de largura de banda de memória reduzida ficaram ainda mais claras. Salvar os resultados intermediários em um formato de maior precisão ainda manteve os ganhos de desempenho da quantização de 8 bits, garantindo que a gente não perdesse precisão durante os cálculos.

#Conclusão

A quantização pode ser uma ferramenta poderosa pra melhorar a eficiência dos modelos de deep learning, especialmente quando implementada corretamente em um sistema como o TVM. Entendendo as forças e fraquezas das tarefas computation e memory-bound, a gente pode aplicar melhor a quantização pra alcançar melhorias significativas de desempenho.

Com ajuste cuidadoso e correção de problemas dentro do modelo, a gente pode transformar a quantização em um ativo em vez de um passivo. Esse trabalho abre caminhos pra mais otimizações e prepara o terreno pra usar essas técnicas poderosas em aplicações do mundo real.