diff --git a/.ci/scripts/test-native.sh b/.ci/scripts/test-native.sh index ebd4277d4bc..2f08fc6dabf 100644 --- a/.ci/scripts/test-native.sh +++ b/.ci/scripts/test-native.sh @@ -20,17 +20,22 @@ case "${component}" in test_regex='^native_runtime_test$' ;; ggml) - cmake_args+=(-DEXECUTORCH_BUILD_NATIVE_GGML=ON -DGGML_NATIVE=OFF) - targets+=(native_ggml_engine_test) - test_regex='^native_ggml_engine_test$' + cmake_args+=( + -DEXECUTORCH_BUILD_NATIVE_GGML=ON + -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON + -DGGML_NATIVE=OFF + ) + targets+=(run_hf_llm_ggml native_ggml_engine_test native_llm_runner_utils_test) + test_regex='^native_(ggml_engine|llm_runner_utils)_test$' ;; vulkan) cmake_args+=( -DEXECUTORCH_BUILD_NATIVE_VULKAN=ON -DEXECUTORCH_BUILD_NATIVE_VISION_RUNNER=ON + -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON -DEXECUTORCH_VULKAN_SHADER_COMPILE_NTHREADS="${build_jobs}" ) - targets+=(native_vulkan_engine run_classification native_vulkan_passes_test native_runner_utils_test) + targets+=(native_vulkan_engine run_classification run_hf_llm_vulkan native_vulkan_passes_test native_runner_utils_test) test_regex='^native_(vulkan_passes|runner_utils)_test$' ;; *) @@ -53,3 +58,10 @@ cmake --build "${build_dir}" \ --target "${targets[@]}" ctest --test-dir "${build_dir}/backends/native/test" \ -R "${test_regex}" --output-on-failure --no-tests=error --timeout 120 + +if [[ "${component}" == "ggml" ]]; then + PYTHONPATH="${PWD}/src:${PYTHONPATH:-}" \ + FLATC_EXECUTABLE="${FLATC_EXECUTABLE:-${PWD}/${build_dir}/third-party/flatc_ep/bin/flatc}" \ + python .ci/scripts/test_native_llm_runner.py \ + --runner "${build_dir}/backends/native/examples/hf_llm/runtime/run_hf_llm_ggml" +fi diff --git a/.ci/scripts/test_native_llm_runner.py b/.ci/scripts/test_native_llm_runner.py new file mode 100644 index 00000000000..9fe9eadf638 --- /dev/null +++ b/.ci/scripts/test_native_llm_runner.py @@ -0,0 +1,109 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Compare GGML LLM runner prefill and decoding with a tiny eager HF Llama.""" + +import argparse +import base64 +import re +import subprocess +import tempfile +from pathlib import Path + +import torch +from executorch.backends.native.examples.hf_llm.export_hf_llm import LastTokenLM +from executorch.exir.native import to_native +from transformers import LlamaConfig, LlamaForCausalLM + + +@torch.no_grad() +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--runner", required=True, type=Path) + args = parser.parse_args() + + torch.manual_seed(0) + torch.set_num_threads(4) + max_seq_length = 48 + max_new_tokens = 8 + config = LlamaConfig( + vocab_size=512, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=max_seq_length, + bos_token_id=256, + eos_token_id=257, + attn_implementation="sdpa", + ) + model = LlamaForCausalLM(config).eval() + model.generation_config.use_cache = True + model.generation_config.cache_implementation = "static" + wrapper = LastTokenLM(model, max_seq_length).eval() + seq = torch.export.Dim("seq", min=1, max=max_seq_length - 1) + exported = torch.export.export( + wrapper, + (torch.zeros((1, 6), dtype=torch.long), torch.arange(6)), + dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}}, + strict=True, + ) + + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + package = root / "llama.ptn" + to_native(exported).save(str(package)) + tokenizer = root / "tokenizer.model" + # The byte vocabulary puts Llama's BOS at 256, followed by 255 specials. + tokenizer.write_text( + "".join( + f"{base64.b64encode(bytes([i])).decode()} {i}\n" for i in range(256) + ) + ) + prompt = "abc" + unreachable_eos_id = config.vocab_size + tokens = [config.bos_token_id, *prompt.encode()] + expected = [] + # Recompute the full prefix without a cache for an independent reference. + for _ in range(max_new_tokens): + logits = model(torch.tensor([tokens]), use_cache=False).logits + token = int(logits[0, -1].argmax()) + expected.append(token) + tokens.append(token) + + result = subprocess.run( + [ + str(args.runner.resolve()), + "--device=cpu", + "--cpu_threads=4", + f"--package={package}", + f"--tokenizer={tokenizer}", + f"--prompt={prompt}", + f"--eos_ids={unreachable_eos_id}", + f"--max_new_tokens={max_new_tokens}", + f"--max_seq_length={max_seq_length}", + "--print_token_ids", + ], + capture_output=True, + text=True, + errors="replace", + timeout=120, + ) + print(result.stdout, end="") + print(result.stderr, end="") + result.check_returncode() + match = re.search(r"^token_ids: ([0-9,]+)$", result.stdout, re.MULTILINE) + if match is None: + raise AssertionError("Runner did not report generated token IDs") + actual = [int(token) for token in match[1].split(",")] + if actual != expected: + raise AssertionError(f"Generated tokens {actual} != eager {expected}") + print(f"GGML HF LLM runner matched all {max_new_tokens} eager token IDs") + + +if __name__ == "__main__": + main() diff --git a/.github/workflows/test-native.yml b/.github/workflows/test-native.yml index 6dd7e5db718..c5e0f0b7c76 100644 --- a/.github/workflows/test-native.yml +++ b/.github/workflows/test-native.yml @@ -7,8 +7,13 @@ on: - .github/workflows/_docker-image.yml - .ci/docker/** - .ci/scripts/test-native.sh + - .ci/scripts/test_native_llm_runner.py - .gitmodules - CMakeLists.txt + - install_requirements.py + - install_requirements.sh + - requirements-dev.txt + - requirements-examples.txt - backends/native/** - backends/vulkan/CMakeLists.txt - backends/vulkan/cmake/** @@ -25,6 +30,10 @@ on: - backends/xnnpack/third-party/FXdiv - backends/xnnpack/third-party/cpuinfo - backends/xnnpack/third-party/pthreadpool + - examples/models/llama/tokenizer/** + - extension/llm/tokenizers + - extension/llm/export/** + - exir/** - runtime/** - third-party/** - tools/cmake/** @@ -93,4 +102,12 @@ jobs: export CMAKE_PREFIX_PATH="/tmp/native-deps:${CMAKE_PREFIX_PATH:-}" export PATH="/tmp/native-deps/bin:${PATH}" + if [[ "${NATIVE_COMPONENT}" == "ggml" ]]; then + ./install_requirements.sh --use-pt-pinned-commit + # The LLM runner test imports executorch from src/ without installing the + # wheel, so install its export-time dependencies directly. + pip install flatbuffers ruamel.yaml tabulate \ + "$(grep -E '^transformers\b' requirements-examples.txt)" + fi + CMAKE_BUILD_PARALLEL_LEVEL=8 bash .ci/scripts/test-native.sh "${NATIVE_COMPONENT}" diff --git a/CMakeLists.txt b/CMakeLists.txt index 7ae92a457b5..3768d04a083 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -980,7 +980,10 @@ if(EXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP) endif() if(EXECUTORCH_BUILD_EXTENSION_LLM) - if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER) + if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER + OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER + ) + # Both LLM runners need tokenizers with regex lookahead support. set(SUPPORT_REGEX_LOOKAHEAD ON) # llama/runner/CMakeLists.txt builds a shared library libllama_runner.so # that transitively depends on tokenizers. Need to build tokenizers with @@ -991,7 +994,9 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) set(CMAKE_POSITION_INDEPENDENT_CODE ON) endif() add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/tokenizers) - if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER) + if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER + OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER + ) set(CMAKE_POSITION_INDEPENDENT_CODE ${ORIGINAL_CMAKE_POSITION_INDEPENDENT_CODE_FLAG} ) diff --git a/backends/native/CMakeLists.txt b/backends/native/CMakeLists.txt index e111c8f5453..2b4b8ec1a7c 100644 --- a/backends/native/CMakeLists.txt +++ b/backends/native/CMakeLists.txt @@ -36,7 +36,7 @@ install( PATTERN "ggml" EXCLUDE ) -if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) +if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER OR EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) add_subdirectory(examples) endif() diff --git a/backends/native/README.md b/backends/native/README.md index cb1f5af1eed..6fafb89f11d 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -33,6 +33,11 @@ quantize/dequantize ops. TorchAO int4 group-wise weights of `linear` and stored packed with their scheme in tensor metadata. Each method must lower wholly to one delegate; `to_native` raises otherwise. +Immutable constant keys are namespaced as `/`, even for a single +method. A key mutated by any method becomes `state/` in every method, so +methods share its runtime state. Method names may not be `state` or start with +`state/`. `to_native` raises if namespaced keys collide. + A `.ptn` is an uncompressed zip with fixed entry names: | Entry | Contents | @@ -110,16 +115,21 @@ target_link_libraries(my_app PRIVATE native_api native_vulkan_engine) With `EXECUTORCH_BUILD_TESTS=ON`, `native_runtime_test` covers the shared runtime and API. `native_vulkan_passes_test` covers the Vulkan graph rewrites without -requiring a GPU. When the vision runner below is enabled, -`native_runner_utils_test` covers its shared runner utilities. Run these with +requiring a GPU. When either runner below is enabled, +`native_runner_utils_test` covers their shared runner utilities. Run these with `ctest --test-dir cmake-out -R native_`. The `Test Native Engines` CI workflow has separate Linux jobs for the shared runtime, GGML, and Vulkan, with each test suite run once. To reproduce with dependencies installed, run `bash .ci/scripts/test-native.sh runtime`, `bash .ci/scripts/test-native.sh ggml`, or -`bash .ci/scripts/test-native.sh vulkan` from the OSS checkout. The Vulkan job -also builds `run_classification`; its tests do not require a GPU. +`bash .ci/scripts/test-native.sh vulkan` from the OSS checkout. The GGML job +also runs `native_llm_runner_utils_test` and executes `run_hf_llm_ggml` on a tiny +seeded Hugging Face Llama with a generated byte tokenizer. Its eight generated +token IDs must match eager PyTorch, covering prefill and KV-cache decoding without +downloading a model. The GGML job requires the Python export dependencies +installed by the `Test Native Engines` workflow. The Vulkan job builds +`run_classification` and `run_hf_llm_vulkan`; its tests do not require a GPU. ### GGML engine @@ -177,3 +187,31 @@ Export the `.ptn`, input, reference logits, and labels with [`export.py`](../../examples/native/vision/export.py). The runner prints the top-5 logits and, when `--expected` is given, exits nonzero if the logits differ from the reference beyond `--rtol`/`--atol` (both default to `1e-3`). + +## Native LLM runner + +`EXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON` builds `run_hf_llm_vulkan` with Native +Vulkan and `run_hf_llm_ggml` with Native GGML. It requires gflags and +`tokenizers::tokenizers` (the vendored tokenizer is built automatically in an +OSS checkout; otherwise provide an installed tokenizers package with regex +lookahead support). It accepts Llama 3 TikToken tokenizer files. + +To build the text runner for GGML CPU, initialize the llama.cpp submodule as above: + +```sh +cmake -S . -B cmake-out -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_NATIVE=ON -DEXECUTORCH_BUILD_NATIVE_GGML=ON \ + -DEXECUTORCH_BUILD_NATIVE_LLM_RUNNER=ON +cmake --build cmake-out --target run_hf_llm_ggml +cmake --install cmake-out --prefix install +install/bin/run_hf_llm_ggml --device cpu --cpu_threads 4 \ + --package llama.ptn --tokenizer tokenizer.model \ + --prompt "The capital of France is" --max_new_tokens 32 +``` + +The GGML runner accepts `--device cpu|gpu`, `--device_index`, `--cpu_threads`, +and `--max_cached_graphs` (default 1, 0 rebuilds each call). For GPU execution, +enable the desired `GGML_*` backend when configuring ExecuTorch, as above. + +With `EXECUTORCH_BUILD_TESTS=ON`, `native_llm_runner_utils_test` exercises the +LLM generation utilities. diff --git a/backends/native/examples/CMakeLists.txt b/backends/native/examples/CMakeLists.txt index 81b683a7d05..2f6235040eb 100644 --- a/backends/native/examples/CMakeLists.txt +++ b/backends/native/examples/CMakeLists.txt @@ -4,5 +4,19 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +if(NOT TARGET gflags AND NOT TARGET gflags::gflags) + find_package(gflags REQUIRED) +endif() +if(TARGET gflags::gflags) + set(_native_runner_gflags gflags::gflags) +else() + set(_native_runner_gflags gflags) +endif() + add_subdirectory(runner_utils) -add_subdirectory(vision) +if(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) + add_subdirectory(vision) +endif() +if(EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) + add_subdirectory(hf_llm/runtime) +endif() diff --git a/backends/native/examples/hf_llm/runtime/CMakeLists.txt b/backends/native/examples/hf_llm/runtime/CMakeLists.txt new file mode 100644 index 00000000000..5e3db657663 --- /dev/null +++ b/backends/native/examples/hf_llm/runtime/CMakeLists.txt @@ -0,0 +1,62 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +if(NOT TARGET tokenizers::tokenizers) + if(EXISTS "${EXECUTORCH_ROOT}/extension/llm/tokenizers/CMakeLists.txt") + set(SUPPORT_REGEX_LOOKAHEAD ON) + add_subdirectory( + ${EXECUTORCH_ROOT}/extension/llm/tokenizers + ${CMAKE_CURRENT_BINARY_DIR}/tokenizers + ) + else() + find_package(tokenizers CONFIG REQUIRED) + endif() +endif() + +add_library(native_llm_runner_utils STATIC LlmRunnerUtils.cpp) +target_compile_features(native_llm_runner_utils PUBLIC cxx_std_20) +target_include_directories( + native_llm_runner_utils PUBLIC $ +) +add_library( + native_llm_runner STATIC + LlmRunner.cpp + ${EXECUTORCH_ROOT}/examples/models/llama/tokenizer/llama_tiktoken.cpp +) +target_link_libraries( + native_llm_runner + PUBLIC native_engine + PRIVATE ${_native_runner_gflags} native_llm_runner_utils native_package + native_runner_utils native_runtime tokenizers::tokenizers +) +foreach(target native_llm_runner native_llm_runner_utils) + target_compile_options( + ${target} PRIVATE $<$:/EHsc> + $<$>:-fexceptions> + ) +endforeach() + +set(_llm_runners) +if(EXECUTORCH_BUILD_NATIVE_VULKAN) + add_executable(run_hf_llm_vulkan run_hf_llm_vulkan.cpp) + target_link_libraries(run_hf_llm_vulkan PRIVATE native_vulkan_engine) + list(APPEND _llm_runners run_hf_llm_vulkan) +endif() +if(EXECUTORCH_BUILD_NATIVE_GGML) + add_executable(run_hf_llm_ggml run_hf_llm_ggml.cpp) + target_link_libraries(run_hf_llm_ggml PRIVATE native_ggml_engine) + list(APPEND _llm_runners run_hf_llm_ggml) +endif() +foreach(target IN LISTS _llm_runners) + target_link_libraries( + ${target} PRIVATE native_llm_runner ${_native_runner_gflags} + ) + target_compile_options( + ${target} PRIVATE $<$:/EHsc> + $<$>:-fexceptions> + ) + install(TARGETS ${target} RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) +endforeach() diff --git a/backends/native/examples/vision/CMakeLists.txt b/backends/native/examples/vision/CMakeLists.txt index af7ea765224..76bfdb282f3 100644 --- a/backends/native/examples/vision/CMakeLists.txt +++ b/backends/native/examples/vision/CMakeLists.txt @@ -4,19 +4,10 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -if(NOT TARGET gflags AND NOT TARGET gflags::gflags) - find_package(gflags REQUIRED) -endif() -if(TARGET gflags::gflags) - set(_native_vision_gflags gflags::gflags) -else() - set(_native_vision_gflags gflags) -endif() - add_executable(run_classification run_classification.cpp) target_link_libraries( run_classification - PRIVATE ${_native_vision_gflags} native_engine native_package + PRIVATE ${_native_runner_gflags} native_engine native_package native_runner_utils native_runtime native_vulkan_engine ) target_compile_options( diff --git a/backends/native/serialization/graph_serialize.py b/backends/native/serialization/graph_serialize.py index 827d19e6370..c841c3560e4 100644 --- a/backends/native/serialization/graph_serialize.py +++ b/backends/native/serialization/graph_serialize.py @@ -31,6 +31,7 @@ from enum import IntEnum from typing import ( Any, + Callable, cast, get_args, get_origin, @@ -1487,8 +1488,8 @@ def _iter_graph_args(value: object) -> Iterator[GraphArg]: yield from _iter_graph_args(item) -def collect_data_keys(program: Program) -> set[str]: - """Return every out-of-line data key the program references. +def map_data_keys(method: Method, fn: Callable[[str], str]) -> None: + """Replace every out-of-line data key in ``method`` with ``fn(key)``, in place. Covers constant references (NamedTensorRef.data_key) plus the quant keys (AffineGroup scale and zero-point) carried on any TensorMeta.quant, whether @@ -1496,30 +1497,39 @@ def collect_data_keys(program: Program) -> set[str]: recursing into HOP subgraphs wherever a schema field carries a ``GraphArg``. PackedQuant carries no external keys. """ - keys: set[str] = set() - def add_meta(meta: Optional[TensorMeta]) -> None: + def map_meta(meta: Optional[TensorMeta]) -> None: if meta is None or meta.quant is None: return scheme = meta.quant.scheme if isinstance(scheme, AffineGroup): - keys.add(scheme.scale_data_key) + scheme.scale_data_key = fn(scheme.scale_data_key) if scheme.zero_point_data_key: - keys.add(scheme.zero_point_data_key) + scheme.zero_point_data_key = fn(scheme.zero_point_data_key) def visit_graph(graph: Graph) -> None: for value in graph.tensor_values or []: - add_meta(value.meta) + map_meta(value.meta) for node in graph.nodes: for graph_arg in _iter_graph_args(node): visit_graph(graph_arg.graph) - for method in program.methods: - for constant in method.constants or []: - keys.add(constant.data_key) - add_meta(constant.meta) - visit_graph(method.graph) + for constant in method.constants or []: + constant.data_key = fn(constant.data_key) + map_meta(constant.meta) + visit_graph(method.graph) + +def collect_data_keys(program: Program) -> set[str]: + """Return every out-of-line data key the program references.""" + keys: set[str] = set() + + def add(key: str) -> str: + keys.add(key) + return key + + for method in program.methods: + map_data_keys(method, add) return keys diff --git a/backends/native/test/CMakeLists.txt b/backends/native/test/CMakeLists.txt index 68b9b119f94..fb321e8736e 100644 --- a/backends/native/test/CMakeLists.txt +++ b/backends/native/test/CMakeLists.txt @@ -79,3 +79,22 @@ if(TARGET native_runner_utils) $<$>:-fexceptions> ) endif() + +if(TARGET native_llm_runner_utils) + add_executable( + native_llm_runner_utils_test + ../examples/hf_llm/runtime/LlmRunnerUtilsTest.cpp + ) + target_link_libraries( + native_llm_runner_utils_test PRIVATE native_llm_runner_utils + GTest::gtest_main + ) + add_test(NAME native_llm_runner_utils_test + COMMAND native_llm_runner_utils_test + ) + target_compile_options( + native_llm_runner_utils_test + PRIVATE $<$:/EHsc> + $<$>:-fexceptions> + ) +endif() diff --git a/exir/native/native.py b/exir/native/native.py index 108bf5e180a..4df444e5340 100644 --- a/exir/native/native.py +++ b/exir/native/native.py @@ -38,13 +38,15 @@ from executorch.backends.native.preprocess import NativeDelegateInfo from executorch.backends.native.serialization.graph_serialize import ( + _compile_to_bytes, _same_tensor, collect_data_keys, deserialize_program, + map_data_keys, merge_programs, validate_program, ) -from executorch.backends.native.serialization.schema import Program +from executorch.backends.native.serialization.schema import OutputKind, Program from executorch.exir._serialize._ptn import write_ptn from torch.utils import _pytree as pytree @@ -52,6 +54,8 @@ from torch.export import ExportedProgram from torch.fx import GraphModule +_STATE_NAMESPACE = "state" + class NativeProgramManager: """One or more methods lowered for the native runtime, ready to save. @@ -255,13 +259,7 @@ def _delegate_info(lowered: object, method_name: str) -> NativeDelegateInfo: def _merge_constants( per_method: dict[str, dict[str, torch.Tensor]], ) -> dict[str, torch.Tensor]: - """Union per-method constants by fqn. - - serialize_program treats an fqn as one buffer across every method, so hold to - that here rather than letting a later method silently clobber an earlier one. - _same_tensor short-circuits on identity and shared storage, which is the usual - case for a weight referenced by several methods. - """ + """Merge constants by their final data keys, rejecting differing values.""" merged: dict[str, torch.Tensor] = {} source: dict[str, str] = {} for method_name, constants in per_method.items(): @@ -347,6 +345,11 @@ def _normalize_programs( ) if not method_name: raise ValueError("to_native: method names must be non-empty strings.") + if method_name.split("/")[0] == _STATE_NAMESPACE: + raise ValueError( + f"to_native: method name {method_name!r} is reserved; " + f"{_STATE_NAMESPACE!r} namespaces mutable state keys." + ) if not isinstance(program, TorchExportedProgram): raise TypeError( f"to_native: method {method_name!r} must map to an " @@ -365,6 +368,11 @@ def to_native( ``programs`` is a single ``ExportedProgram`` (lowered as the sole ``forward`` method) or a dict mapping method name to ``ExportedProgram``. + Immutable constant keys are namespaced as ``/``, including for a + single method. Keys mutated by any method become ``state/`` in every + method, so state remains shared across methods. Equal immutable tensors + deduplicate on save. + Nothing about lowering is configurable yet. The passes, the partitioner, and the edge compile config are all owned here, because correctness of the package depends on them. Exposing the compile config is worth reconsidering. @@ -379,9 +387,10 @@ def to_native( Raises: TypeError: If ``programs`` is not an ``ExportedProgram`` or a method-name dictionary containing only ``ExportedProgram`` values. - ValueError: If the method dictionary is empty, a method name is empty, a - method does not fully delegate to the native backend, a constant - differs between methods, or a data key the graph references has no + ValueError: If the method dictionary is empty, a method name is empty or + in the reserved ``state`` namespace, a method does not fully delegate + to the native backend, a mutable constant differs between methods, + namespaced keys conflict, or a data key the graph references has no backing tensor. """ method_programs = _normalize_programs(programs) @@ -447,13 +456,44 @@ def to_native( method_blobs[method_name] = module.processed_bytes method_constants[method_name] = _delegate_info(module, method_name).constants - ptg_blob = merge_programs(method_blobs) - constants = _merge_constants(method_constants) + program = deserialize_program(merge_programs(method_blobs)) + state_keys = _mutable_data_keys(program) | { + buffer.fqn + for method in program.methods + for buffer in method.mutable_buffers or [] + } + state = {key: f"{_STATE_NAMESPACE}/{key}" for key in state_keys} + data_keys = set(state.values()) + for method in program.methods: + constants_for_method = method_constants[method.name] + namespaced = { + key: f"{method.name}/{key}" + for key in constants_for_method + if key not in state + } + conflicts = set(namespaced.values()) & data_keys + if conflicts: + raise ValueError( + "to_native: namespaced constant keys already exist: " + f"{sorted(conflicts)}" + ) + data_keys.update(namespaced.values()) + mapping = state | namespaced + map_data_keys(method, lambda key, mapping=mapping: mapping.get(key, key)) + for buffer in method.mutable_buffers or []: + buffer.fqn = state[buffer.fqn] + for spec in method.output_specs or []: + if spec.kind == OutputKind.BUFFER_MUTATION and spec.target is not None: + spec.target = state.get(spec.target, spec.target) + method_constants[method.name] = { + mapping.get(key, key): tensor + for key, tensor in constants_for_method.items() + } - program = deserialize_program(ptg_blob) + constants = _merge_constants(method_constants) _validate_merged_program(program, constants) - return NativeProgramManager(ptg=ptg_blob, constants=constants) + return NativeProgramManager(ptg=_compile_to_bytes(program), constants=constants) __all__ = ["NativeProgramManager", "to_native"] diff --git a/exir/native/test/test_native.py b/exir/native/test/test_native.py index f9c10812780..c2ee8591b7d 100644 --- a/exir/native/test/test_native.py +++ b/exir/native/test/test_native.py @@ -24,10 +24,16 @@ Graph, Method, MutableBufferSpec, + OutputKind, Program, ) from executorch.exir import to_edge_transform_and_lower -from executorch.exir._serialize._ptn import ALIASES_ENTRY, PTG_ENTRY, SAFETENSORS_ENTRY +from executorch.exir._serialize._ptn import ( + ALIASES_ENTRY, + PTG_ENTRY, + read_ptn, + SAFETENSORS_ENTRY, +) from executorch.exir.lowered_backend_module import ( executorch_call_delegate, get_lowered_submodules, @@ -60,7 +66,7 @@ def _test_equivalence(self, programs): flow keeps one single-method delegate blob per method and would ship the same constants in a .ptd. The graphs are compared modulo that wrapping, and the full key sets (safetensors owners + aliases) must match the - delegates' named-data keys. + delegates' named-data keys after namespacing immutable constants. """ method_programs = ( programs if isinstance(programs, dict) else {"forward": programs} @@ -87,6 +93,12 @@ def _test_equivalence(self, programs): ) native_program = deserialize_program(ptg_blob) + mutated_keys = { + c.data_key + for method in native_program.methods + for c in method.constants or [] + if c.mutated + } native_graphs = {m.name: m.graph for m in native_program.methods} native_keys = st_keys | set(aliases.keys()) @@ -113,9 +125,17 @@ def _test_equivalence(self, programs): # Constants may be embedded in the .pte (pte_data) or shipped to a # .ptd (external_data) depending on external_constants_tag; collect # both so the key set matches to_native's .ptn/.safetensors. - et_keys |= set(nds.pte_data.keys()) + method_keys = set(nds.pte_data.keys()) for key_to_entry in nds.external_data.values(): - et_keys |= set(key_to_entry.keys()) + method_keys |= set(key_to_entry.keys()) + et_keys |= { + ( + f"state/{key}" + if f"state/{key}" in mutated_keys + else f"{name}/{key}" + ) + for key in method_keys + } # Emission is outside the comparison, but keep proving the ExecuTorch path # still completes end to end through the unchanged NamedDataStore branch. @@ -172,10 +192,17 @@ def test_multi_method(self): "prefill": torch.export.export(model, (torch.randn(2, 8),)), "decode": torch.export.export(model, (torch.randn(1, 8),)), } - # Shared weights resolve to the same fqns across methods, so the key sets - # still line up between the two paths. - _, st_keys, _ = self._test_equivalence(programs) - self.assertEqual(len(st_keys), 2) # one weight + bias, shared by fqn + native_keys, st_keys, aliases = self._test_equivalence(programs) + self.assertEqual( + native_keys, + {"prefill/weight", "prefill/bias", "decode/weight", "decode/bias"}, + ) + self.assertEqual(len(st_keys), 2) + for key in ("weight", "bias"): + self.assertEqual( + aliases.get(f"prefill/{key}", f"prefill/{key}"), + aliases.get(f"decode/{key}", f"decode/{key}"), + ) def test_conv_relu(self): model = nn.Sequential( @@ -251,7 +278,7 @@ def forward(self, x): ep = torch.export.export(ConstAdd(), (torch.randn(2, 4),)) _, st_keys, _ = self._test_equivalence(ep) - self.assertTrue(any("lifted" in k for k in st_keys)) + self.assertEqual(st_keys, {"forward/_lifted_tensor_constant0"}) class ToNativeInputValidationTest(unittest.TestCase): @@ -271,6 +298,12 @@ def test_empty_method_name_is_rejected(self): with self.assertRaisesRegex(ValueError, "non-empty"): to_native({"": object()}) + def test_state_namespace_method_name_is_rejected(self): + for name in ("state", "state/a"): + with self.subTest(name=name): + with self.assertRaisesRegex(ValueError, "reserved"): + to_native({name: object()}) + def test_wrong_method_value_is_rejected(self): with self.assertRaisesRegex(TypeError, "must map to an ExportedProgram"): to_native({"forward": object()}) @@ -405,6 +438,182 @@ def test_dropped_public_input_is_rejected(self): self._check(self._graph_module(drop_input=True)) +class ConstantNamespaceTest(unittest.TestCase): + def _scalar_shift(self, offset: float): + class Shift(nn.Module): + def forward(self, x): + return x + offset + + return torch.export.export(Shift(), (torch.randn(2, 4),)) + + def _read(self, native) -> tuple[Program, dict[str, torch.Tensor]]: + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "m.ptn") + native.save(path) + ptg, constants = read_ptn(path) + return deserialize_program(ptg), constants + + def _lifted_keys(self, program: Program, method_name: str) -> list[str]: + (method,) = [m for m in program.methods if m.name == method_name] + return [c.data_key for c in method.constants or [] if "lifted" in c.data_key] + + def test_distinct_lifted_constants_are_namespaced(self): + programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(2.0)} + native = to_native(programs) + self.assertEqual(native.methods, {"a", "b"}) + program, constants = self._read(native) + (a_key,) = self._lifted_keys(program, "a") + (b_key,) = self._lifted_keys(program, "b") + self.assertEqual(a_key, "a/_lifted_tensor_constant0") + self.assertEqual(b_key, "b/_lifted_tensor_constant0") + torch.testing.assert_close( + constants[a_key], torch.full_like(constants[a_key], 1.0), rtol=0, atol=0 + ) + torch.testing.assert_close( + constants[b_key], torch.full_like(constants[b_key], 2.0), rtol=0, atol=0 + ) + + def test_identical_lifted_constants_dedup_on_save(self): + programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(1.0)} + native = to_native(programs) + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "m.ptn") + native.save(path) + ptg, constants = read_ptn(path) + with zipfile.ZipFile(path) as pkg: + owners = _safetensors_keys(pkg.read(SAFETENSORS_ENTRY)) + aliases = json.loads(pkg.read(ALIASES_ENTRY)) + + program = deserialize_program(ptg) + (a_key,) = self._lifted_keys(program, "a") + (b_key,) = self._lifted_keys(program, "b") + self.assertEqual(a_key, "a/_lifted_tensor_constant0") + self.assertEqual(b_key, "b/_lifted_tensor_constant0") + self.assertEqual(owners, {a_key}) + self.assertEqual(aliases, {b_key: a_key}) + self.assertIs(constants[a_key], constants[b_key]) + + def test_distinct_parameters_are_namespaced(self): + programs = { + "a": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)), + "b": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)), + } + program, constants = self._read(to_native(programs)) + self.assertEqual(set(constants), {"a/weight", "a/bias", "b/weight", "b/bias"}) + for method in program.methods: + self.assertEqual( + {c.data_key for c in method.constants}, + {f"{method.name}/weight", f"{method.name}/bias"}, + ) + for key, tensor in programs[method.name].state_dict.items(): + torch.testing.assert_close(constants[f"{method.name}/{key}"], tensor) + + def test_existing_immutable_key_is_also_namespaced(self): + class ConstWithBuffer(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.ones(4)) + self.register_buffer("forward/weight", torch.zeros(4)) + + def forward(self, x): + return x + self.weight + self.get_buffer("forward/weight") + + ep = torch.export.export(ConstWithBuffer(), (torch.randn(4),)) + program, constants = self._read(to_native(ep)) + self.assertEqual(set(constants), {"forward/weight", "forward/forward/weight"}) + self.assertEqual( + {c.data_key for c in program.methods[0].constants}, set(constants) + ) + torch.testing.assert_close(constants["forward/weight"], torch.ones(4)) + torch.testing.assert_close(constants["forward/forward/weight"], torch.zeros(4)) + + def test_mutable_key_is_namespaced_as_state(self): + class ConstWithMutableBuffer(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.ones(4)) + self.register_buffer("forward/weight", torch.zeros(4)) + + def forward(self, x): + buffer = self.get_buffer("forward/weight") + buffer.add_(x) + return buffer + self.weight + + ep = torch.export.export(ConstWithMutableBuffer(), (torch.randn(4),)) + _, constants = self._read(to_native(ep)) + self.assertEqual(set(constants), {"forward/weight", "state/forward/weight"}) + torch.testing.assert_close(constants["forward/weight"], torch.ones(4)) + torch.testing.assert_close(constants["state/forward/weight"], torch.zeros(4)) + + def test_zero_initialized_mutable_buffer_is_namespaced_as_state(self): + program, constants = self._read( + to_native( + self._buffer_program(torch.zeros(4), mutate=True, persistent=False) + ) + ) + self.assertEqual(constants, {}) + (method,) = program.methods + self.assertEqual([b.fqn for b in method.mutable_buffers], ["state/shared"]) + self.assertEqual(self._mutation_targets(method), ["state/shared"]) + + def _mutation_targets(self, method: Method) -> list[str]: + return [ + spec.target + for spec in method.output_specs or [] + if spec.kind == OutputKind.BUFFER_MUTATION + ] + + def _buffer_program(self, buffer, mutate, key="shared", persistent=True): + class Buffer(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer(key, buffer, persistent=persistent) + + def forward(self, x): + value = self.get_buffer(key) + if mutate: + value.add_(x) + return x + value + + return torch.export.export(Buffer(), (torch.ones(4),)) + + def test_namespaced_keys_from_different_methods_must_not_collide(self): + programs = { + "a": self._buffer_program(torch.ones(4), mutate=False, key="b/shared"), + "a/b": self._buffer_program(torch.ones(4), mutate=False), + } + with self.assertRaisesRegex(ValueError, "namespaced constant keys"): + to_native(programs) + + def test_mutable_buffer_stays_shared_with_read_only_method(self): + shared = torch.zeros(4) + programs = { + "read": self._buffer_program(shared, mutate=False), + "update": self._buffer_program(shared, mutate=True), + } + program, constants = self._read(to_native(programs)) + self.assertEqual(set(constants), {"state/shared"}) + torch.testing.assert_close(constants["state/shared"], shared) + for method in program.methods: + (constant,) = method.constants + self.assertEqual(constant.data_key, "state/shared") + self.assertEqual(constant.mutated, method.name == "update") + self.assertEqual( + self._mutation_targets(method), + ["state/shared"] if method.name == "update" else [], + ) + + def test_conflicting_mutable_buffer_values_are_rejected(self): + programs = { + "read": self._buffer_program(torch.ones(4), mutate=False), + "update": self._buffer_program(torch.zeros(4), mutate=True), + } + with self.assertRaisesRegex( + ValueError, "constant 'state/shared' differs between methods" + ): + to_native(programs) + + class MutableBufferDedupTest(unittest.TestCase): def test_batch_norm_statistics_mutated_by_another_method(self): class Classifier(nn.Module): diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index 039c5ce3712..acd332f0bda 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -195,9 +195,14 @@ cmake_dependent_option( "Build the native vision classification runner (run_classification)" OFF "EXECUTORCH_BUILD_NATIVE_VULKAN" OFF ) +cmake_dependent_option( + EXECUTORCH_BUILD_NATIVE_LLM_RUNNER "Build the native LLM runners" OFF + "EXECUTORCH_BUILD_NATIVE_VULKAN OR EXECUTORCH_BUILD_NATIVE_GGML" OFF +) announce_configured_options(EXECUTORCH_BUILD_NATIVE_VULKAN) announce_configured_options(EXECUTORCH_BUILD_NATIVE_GGML) announce_configured_options(EXECUTORCH_BUILD_NATIVE_VISION_RUNNER) +announce_configured_options(EXECUTORCH_BUILD_NATIVE_LLM_RUNNER) define_overridable_option( EXECUTORCH_BUILD_WEBGPU "Build the WebGPU backend" BOOL OFF )