Skip to content

[ET-VK] Enable Partial GPU lowering via Vulkan in stories model export #2368

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 4 commits into from
Closed

[ET-VK] Enable Partial GPU lowering via Vulkan in stories model export #2368

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e4bec07
[ET-VK] Enable Partial GPU lowering via Vulkan in stories model export
SS-JIA Mar 12, 2024
5ad17d3
Update on "[ET-VK] Enable Partial GPU lowering via Vulkan in stories …
SS-JIA Mar 12, 2024
d4507ea
Update on "[ET-VK] Enable Partial GPU lowering via Vulkan in stories …
SS-JIA Mar 12, 2024
81b6ade
Update on "[ET-VK] Enable Partial GPU lowering via Vulkan in stories …
SS-JIA Mar 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions examples/models/llama2/TARGETS
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@ runtime.python_library(
"//executorch/backends/transforms:duplicate_dynamic_quant_chain",
"//executorch/backends/xnnpack:xnnpack_backend",
"//executorch/backends/xnnpack/partition:xnnpack_partitioner",
"//executorch/backends/vulkan/partitioner:vulkan_partitioner",
"//executorch/examples/models:model_base",
"//executorch/examples/models:models",
"//executorch/examples/portable:utils",
Expand Down
13 changes: 13 additions & 0 deletions examples/models/llama2/export_llama_lib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import pkg_resources
import torch
from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
XnnpackDynamicallyQuantizedPartitioner,
)
Expand Down Expand Up @@ -359,6 +360,7 @@ def build_args_parser() -> argparse.ArgumentParser:
parser.add_argument("-2", "--fairseq2", action="store_true")
parser.add_argument("-v", "--verbose", action="store_true")
parser.add_argument("-X", "--xnnpack", action="store_true")
parser.add_argument("-V", "--vulkan", action="store_true")

parser.add_argument(
"--generate_etrecord",
Expand Down Expand Up @@ -463,6 +465,17 @@ def _export_llama(modelname, args) -> str: # noqa: C901
# partitioners[XnnpackPartitioner.__name__] = XnnpackPartitioner()
modelname = f"xnnpack_{modelname}"

if args.vulkan:
assert (
args.dtype_override is None
), "Vulkan backend does not support non fp32 dtypes at the moment"
assert (
args.quantization_mode is None
), "Vulkan backend does not support quantization at the moment"

partitioners[VulkanPartitioner.__name__] = VulkanPartitioner()
modelname = f"vulkan_{modelname}"

builder_exported_to_edge = (
load_llama_model(
checkpoint=checkpoint_path,
Expand Down
6 changes: 5 additions & 1 deletion examples/models/llama2/runner/targets.bzl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,11 @@ def define_common_targets():
"//executorch/extension/module:module" + aten_suffix,
"//executorch/kernels/quantized:generated_lib" + aten_suffix,
"//executorch/runtime/core/exec_aten:lib" + aten_suffix,
] + (_get_operator_lib(aten)),
] + (_get_operator_lib(aten)) + ([
# Vulkan API currently cannot build on some platforms (e.g. Apple, FBCODE)
# Therefore enable it explicitly for now to avoid failing tests
"//executorch/backends/vulkan:vulkan_backend_lib",
] if native.read_config("llama", "use_vulkan", "0") == "1" else []),
external_deps = [
"libtorch",
] if aten else [],
Expand Down