Optimizing machine learning models: a droplet search approach to efficient kernel scheduling

Abstract

O escalonamento do kernel, um estágio crítico da otimização computacional, envolve a tarefa de determinar a implementação mais eficiente para um kernel computacional. Essa jornada de otimização envolve experimentar parâmetros do compilador, como tamanhos de janelas lado a lado e fatores de desenrolamento, para discernir a configuração ideal. A abordagem introduzida neste trabalho conceitua esses parâmetros como pontos dentro de um espaço de coordenadas cuidadosamente definido. Embora a função que mapeia esses pontos para os tempos de execução do kernel normalmente resulte em uma superfície não convexa, a evidência empírica apoia uma nova hipótese denominada "droplet expectation". Esta hipótese postula que a origem desta superfície, representando um kernel não otimizado, e seu ótimo global, significando o kernel mais rápido, residem coletivamente dentro dos limites de uma região convexa. Aproveitando esse insight, o trabalho propõe a metodologia Droplet Search, baseada no algoritmo de descida por coordenadas. Assumindo a validade da expectativa do droplet, o Droplet Search demonstra a capacidade de convergir rapidamente para a configuração ideal do kernel. Implementado no Apache TVM desde abril de 2023, o Droplet Search passa por uma avaliação rigorosa por meio de extensa experimentação envolvendo seis grandes modelos de aprendizagem profunda implantados em uma variedade de dispositivos de computação, incluindo ARM, Intel, AMD e NVIDIA. Os resultados obtidos desses experimentos mostram que o Droplet Search não apenas corresponde à eficácia de outras técnicas de pesquisa do AutoTVM, mas também as supera em eficiência, apresentando uma melhoria de velocidade que varia entre duas a dez vezes. Além de seus recursos de otimização acelerada, os modelos gerados pelo Droplet Search exibem uma vantagem competitiva quando justapostos aos produzidos pelo AutoScheduler (Ansor) do TVM. Esta é uma revelação digna de nota, considerando que o Ansor emprega quatro a cinco vezes mais transformações de código do que o AutoTVM. As descobertas empíricas aqui apresentadas ressaltam a trajetória promissora do Droplet Search como uma ferramenta potente no domínio do agendamento do kernel, redefinindo benchmarks ao oferecer processos de otimização acelerados e resultados competitivos, tudo realizado com uma sobrecarga computacional visivelmente reduzida.