Add torchao mps lowbit ops to llama runner

[manuelcandales]

Setup ET: https://pytorch.org/executorch/stable/getting-started-setup

Install llama runner requirements

bash examples/models/llama/install_requirements.sh

Build ET with MPS ON:

cmake -DPYTHON_EXECUTABLE=python \
    -DCMAKE_INSTALL_PREFIX=cmake-out \
    -DEXECUTORCH_ENABLE_LOGGING=1 \
    -DCMAKE_BUILD_TYPE=Release \
    -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
    -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
    -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
    -DEXECUTORCH_BUILD_MPS=ON \
    -DEXECUTORCH_BUILD_XNNPACK=ON \
    -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
    -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
    -Bcmake-out .
cmake --build cmake-out -j16 --target install --config Release

Build llama runner with torchao mps ops

cmake -DPYTHON_EXECUTABLE=python \
    -DCMAKE_PREFIX_PATH=$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())') \
    -DCMAKE_INSTALL_PREFIX=cmake-out \
    -DCMAKE_BUILD_TYPE=Release \
    -DEXECUTORCH_BUILD_KERNELS_CUSTOM=ON \
    -DEXECUTORCH_BUILD_TORCHAO=ON \
    -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \
    -DEXECUTORCH_BUILD_XNNPACK=ON \
    -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
    -Bcmake-out/examples/models/llama \
    examples/models/llama
cmake --build cmake-out/examples/models/llama -j16 --target install --config Release

Export model. Note: qmode can be any of torchao:fpa1w ... torchao:fpa7w, and group_size can be any of: 32, 64, 128, 256

CMAKE_INSTALL_PREFIX=$PWD/cmake-out python -m examples.models.llama.export_llama \
--checkpoint /path/to/model.pth \
--params  /path/to/params.json \
-kv --use_sdpa_with_kv_cache --disable_dynamic_shape --metadata '{"get_bos_id":128000, "get_eos_ids":[128009, 128001]}' \
-qmode "torchao:fpa4w" --group_size 32 \
--output_name path/to/output.pte

Run:

cmake-out/examples/models/llama/llama_main --model_path=/path/to/output.pte --tokenizer_path=/path/to/tokenizer.model --prompt="Once upon a time,"

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Cancelled Job

As of commit d47141f with merge base daf9aee ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / linux / linux-job (gh)
  curl: (22) The requested URL returned error:
• pull / unittest-arm / linux-job (gh)
  RuntimeError: Command docker exec -t bcbce14150776a156ffad265b9508db49218afdbf2eddd961326cf6e81234108 /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / test-mediatek-models-linux / linux-job (gh)
  ##[error]The operation was canceled.

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

[kimishpatel]

Some nits but largely looks good. Please include the output for llama runner in the summary.

Consider fixing the nits. Particularly I dont think we should add TORCHAO_MPS. Just overload TORCHAO one

[kimishpatel]

Please using logging