smpanaro

In practice, how often do the models use the ANE? It sounds like you are optimizing for speed which in my experience always favors GPU.

For single batch inference of anything remotely LLM you'll hit the memory bound way before FLOPs, so I haven't actually looked at FLOPs much. For raw performance GPU is certainly better. ANE is more energy efficient, but you need larger batches to really benefit.

Maybe cache is the wrong word. This is a limit to how much can be mmap'd for the ANE at once. It's not too hard to hit on M1 if your model is in the GB range. Chunking the model into smaller pieces makes it more likely to "fit", but if it doesn't fit you have to unmap/remap in each forward pass which will be noticeable.

Awesome to hear about ModernBERT! Big fan of your work as well :)

coremltools is the only way to run on ANE, so less of a trick and more of a requirement.

The tricks are more around optimizing for the hardware capabilities/constraints. For instance:

- conv2d is faster than linear (see Apple's post [0]) so you rewrite the model for that (example from the repo [1])

- inputs/outputs are static shapes, so KV cache requires some creativity (I wrote about that here [2])

- compute is float16 (not bfloat16) so occasionally you have to avoid activation overflows

[0] https://machinelearning.apple.com/research/neural-engine-tra...

[1] https://github.com/Anemll/Anemll/blob/4bfa0b08183a437e759798...

[2] https://stephenpanaro.com/blog/kv-cache-for-neural-engine

What do you mean by less wide? The main bottleneck for transformers is memory bandwidth. ANE has a much lower ceiling than CPU/GPU (yes, despite unified memory).

Chunking is actually beneficial as long as all the chunks can fit into the ANE’s cache. It speeds up compilation for large network graphs and cached loads are negligible cost. On M1 the cache limit is 3-4GB, but it is higher on M2+.

Not a public follow-up but the iOS 17 speech-to-text model has a clever approach to KV caching that works within the ANE’s constraints (fixed size inputs).

I wrote about it here[0] but the gist is you can have a fixed size cache and slide it in chunks with each inference. Not as efficient as a cache that grows by one each time of course.

[0] https://stephenpanaro.com/blog/inside-apples-2023-transforme...

I bet these can all run on ANE. I’ve run gpt2-xl 1.5B on ANE [1] and WhisperKit [2] also runs larger models on it.

The smaller ones (1.1B and below) will be usably fast and with quantization I suspect the 3B one will be as well. GPU will still be faster but power for speed is the trade-off currently.

[1] 7 tokens/sec https://x.com/flat/status/1719696073751400637 [2] https://www.takeargmax.com/blog/whisperkit

Has perplexity fallen out of favor? I didn't see it mentioned anywhere. I tried using lm-eval for the 2B model but the results seem wrong (46.1288).

The jump was due to them fixing a bug. There’s a footnote about it on the bottom of page 5.

In the Discord, they mentioned a TinyLLaMa v2, presumably that would have this bug (and another bug, footnote page 4) fixed.

MobileVLM [1] is another recent small multimodal model. They trained their own 1.4B/2.7B LLaMa from scratch using RedPajama and Vicuna instead of leveraging Phi-2.

The papers only have one common benchmark (GQA, MobileVLM scores better) so hard to say how they compare otherwise.

[1] https://arxiv.org/abs/2312.16886

Multi-device: Operations can run on any of the supported devices (currently, the CPU and GPU).

Probably reading into this too much, but is this hinting at future Neural Engine support?

It’d be nice to access that without CoreML.

Yeah, I believe it is. You trade off speed for lower power usage and CPU. 8 tokens/sec is usable though.

You can do autoregressive decoding with KV caching on the Neural Engine. You have to make a bit of a trade off and use fixed size inputs [1] but the speed up over no caching is meaningful.

There's a Whisper (Encoder-Decoder) [2] implementation if you want to see it in practice. Shameless plug, but I have a repo [3] where I'm working on autoregressive text generation on the Neural Engine. I'm running gpt2-xl (1.5B params) locally with KV caching at 120ms/token (vs. 450ms without caching). Will push an update soon.

Without quantization you can't go much higher than 1.5B params on M1's Neural Engine. M2 seems to have a higher ceiling but I haven't measured. I'm optimistic (but have not tried) that the new runtime quantization added to CoreML this year will allow for larger (and maybe faster) models on both.

[1] Technically you should be able to use 1 input with an enumerated set of sizes but I haven't been able to get it to work on the Neural Engine. This would likely be even faster. [2] https://github.com/wangchou/whisper.coreml/ [3] https://github.com/smpanaro/more-ane-transformers/

I've been looking at these files too and have another data point for unilm.bundle being the new text prediction.

If you take an iOS simulator, turn off "Settings > General > Keyboard > Predictive", reboot it and then watch the console logs as you turn that switch back on, you'll see the "kbd" process load the models out of that bundle.

There is an input for temperature in the CPU model. If you can find and hook the call to predict you can probably see what is being passed.

Interestingly the Neural Engine version of the model does not take a temperature input, but it does output the raw embeddings in addition to the logits.

Sorry for the confusing phrasing about STT vs TTS. I'm not familiar with cases where you would use something like this 'at the edge' instead of say a laptop. I was thinking maybe some sort of offline setup with a microphone -- but in that case the audio is just real-time. Do you have some use cases in mind?

1/4 of the price for 1/3 of the speed is a good deal! Presumably still faster than faster-whisper on the same hardware?

What's an example use case for something like this? "At the edge" makes me think offline but are you generating audio at anything faster than real time in that case?

Would be curious to see an even lower cost/lower power option. Seems this one is $120-170.

Also not an expert, but I believe it is a little bit of both for inference.

If you are generating token-by-token naively, you do need to pay the n^2 cost since every token must attend to all other tokens. Generating a sequence of 5 tokens starting from 2 (infer 3rd, infer 4th, infer 5th...) will be much faster than starting from 1024 (infer 1025th, infer 1026th, ...) since your n is smaller. But each time time it is n^2.

However that is a naive approach. There is a common optimization, KV caching[1], (on by default for HuggingFace models) that caches all the work from the prior step so you only have to compute the attention for the new token. So you get something like (infer 1025th, cache 1025 attend new 1 token over the other 1025, cache 1026 attend 1 new token over the other 1026, ...). Not quite constant time but much better than n^2.

I would imagine there are other optimizations too, but this is the one I've heard of.

[1] fairly code-y, but links to some other posts at the start: https://www.dipkumar.dev/becoming-the-unbeatable/posts/gpt-k...