ANE-friendly static llama

[metascroy]

This directory contains ANE-friendly Llama models.

Export model with:

python export.py -n /path/to/output/model.pte -p /path/to/params.json -c /path/to/model.pth --seq_length 64 --max_seq_length 1024 --coreml-quantize c4w

The runner is written in python and is only intended to serve as an example for how the model inputs should be processed; it is not performant.

Run model with:

python run.py -m /path/to/model.pte -p /path/to/params.json -t /path/to/tokenizer.model --seq_length 64 --max_seq_length 1024 --prompt "Once upon a time," --n_steps 512

The model here is based on a "sliding" cache, where old tokens are evicted from the cache. By default, the cache size is max_seq_length - seq_length, but you can explicitly pass in a smaller cache size (e.g., --cache_size 512). This can speed up computation and reduce memory. Keep in mind that once cache_size is reached, older tokens get evicted from the cache and do not participate in attention.

cc @kimishpatel @YifanShenSZ @cymbalrush

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/8436

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 92a1be8 with merge base 52a3a9a ():

NEW FAILURE - The following job has failed:

• Check Labels / Check labels (gh)
  RuntimeError: Error checking labels: PR does not have required labels

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[YIWENX14]

Add .squeeze(1) to make h of 2d shape?

[YifanShenSZ]

Generally LGTM! The major changes I found are:

1. The dedicated KV cache
2. The split linear

Anything I missed?

[YifanShenSZ]

This still looks like an index put to me? Does it successfully run on ANE?

[metascroy]

InputManager and its methods (_update_cache, etc) are not part of the model. We will have a C++ implementation of it that runs on CPU. This python implementation is intended to serve as a reference for the C++ one.

[YifanShenSZ]

Awesome! So split linear is found to be more performant on ANE? Empirically 1024 is found to be the best?

PS: On our end we found split softmax would be more performant apple/coremltools#2418

[metascroy]

We haven't done extensive testing yet on different hardware to find the right splitting values yet. I just noticed that 1024 works better on my M1 Pro recently.

On the SDPA pass: apple/coremltools#2418:

We observed something similar. We can get better Llama performance by processing tokens in smaller seq_length chunks (e.g., 256) (this not only chunks the SDPA, but all ops). This is easy enough to do, but it only chunks the Q seq_length (target_seq_length) in SDPA. It doesn't chunks the source_seq_length (which is more realistically the bigger value from the K/V caches, e.g., max_context_length). I suspect the chunking will help here too. But unlike chunking the target_seq_length, chunking the source_seq_length will require decomposing the SDPA op. Do you have plans to add support for this?

[YifanShenSZ]

Core ML multifunction is already runable now via ExecuTorch?

[metascroy]

No, I need to remove this script from the PR. I'll do an update.

[metascroy]

Generally LGTM! The major changes I found are:

1. The dedicated KV cache
2. The split linear

Anything I missed?

Those are the main changes. The structure of the KV caches and attention mask are different from how they're usually constructed. For example, the current tokens are always on the right-most side of the attention mask, whereas usually they are somewhere in the middle. This means we can get the K-value with one static concat (k = concat(k_cache, k_curr)).

The model also distinguishes between the max_seq_length and the cache_size. cache_size can be less than max_seq_length. The effect of this is older tokens are evicted from the cache and don't serve in attention.

[cccclai]

LGTM! Let's build on top of it. Also there are lots of errors in CI

[metascroy]

LGTM! Let's build on top of it. Also there are lots of errors in CI

I'll take a look at the CI failures.