Qualcomm AI Engine Direct - Add 4-bit Embedding Quantization Option

[shewu-quic]

Summary:

• Introduce 4-bit embedding quantization for prefill, kv, and hybrid mode
• Fixed an assertion condition bug in the annotate_and_quant_scalar pass
• Refactor passes in capture_program
• Add topological sorting for passes in capture_program
• Refactor export pte flow in hybrid mode

Reproduce command

Export pte with 4-bit embedding

python examples/qualcomm/oss_scripts/llama3_2/llama.py -a ../artifacts/ -b build-android -H ${HOST} -s ${SERIAL} -m SM8650 --checkpoint ${checkpoint} --params ${param} --tokenizer_model tokenizer.model --prompt $'Could you tell me about Facebook?' --temperature 0 --model_size 1B --model_mode hybrid --prefill_seq_len 32 --kv_seq_len 512--ptq 16a4w --compile_only --embedding-quantize 4,32 --num_sharding 4

Run

python examples/qualcomm/oss_scripts/llama3_2/llama.py  -a ../artifacts/ -b build-android -H ${HOST} -s ${SERIAL}  -m SM8650 --checkpoint ${checkpoint} --params ${param} --tokenizer_model ../tokenizer.model --temperature 0 --model_size 1B --model_mode hybrid  --ptq 16a4w --prefill_seq_len 32 --kv_seq_len 512 --pre_gen_pte  ${pre_gen_pte}  --prompt "Could you tell me about Facebook?" --embedding-quantize 4,32 --num_sharding 4

Check story llama results

hybrid (kv=128, prefill=32) + 4-bit embedding (pte size: 142MB)

Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, scary dog. The dog barked and growled at her. Lily was scared and didn't know what to do.
Suddenly, a kind man came by and said, "Don't worry, I'll help you." He picked up Lily and carried her away from the dog. The man was very strong and brave.
After they were safe, the man said, "You were very brave to run away

hybrid (kv=128, prefill=32) (pte size: 223MB)

Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, scary dog. The dog barked and growled at her. Lily was scared and didn't know what to do.
Suddenly, a kind man came by and said, "Don't worry, I'll help you." He picked up Lily and carried her away from the dog. The man was very strong and brave.
After they were safe, the man said, "You were very brave to run away

prefill=32 + 4-bit embedding (pte size: 70MB)

Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a

prefill=32 (pte size: 104MB)

Once upon a time, there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a

llama 3.2 1B performance

hybrid (kv=512, prefill=32, num_sharding=4) + 4-bit embedding

I 00:00:01.529260 executorch:runner.cpp:456] 	Prompt Tokens: 17    Generated Tokens: 28
I 00:00:01.529274 executorch:runner.cpp:462] 	Model Load Time:		1.062000 (seconds)
I 00:00:01.529286 executorch:runner.cpp:472] 	Total inference time:		0.464000 (seconds)		 Rate: 	60.344828 (tokens/second)
I 00:00:01.529295 executorch:runner.cpp:480] 		Prompt evaluation:	0.042000 (seconds)		 Rate: 	404.761905 (tokens/second)
I 00:00:01.529303 executorch:runner.cpp:491] 		Generated 28 tokens:	0.422000 (seconds)		 Rate: 	66.350711 (tokens/second)
I 00:00:01.529313 executorch:runner.cpp:499] 	Time to first generated token:	0.042000 (seconds)
I 00:00:01.529321 executorch:runner.cpp:506] 	Sampling time over 28 tokens:	0.026000 (seconds)

hybrid (kv=512, prefill=32, num_sharding=4)

I 00:00:01.798075 executorch:runner.cpp:456] 	Prompt Tokens: 17    Generated Tokens: 39
I 00:00:01.798088 executorch:runner.cpp:462] 	Model Load Time:		1.171000 (seconds)
I 00:00:01.798098 executorch:runner.cpp:472] 	Total inference time:		0.624000 (seconds)		 Rate: 	62.500000 (tokens/second)
I 00:00:01.798107 executorch:runner.cpp:480] 		Prompt evaluation:	0.040000 (seconds)		 Rate: 	425.000000 (tokens/second)
I 00:00:01.798115 executorch:runner.cpp:491] 		Generated 39 tokens:	0.584000 (seconds)		 Rate: 	66.780822 (tokens/second)
I 00:00:01.798123 executorch:runner.cpp:499] 	Time to first generated token:	0.040000 (seconds)
I 00:00:01.798130 executorch:runner.cpp:506] 	Sampling time over 39 tokens:	0.029000 (seconds)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit e530130 with merge base a5c7609 ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-eval_llama-mmlu-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / test-llava-runner-linux / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[shewu-quic]

Hi @cccclai ,

This PR is to support 4-bit embedding cpu on static llama and refactor the passes in capture_program.
Could you help please take a look?

Thanks

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[cccclai]

Looks good - probably need to rebase because #7618 is merged.

[cccclai]

Did you run into any issue because the pass order isn't correct?

[shewu-quic]

No, this function is to prevent users from unexpectedly modifying the order of passes in scripts.
I think it would be flexible that user wanna introduce their own implementation.

[shewu-quic]

Hi @cccclai,
Thanks for review! Just rebased.

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[cccclai]

It needs rebase again, given #7618 just reland and #7694 is also landed

[facebook-github-bot]

@cccclai has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[cccclai]

Looks good, thank you!