Explore as a Storm, Exploit as a Raindrop: On the Benefit of Fine-Tuning Kernel Schedulers with Coordinate Descent

要約

機械学習モデルはカーネルで構成されます。カーネルは、テンソル (自然数の線形結合によってインデックス付けされたデータ) に演算を適用するアルゴリズムです。
カーネルの例には、畳み込み、転置、ベクトル積などがあります。
カーネルを実装するにはさまざまな方法があります。
これらの実装は、カーネルの最適化スペースを形成します。
カーネルのスケジューリングは、目的関数 (通常は実行速度) を考慮して、最適な実装を見つける問題です。
Ansor、Halide、AutoTVM などのカーネル オプティマイザーは、探索と活用の 2 つのフェーズを組み合わせた検索ヒューリスティックによってこの問題を解決します。
最初のステップでは、さまざまなカーネル最適化スペースを評価します。
後者は、同じ空間内のカーネルを調査することによって、最適な実装を改善しようとします。
たとえば、Ansor は探索のためのスケッチによるカーネル生成を組み合わせ、進化的アルゴリズムを活用して最良のスケッチを活用します。
この研究では、AutoTVM アルゴリズムである Droplet Search を Ansor の探索フェーズに組み込むことで、カーネルの品質を向上させながら Ansor の検索時間を短縮できる可能性を実証します。
このアプローチには、Ansor によって調査されるサンプルの数を制限し、最適なものを選択し、それを座標降下アルゴリズムで活用することが含まれます。
Ansor が生成する最初の 300 カーネルにこのアプローチを適用すると、通常、Ansor に 10,000 カーネルを分析させるよりも短い時間でより良いカーネルが得られます。
この結果は、AMD Ryzen 7 (x86)、NVIDIA A100 tensor コア、NVIDIA RTX 3080 GPU の 4 つのアーキテクチャ上で実行される 20 のよく知られたディープラーニング モデル (AlexNet、ResNet、VGG、DenseNet など) で再現されています。
、ARM A64FX。
この組み合わせアプローチを使用したパッチは、2024 年 2 月に Ansor で承認されました。この検索方法の汎用性の証拠として、同様のパッチが同様に良好な結果を達成し、2024 年 6 月に TVM の MetaSchedule に提出されました。

要約(オリジナル)

Machine-learning models consist of kernels, which are algorithms applying operations on tensors — data indexed by a linear combination of natural numbers. Examples of kernels include convolutions, transpositions, and vectorial products. There are many ways to implement a kernel. These implementations form the kernel’s optimization space. Kernel scheduling is the problem of finding the best implementation, given an objective function — typically execution speed. Kernel optimizers such as Ansor, Halide, and AutoTVM solve this problem via search heuristics, which combine two phases: exploration and exploitation. The first step evaluates many different kernel optimization spaces. The latter tries to improve the best implementations by investigating a kernel within the same space. For example, Ansor combines kernel generation through sketches for exploration and leverages an evolutionary algorithm to exploit the best sketches. In this work, we demonstrate the potential to reduce Ansor’s search time while enhancing kernel quality by incorporating Droplet Search, an AutoTVM algorithm, into Ansor’s exploration phase. The approach involves limiting the number of samples explored by Ansor, selecting the best, and exploiting it with a coordinate descent algorithm. By applying this approach to the first 300 kernels that Ansor generates, we usually obtain better kernels in less time than if we let Ansor analyze 10,000 kernels. This result has been replicated in 20 well-known deep-learning models (AlexNet, ResNet, VGG, DenseNet, etc.) running on four architectures: an AMD Ryzen 7 (x86), an NVIDIA A100 tensor core, an NVIDIA RTX 3080 GPU, and an ARM A64FX. A patch with this combined approach was approved in Ansor in February 2024. As evidence of the generality of this search methodology, a similar patch, achieving equally good results, was submitted to TVM’s MetaSchedule in June 2024.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google