Apple Neural Engine (ANE) instead of / additionally to GPU on M1, M2 chips

username: Elias_Vansteenkiste
cooked: <p>According to the <a href="https://pytorch.org/docs/stable/notes/mps.html" rel="noopener nofollow ugc">docs</a>, MPS backend is using the GPU on M1, M2 chips via metal compute shaders.</p>
<blockquote>
<p><code>mps</code> device enables high-performance training on GPU for MacOS devices with Metal programming framework. It introduces a new device to map Machine Learning computational graphs and primitives on highly efficient Metal Performance Shaders Graph framework and tuned kernels provided by Metal Performance Shaders framework respectively.</p>
<p>The new MPS backend extends the PyTorch ecosystem and provides existing scripts capabilities to setup and run operations on GPU.</p>
</blockquote>
<p>According to the following repository, <a href="https://github.com/woolfel/ml-macos-performance/blob/main/results.md" rel="noopener nofollow ugc">ml-macos-performanc</a>, inference on the ANE is 7x faster</p>
<pre><code class="lang-auto">densenet121_keras_applications Latency ANE : 0.0012743692083333827 RPS ANE : 784.7019478034924

densenet121_keras_applications Latency GPU : 0.008270947500000033 RPS GPU : 120.90513209036763

densenet121_keras_applications Latency CPU : 0.015347813229166719 RPS CPU : 65.15586195039286
</code></pre>
<p>I was wondering would the performance of training be better if we can also use the Apple’s Neural Engine?</p>
<p>There is obviously some restrictions, see <a href="https://github.com/hollance/neural-engine/blob/master/docs/unsupported-layers.md" rel="noopener nofollow ugc">unsupported neural engine layers</a>, but that should be similar to Google’s TPUs.</p>
<p>It seems that there is maybe too much protection/security around directly calling the ANE,<br>
according to <a href="https://github.com/geohot/tinygrad/tree/master/accel/ane" rel="noopener nofollow ugc">George Hotz’s tinygrad</a>, maybe that is the biggest blocker?</p>

username: sujantkumarkv
cooked: <p>Hi,<br>
thanks for the writeup; btw the tinygrad’s link gives a 404 <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>
<p>I have been thinking to apply FlashAttention for faster training locally on macbooks but it currently only supports cuda plus MPS is less mature with implementations afaik.</p>
<p>The project is in ideation stages, <a href="https://github.com/Lightning-AI/lit-gpt/issues/429" rel="noopener nofollow ugc">here</a>.</p>
<p>I don’t have all the answers ofcourse, and this will be an opensource collaborative attempt. I’m researching what are the missing pieces I need to look for.</p>
<p>The goal is clear: “<strong>Make training faster on macbooks with Flash Attention</strong>” and may need various pieces for that: MPS, Pytorch, ANE etc.</p>
<p>I appreciate absolutely any help/comments/inputs on this from the community.</p>

username: pypdeveloper
cooked: <p>It seems Apples new ML framework MLX also doesn’t support the use of the ANE for inference, maybe they are still trying to workout the API, for python and C applications? Also is it possible to use the ANE for training?</p>