Try some QD8-BF16 Experiments

[mcr229]

We've prototyped some new QD8-BF16-QB4W kernels in XNNPACK. Let's try leveraging them in ExecuTorch and see what our performance looks like:

Exports:

f32
python -m examples.models.llama.export_llama --checkpoint ~/Desktop/oss/models/llama3.2-1b/Llama3.2-1B/consolidated.00.pth -p ~/Desktop/oss/models/llama3.2-1b/Llama3.2-1B/params.json -d fp32 -kv -X -qmode 8da4w --group_size 64 --use_sdpa_with_kv_cache -n llama3_qd8_f32.pte --verbose

bf16
python -m examples.models.llama.export_llama --checkpoint ~/Desktop/oss/models/llama3.2-1b/Llama3.2-1B/consolidated.00.pth -p ~/Desktop/oss/models/llama3.2-1b/Llama3.2-1B/params.json -d bf16 -kv -X -qmode 8da4w --group_size 64 --use_sdpa_with_kv_cache -n llama3_qd8_bf16.pte  --verbose

There is no change in size since this only affects activations. To make the comparisons more fair, we removed delegation of all other operators aside from qd8-bf16-qb4w. This is because in XNNPACK we are still lacking those bf16 operators.

e1q:/data/local/tmp/bf16 $ ./llama_main --model_path=llama3_qd8_f32.pte --tokenizer_path=../llama_artifacts/tokenizer.model                                                                                                                                                                                   
The answer to the ultimate question isI 00:00:03.375228 executorch:text_prefiller.cpp:96] Prefill token result numel(): 128256
:I 00:00:03.375524 executorch:runner.cpp:282] RSS after prompt prefill: 1812.015625 MiB (0 if unsupported)
 the original three . If you can’t find a single instance of an animal or a person who doesn’t belong to one of those three groups, then you don’t belong to one of those three groups.
We can see the evidence of our species by observing our culture. There is no such thing as a ‘random’ or ‘pure’ culture. Even in cultures where a man can beat a woman and the woman will let him, she is a woman. A man can be raped, but a woman is raped. A man can kill someone, but a woman is killed. A man can get
I 00:00:05.265945 executorch:runner.cpp:302] RSS after finishing text generation: 1812.015625 MiB (0 if unsupported)
PyTorchObserver {"prompt_tokens":7,"generated_tokens":120,"model_load_start_ms":1734622925574,"model_load_end_ms":1734622928866,"inference_start_ms":1734622928866,"inference_end_ms":1734622930826,"prompt_eval_end_ms":1734622928936,"first_token_ms":1734622928936,"aggregate_sampling_time_ms":138,"SCALING_FACTOR_UNITS_PER_SECOND":1000}
I 00:00:05.265997 executorch:stats.h:108]       Prompt Tokens: 7    Generated Tokens: 120
I 00:00:05.265999 executorch:stats.h:114]       Model Load Time:                3.292000 (seconds)
I 00:00:05.266002 executorch:stats.h:124]       Total inference time:           1.960000 (seconds)               Rate:  61.224490 (tokens/second)
I 00:00:05.266014 executorch:stats.h:132]               Prompt evaluation:      0.070000 (seconds)               Rate:  100.000000 (tokens/second)
I 00:00:05.266018 executorch:stats.h:143]               Generated 120 tokens:   1.890000 (seconds)               Rate:  63.492063 (tokens/second)
I 00:00:05.266025 executorch:stats.h:151]       Time to first generated token:  0.070000 (seconds)
I 00:00:05.266036 executorch:stats.h:158]       Sampling time over 127 tokens:  0.138000 (seconds)

e1q:/data/local/tmp/bf16 $ ./llama_main --model_path=llama3_qd8_bf16.pte --tokenizer_path=../llama_artifacts/tokenizer.model
The answer to the ultimate question isI 00:00:03.217430 executorch:text_prefiller.cpp:96] Prefill token result numel(): 128256
:I 00:00:03.217586 executorch:runner.cpp:282] RSS after prompt prefill: 1305.042969 MiB (0 if unsupported)
 “I don’t know.”<|end_of_text|><|begin_of_text|>
The answer to the ultimate question is: “I don’t know.”
“Then there’s a little problem: We’re living in an era of faith and a period of disbelief.”
A woman at a press conference says that it is too early to declare that a no-fault approach is right. In a situation of so many possibilities, how can we possibly make a decision?
The answer to the ultimate question is: “I don’t know.”
The answer to the ultimate question is: “I don’t know.”
The answer to the ultimate question is:
I 00:00:05.314108 executorch:runner.cpp:302] RSS after finishing text generation: 1305.042969 MiB (0 if unsupported)
PyTorchObserver {"prompt_tokens":7,"generated_tokens":120,"model_load_start_ms":1734622914462,"model_load_end_ms":1734622917579,"inference_start_ms":1734622917579,"inference_end_ms":1734622919757,"prompt_eval_end_ms":1734622917661,"first_token_ms":1734622917661,"aggregate_sampling_time_ms":207,"SCALING_FACTOR_UNITS_PER_SECOND":1000}
I 00:00:05.316944 executorch:stats.h:108]       Prompt Tokens: 7    Generated Tokens: 120
I 00:00:05.316948 executorch:stats.h:114]       Model Load Time:                3.117000 (seconds)
I 00:00:05.316951 executorch:stats.h:124]       Total inference time:           2.178000 (seconds)               Rate:  55.096419 (tokens/second)
I 00:00:05.316955 executorch:stats.h:132]               Prompt evaluation:      0.082000 (seconds)               Rate:  85.365854 (tokens/second)
I 00:00:05.316958 executorch:stats.h:143]               Generated 120 tokens:   2.096000 (seconds)               Rate:  57.251908 (tokens/second)
I 00:00:05.316961 executorch:stats.h:151]       Time to first generated token:  0.082000 (seconds)
I 00:00:05.316964 executorch:stats.h:158]       Sampling time over 127 tokens:  0.207000 (seconds)

Some things to notice here is that the BF16 model uses 1/3 of the memory as the fp32 model. Additionally, we see some performance drops in BF16. This is likely because the GEMM kernel contains an extra shift to perform bf16 (things are still calculated in f32, just right shifted before storing). Additionaly the Quantize kernel for bf16 --> qd8 is still a naive implementation so it is a bit slower. Another thing to notice is that the results seem to be nonsensical.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit d4845a0 with merge base fff7b3c ():
💚 Looks good so far! There are no failures yet. 💚

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

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

[digantdesai]

can we check against bf16 llm w/ and w/o xnnpack qd8-bf16-qb4w (running q/dq if possible or even non quantized)?

[meta-cla]

Hi @mcr229!

Thank you for your pull request.

We require contributors to sign our Contributor License Agreement, and yours needs attention.

You currently have a record in our system, but the CLA is no longer valid, and will need to be resubmitted.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.