From 1ac7b5af00a93d1e1569379ecdfe50d1f86b88dd Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Wed, 28 Oct 2026 10:25:22 -0700 Subject: [PATCH] Add Whisper HF export (encoder, cross_kv, static-cache decoder) (#23901) Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/23901 Add `examples/audio/export_hf_whisper.py` to export `WhisperForConditionalGeneration` directly through Transformers into one `.ptn` with three methods: `encoder` maps mel features to hidden states, `cross_kv` projects per-layer cross-attention keys then values, and `decoder` consumes token IDs, explicit cache positions, and those projections to produce logits. Self-attention uses zero-initialized mutable static-cache buffers. Positions and the causal mask derive from `cache_position`, so the runner needs no attention-mask input. `--dump_dir` writes runner inputs and an eager greedy transcript with a fresh decoder, preserving the exported cache buffers. An unspecified forced language token is resolved from the example features through Transformers language detection, retaining subsequent forced-token positions. Sparse forced prefixes are rejected. Generation is capped at `max_target_positions` and skips the unused decoder call after the last requested token. The export format is independent of the runtime engine. The exporter and the existing LLM example share lowering reports through `examples/hf_llm/hf_utils.py`. Tests validate native delegation, zero-initialized decoder caches, three-method packaging, and eager numerics. Runtime execution support and runner integration follow in later diffs. Random-weight loading uses checkpoint generation settings when present, warns when the file is absent, and propagates malformed-config errors. Vector generation supports an unset EOS token by using an empty stop set. Authored with assistance from Codex. Reviewed By: digantdesai Differential Revision: D125295057 --- backends/native/examples/audio/BUCK | 28 ++ .../examples/audio/export_hf_whisper.py | 427 ++++++++++++++++++ backends/native/examples/hf_llm/BUCK | 13 + .../native/examples/hf_llm/export_hf_llm.py | 44 +- backends/native/examples/hf_llm/hf_utils.py | 61 +++ backends/native/test/BUCK | 14 + backends/native/test/test_hf_whisper.py | 241 ++++++++++ 7 files changed, 786 insertions(+), 42 deletions(-) create mode 100644 backends/native/examples/audio/BUCK create mode 100644 backends/native/examples/audio/export_hf_whisper.py create mode 100644 backends/native/examples/hf_llm/hf_utils.py create mode 100644 backends/native/test/test_hf_whisper.py diff --git a/backends/native/examples/audio/BUCK b/backends/native/examples/audio/BUCK new file mode 100644 index 00000000000..f6e7ab8a11f --- /dev/null +++ b/backends/native/examples/audio/BUCK @@ -0,0 +1,28 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target( + _kind = runtime.python_library, + name = "hf_whisper", + srcs = ["export_hf_whisper.py"], + visibility = ["//executorch/backends/native/test/..."], + deps = [ + "//caffe2:torch", + "//executorch/backends/native/examples/hf_llm:hf_utils", + "//executorch/exir:lib", + "//executorch/exir/native:native", + "fbsource//third-party/pypi/transformers:transformers", + ], +) + +fbcode_target( + _kind = runtime.python_binary, + name = "export_hf_whisper", + main_function = "executorch.backends.native.examples.audio.export_hf_whisper.main", + deps = [ + ":hf_whisper", + "fbsource//third-party/pypi/safetensors:safetensors", + ], +) diff --git a/backends/native/examples/audio/export_hf_whisper.py b/backends/native/examples/audio/export_hf_whisper.py new file mode 100644 index 00000000000..6f0dabf7892 --- /dev/null +++ b/backends/native/examples/audio/export_hf_whisper.py @@ -0,0 +1,427 @@ +# 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. + +"""Export a Hugging Face Whisper model through the native backend. + +Builds ``WhisperForConditionalGeneration`` from a Hugging Face model +directory (or only its config, with random weights) and exports three methods +into one ``.ptn`` for the native runtime: + +- ``encoder``: ``input_features [1, num_mel_bins, 2 * max_source_positions]`` + (30 s of log-mel at the usual 100 fps) -> ``encoder_hidden_states + [1, max_source_positions, d_model]``. +- ``cross_kv``: ``encoder_hidden_states`` -> per-layer cross-attention + keys then values, ``2 * decoder_layers`` tensors of + ``[1, decoder_attention_heads, max_source_positions, head_dim]``, + projected and reshaped exactly as ``WhisperAttention`` would on first use. +- ``decoder``: ``(decoder_input_ids [1, T], cache_position [T], + *cross_keys, *cross_values)`` -> ``logits [1, T, vocab_size]`` with a + static self-attention KV cache of ``max_target_positions`` rows held in + mutable buffers, mirroring the decoder-only wrapper in ``export_hf_llm``. + The runner projects cross K/V once, prefills the forced prefix, then + decodes one token per call; positions and the causal mask derive from + ``cache_position`` in-graph, so there is no attention-mask input. Cross K/V + reach the decoder as ordinary inputs reused throughout decoding. + +Reports which operators the native partitioner leaves behind, like +``export_hf_llm``. ``--dump_dir`` additionally writes the example mel input +and an eager greedy transcript (``input_features.bin``, ``prompt_ids.txt``, +``expected_ids.txt``) for the ``examples/audio`` runner. +""" + +from __future__ import annotations + +import argparse +import logging +import os + +import torch +import torch.nn as nn +from executorch.backends.native.examples.hf_llm.hf_utils import report_lowering +from executorch.exir.native import to_native +from torch.export import Dim, export, ExportedProgram +from transformers import ( + AutoConfig, + GenerationConfig, + WhisperConfig, + WhisperForConditionalGeneration, +) +from transformers.cache_utils import DynamicCache, EncoderDecoderCache, StaticCache +from transformers.integrations.executorch import register_dynamic_cache_export_support +from transformers.utils import GENERATION_CONFIG_NAME +from transformers.utils.hub import cached_file + +logger: logging.Logger = logging.getLogger(__name__) + + +def _load( + path: str, + encoder_layers: int, + decoder_layers: int, + random_weights: bool, +) -> WhisperForConditionalGeneration: + config = AutoConfig.from_pretrained(path) + if not isinstance(config, WhisperConfig): + raise ValueError(f"{path!r} is a {type(config).__name__}, not Whisper.") + if encoder_layers: + config.encoder_layers = encoder_layers + if decoder_layers: + config.decoder_layers = decoder_layers + kwargs = {"torch_dtype": torch.float32, "attn_implementation": "sdpa"} + if random_weights: + model = WhisperForConditionalGeneration(config) + if hasattr(model, "set_attn_implementation"): + model.set_attn_implementation("sdpa") + # The constructor synthesizes a generation config from the model + # config; prefer the checkpoint's (it carries forced decoder ids). + generation_config_file = cached_file( + path, + GENERATION_CONFIG_NAME, + _raise_exceptions_for_missing_entries=False, + ) + if generation_config_file is None: + logger.warning( + "No generation config found in %s; using the model config's defaults", + path, + ) + else: + model.generation_config = GenerationConfig.from_pretrained(path) + else: + model = WhisperForConditionalGeneration.from_pretrained( + path, config=config, **kwargs + ) + return model.eval() + + +class WhisperEncoder(nn.Module): + """Log-mel ``input_features`` -> encoder hidden states.""" + + def __init__(self, encoder: nn.Module) -> None: + super().__init__() + self.encoder = encoder + + def forward(self, input_features: torch.Tensor) -> torch.Tensor: + return self.encoder(input_features=input_features).last_hidden_state + + +class WhisperCrossKV(nn.Module): + """Project encoder hidden states to per-layer cross-attention K/V. + + Returns keys for every decoder layer followed by values, each + ``[1, heads, source_positions, head_dim]`` — the tensors + ``WhisperAttention`` computes from ``encoder_hidden_states`` on first + use (``k_proj``/``v_proj``, head-major, contiguous; cross K/V are not + scaled, only queries are). + """ + + def __init__(self, model: WhisperForConditionalGeneration) -> None: + super().__init__() + self.layers = model.get_decoder().layers + config = model.config + self.num_heads = config.decoder_attention_heads + self.head_dim = config.d_model // config.decoder_attention_heads + self.source_positions = config.max_source_positions + + def forward(self, encoder_hidden_states: torch.Tensor) -> tuple[torch.Tensor, ...]: + # Compute in return order (all keys, then all values): the native + # partitioner emits delegate outputs in computation order, and + # to_native requires the method to return them in that order. + shape = (1, self.source_positions, self.num_heads, self.head_dim) + keys = [ + layer.encoder_attn.k_proj(encoder_hidden_states) + .view(*shape) + .transpose(1, 2) + .contiguous() + for layer in self.layers + ] + values = [ + layer.encoder_attn.v_proj(encoder_hidden_states) + .view(*shape) + .transpose(1, 2) + .contiguous() + for layer in self.layers + ] + return (*keys, *values) + + +class WhisperDecoder(nn.Module): + """Static-cache decoder pass over precomputed cross K/V -> logits. + + Self-attention K/V live in a ``StaticCache`` registered as zero-initialized + mutable buffers. The cross-attention cache is marked updated and its + layers are bound to the ``cross_kv`` inputs (keys for every layer, then + values) on each call. ``WhisperAttention`` reads them directly without + projecting encoder states. The cache is stateful: one instance decodes + one utterance. + """ + + def __init__(self, model: WhisperForConditionalGeneration) -> None: + super().__init__() + self.decoder = model.get_decoder() + self.proj_out = model.proj_out + config = model.config + head_dim = config.d_model // config.decoder_attention_heads + self.self_cache = StaticCache( + config=config, max_cache_len=config.max_target_positions + ) + self.self_cache.early_initialization( + batch_size=1, + num_heads=config.decoder_attention_heads, + head_dim=head_dim, + dtype=torch.float32, + device=torch.device("cpu"), + ) + self.cross_cache = DynamicCache(config=config) + self.cache = EncoderDecoderCache(self.self_cache, self.cross_cache) + for layer_idx in range(config.decoder_layers): + self.cache.is_updated[layer_idx] = True + register_dynamic_cache_export_support() + for i in range(len(self.self_cache)): + self.register_buffer( + f"key_cache_{i}", self.self_cache.layers[i].keys, persistent=False + ) + self.register_buffer( + f"value_cache_{i}", + self.self_cache.layers[i].values, + persistent=False, + ) + + def forward( + self, + decoder_input_ids: torch.Tensor, + cache_position: torch.Tensor, + *cross_kv: torch.Tensor, + ) -> torch.Tensor: + num_layers = len(self.cross_cache.layers) + for layer_idx, layer in enumerate(self.cross_cache.layers): + layer.keys = cross_kv[layer_idx] + layer.values = cross_kv[num_layers + layer_idx] + layer.is_initialized = True + hidden = self.decoder( + input_ids=decoder_input_ids, + # Gate only: WhisperDecoderLayer runs its cross-attention block + # iff this is not None, but with the cross cache marked updated + # WhisperAttention reads K/V from the cache and never touches + # these values. Any input tensor serves; reuse cross keys 0. + encoder_hidden_states=cross_kv[0], + past_key_values=self.cache, + use_cache=True, + cache_position=cache_position, + ).last_hidden_state + return self.proj_out(hidden) + + +def prompt_ids( + model: WhisperForConditionalGeneration, + input_features: torch.Tensor | None = None, +) -> list[int]: + """Forced decoder prefix, detecting an unspecified language from the input.""" + ids = [model.config.decoder_start_token_id] + forced = getattr(model.generation_config, "forced_decoder_ids", None) or [] + for position, token in sorted(forced): + if position != len(ids): + raise ValueError("forced_decoder_ids must form a contiguous prefix") + if token is None: + if position != 1: + raise ValueError("only the language token at position 1 may be None") + if input_features is None: + raise ValueError("input_features required for language detection") + token = int( + model.detect_language( + input_features=input_features, + num_segment_frames=2 * model.config.max_source_positions, + ).item() + ) + ids.append(token) + return ids + + +def example_input_features(config: WhisperConfig) -> torch.Tensor: + """Seeded stand-in mel input; the feature extractor pads audio to this shape.""" + generator = torch.Generator().manual_seed(0) + return torch.randn( + (1, config.num_mel_bins, 2 * config.max_source_positions), + generator=generator, + ) + + +def build_exported_programs( + model: WhisperForConditionalGeneration, + example_features: torch.Tensor | None = None, +) -> tuple[WhisperEncoder, WhisperCrossKV, WhisperDecoder, dict[str, ExportedProgram]]: + """Export the encoder/cross-KV/decoder wrappers around ``model``. + + The encoder and cross-KV shapes are static (Whisper always consumes a + full 30 s mel window); only the decoder token axis is dynamic. + """ + config = model.config + encoder = WhisperEncoder(model.get_encoder()).eval() + cross_kv = WhisperCrossKV(model).eval() + decoder = WhisperDecoder(model).eval() + if example_features is None: + example_features = example_input_features(config) + batch = example_features.shape[0] + + with torch.no_grad(): + hidden = encoder(example_features) + cross = cross_kv(hidden) + # Trace with a handful of tokens: a 1-token example would specialize + # the dynamic token axis to a constant. + example_ids = prompt_ids(model, example_features) + example_ids += [example_ids[-1]] * max(0, 4 - len(example_ids)) + prefix = torch.tensor([example_ids], dtype=torch.long) + prefix = prefix.expand(batch, -1).contiguous() + cache_position = torch.arange(len(example_ids), dtype=torch.long) + tokens = Dim("tokens", min=1, max=config.max_target_positions) + encoder_ep = export( + encoder, (example_features,), dynamic_shapes=None, strict=True + ) + cross_kv_ep = export(cross_kv, (hidden,), dynamic_shapes=None, strict=True) + decoder_ep = export( + decoder, + (prefix, cache_position, *cross), + dynamic_shapes=( + {1: tokens}, + {0: tokens}, + (None,) * len(cross), + ), + strict=True, + ) + programs = { + "encoder": encoder_ep, + "cross_kv": cross_kv_ep, + "decoder": decoder_ep, + } + return encoder, cross_kv, decoder, programs + + +def _greedy_transcript( + encoder: WhisperEncoder, + cross_kv: WhisperCrossKV, + decoder: WhisperDecoder, + example_features: torch.Tensor, + prefix: list[int], + eos_ids: set[int], + max_new_tokens: int, +) -> list[int]: + ids = list(prefix) + capacity = decoder.decoder.config.max_target_positions + if not ids or len(ids) > capacity: + raise ValueError( + "the prefix must contain between 1 and max_target_positions tokens" + ) + if max_new_tokens < 0: + raise ValueError("max_new_tokens must be nonnegative") + max_new_tokens = min(max_new_tokens, capacity - len(ids)) + if max_new_tokens == 0: + return ids + with torch.no_grad(): + cross = cross_kv(encoder(example_features)) + window = torch.tensor([ids], dtype=torch.long) + positions = torch.arange(len(ids), dtype=torch.long) + logits = decoder(window, positions, *cross) + for step in range(max_new_tokens): + next_id = int(logits[0, -1].argmax()) + ids.append(next_id) + if next_id in eos_ids or step + 1 == max_new_tokens: + break + logits = decoder( + torch.tensor([[next_id]], dtype=torch.long), + torch.tensor([len(ids) - 1], dtype=torch.long), + *cross, + ) + return ids + + +def _dump_test_vectors( + dump_dir: str, + model: WhisperForConditionalGeneration, + encoder: WhisperEncoder, + example_features: torch.Tensor, + max_new_tokens: int, +) -> None: + os.makedirs(dump_dir, exist_ok=True) + # A fresh decoder: the exported one's cache buffers must stay zero for + # the serialized package, and eager decoding would fill them. + decoder = WhisperDecoder(model).eval() + prefix = prompt_ids(model, example_features) + eos = model.config.eos_token_id + eos_ids = {eos} if isinstance(eos, int) else set(eos or []) + transcript = _greedy_transcript( + encoder, + WhisperCrossKV(model).eval(), + decoder, + example_features, + prefix, + eos_ids, + max_new_tokens, + ) + with open(os.path.join(dump_dir, "input_features.bin"), "wb") as f: + f.write(example_features.contiguous().numpy().tobytes()) + for name, ids in (("prompt_ids.txt", prefix), ("expected_ids.txt", transcript)): + with open(os.path.join(dump_dir, name), "w") as f: + f.write("".join(f"{token}\n" for token in ids)) + print(f"prompt ids: {prefix}") + print(f"eos ids: {sorted(eos_ids)}") + print(f"eager greedy ids ({len(transcript)}): {transcript}") + print(f"wrote test vectors to {dump_dir}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", + default="openai/whisper-tiny", + help="Hugging Face model directory or hub ID.", + ) + parser.add_argument( + "--random_weights", + action="store_true", + help="Use only the model's config, with random weights.", + ) + parser.add_argument("--encoder_layers", type=int, default=0) + parser.add_argument("--decoder_layers", type=int, default=0) + parser.add_argument("--output", default="") + parser.add_argument( + "--dump_dir", + default="", + help="Also write runner test vectors (mel input, prompt, eager ids).", + ) + parser.add_argument( + "--max_new_tokens", + type=int, + default=32, + help="Maximum new test-vector tokens, capped by the decoder context.", + ) + args = parser.parse_args() + logging.basicConfig(level=logging.INFO) + torch.manual_seed(0) + + model = _load( + args.model, args.encoder_layers, args.decoder_layers, args.random_weights + ) + example_features = example_input_features(model.config) + encoder, _cross_kv, _decoder, programs = build_exported_programs( + model, example_features + ) + for name, ep in programs.items(): + report_lowering(ep, name) + + if args.dump_dir: + _dump_test_vectors( + args.dump_dir, + model, + encoder, + example_features, + args.max_new_tokens, + ) + if args.output: + to_native(programs).save(args.output) + print(f"saved {args.output}") + + +if __name__ == "__main__": + main() diff --git a/backends/native/examples/hf_llm/BUCK b/backends/native/examples/hf_llm/BUCK index 535e80081a6..4a3f6794f94 100644 --- a/backends/native/examples/hf_llm/BUCK +++ b/backends/native/examples/hf_llm/BUCK @@ -3,12 +3,25 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target( + _kind = runtime.python_library, + name = "hf_utils", + srcs = ["hf_utils.py"], + visibility = ["//executorch/backends/native/examples/audio/..."], + deps = [ + "//caffe2:torch", + "//executorch/backends/native:lib", + "//executorch/exir:lib", + ], +) + fbcode_target( _kind = runtime.python_library, name = "export_hf_llm_lib", srcs = ["export_hf_llm.py"], deps = [ ":gguf_hf", + ":hf_utils", "//caffe2:torch", "//executorch/backends/native:lib", "//executorch/exir:lib", diff --git a/backends/native/examples/hf_llm/export_hf_llm.py b/backends/native/examples/hf_llm/export_hf_llm.py index 8a05ee90ac7..abb5c0925b8 100644 --- a/backends/native/examples/hf_llm/export_hf_llm.py +++ b/backends/native/examples/hf_llm/export_hf_llm.py @@ -17,24 +17,15 @@ """ import argparse -import collections import logging -import operator import torch import torch.nn as nn -from executorch.backends.native import get_default_compile_config from executorch.backends.native.examples.hf_llm.gguf_hf import ( gguf_architecture, gguf_to_hf, ) -from executorch.backends.native.partitioner import NativePartitioner -from executorch.backends.native.passes import get_default_passes -from executorch.exir import to_edge_transform_and_lower -from executorch.exir.lowered_backend_module import ( - executorch_call_delegate, - get_lowered_submodules, -) +from executorch.backends.native.examples.hf_llm.hf_utils import report_lowering from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples from executorch.extension.llm.export.load import assign_one, iter_checkpoint, RawIter from torch.export import Dim, export @@ -270,15 +261,6 @@ def golden_samples( return SampleSet(samples, source="eager/cpu", tol=tol) -def _histogram(gm: torch.fx.GraphModule) -> collections.Counter: - return collections.Counter( - str(n.target) - for n in gm.graph.nodes - if n.op == "call_function" - and n.target not in (executorch_call_delegate, operator.getitem) - ) - - def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument( @@ -360,29 +342,7 @@ def main() -> None: dynamic_shapes={"input_ids": {1: seq}, "cache_position": {0: seq}}, strict=not args.non_strict, ) - print("== exported aten ops ==", flush=True) - for k, v in sorted(_histogram(ep.graph_module).items()): - print(f" {v:4d} {k}") - - edge = to_edge_transform_and_lower( - ep, - transform_passes=get_default_passes(), - partitioner=[ - NativePartitioner(external_constants_tag=None, _serialize_as_ptn=True) - ], - compile_config=get_default_compile_config(), - ) - gm = edge.exported_program().graph_module - lowered = get_lowered_submodules(gm) - print(f"== {len(lowered)} native delegate(s) ==") - delegated = collections.Counter() - for _, module, _ in lowered: - delegated += _histogram(module.original_module.graph_module) - for k, v in sorted(delegated.items()): - print(f" {v:4d} {k}") - print("== residual (not delegated) ==") - for k, v in sorted(_histogram(gm).items()): - print(f" {v:4d} {k}") + report_lowering(ep, "forward") if args.output: to_native(ep).save(args.output) diff --git a/backends/native/examples/hf_llm/hf_utils.py b/backends/native/examples/hf_llm/hf_utils.py new file mode 100644 index 00000000000..c0f8709d9b0 --- /dev/null +++ b/backends/native/examples/hf_llm/hf_utils.py @@ -0,0 +1,61 @@ +# 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. + +"""Report native lowering for exported Hugging Face models.""" + +import collections +import operator + +import torch +from executorch.backends.native import get_default_compile_config +from executorch.backends.native.partitioner import NativePartitioner +from executorch.backends.native.passes import get_default_passes +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import ( + executorch_call_delegate, + get_lowered_submodules, +) +from torch.export import ExportedProgram + + +def histogram(gm: torch.fx.GraphModule) -> collections.Counter: + """Count call_function targets, ignoring delegate plumbing.""" + return collections.Counter( + str(n.target) + for n in gm.graph.nodes + if n.op == "call_function" + and n.target not in (executorch_call_delegate, operator.getitem) + ) + + +def report_lowering(ep: ExportedProgram, label: str) -> None: + """Lower ``ep`` for PTN and print delegated vs. residual op histograms.""" + print(f"== {label}: exported aten ops ==", flush=True) + for k, v in sorted(histogram(ep.graph_module).items()): + print(f" {v:4d} {k}") + + edge = to_edge_transform_and_lower( + ep, + transform_passes=get_default_passes(), + partitioner=[ + NativePartitioner(external_constants_tag=None, _serialize_as_ptn=True) + ], + compile_config=get_default_compile_config(), + ) + gm = edge.exported_program().graph_module + lowered = get_lowered_submodules(gm) + print(f"== {label}: {len(lowered)} native delegate(s) ==") + delegated: collections.Counter = collections.Counter() + for _, module, _ in lowered: + delegated += histogram(module.original_module.graph_module) + for k, v in sorted(delegated.items()): + print(f" {v:4d} {k}") + print(f"== {label}: residual (not delegated) ==") + residual = histogram(gm) + if not residual: + print(" (none)") + for k, v in sorted(residual.items()): + print(f" {v:4d} {k}") diff --git a/backends/native/test/BUCK b/backends/native/test/BUCK index c3f3eeb546d..a1e520843ae 100644 --- a/backends/native/test/BUCK +++ b/backends/native/test/BUCK @@ -87,3 +87,17 @@ fbcode_target( "//executorch/extension/llm/export:gguf", ], ) + +fbcode_target( + _kind = runtime.python_test, + name = "test_hf_whisper", + srcs = ["test_hf_whisper.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/native:lib", + "//executorch/backends/native/examples/audio:hf_whisper", + "//executorch/exir:lib", + "//executorch/exir/native:native", + "fbsource//third-party/pypi/transformers:transformers", + ], +) diff --git a/backends/native/test/test_hf_whisper.py b/backends/native/test/test_hf_whisper.py new file mode 100644 index 00000000000..7e53594c9f7 --- /dev/null +++ b/backends/native/test/test_hf_whisper.py @@ -0,0 +1,241 @@ +# 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. + +"""Whisper export contracts shared by the native runtime engines. + +A small random-weight model exercises native delegation, three-method +packaging, cache initialization, and eager numerics through prefill and +incremental decoding. Runtime operator support is covered by engine tests. +""" + +import operator +import os +import tempfile +import unittest + +import torch +from executorch.backends.native import get_default_compile_config +from executorch.backends.native.examples.audio.export_hf_whisper import ( + _dump_test_vectors, + _greedy_transcript, + _load, + build_exported_programs, + example_input_features, + prompt_ids, + WhisperDecoder, +) +from executorch.backends.native.partitioner import NativePartitioner +from executorch.backends.native.passes import get_default_passes +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import ( + executorch_call_delegate, + get_lowered_submodules, +) +from executorch.exir.native import to_native +from transformers import ( + GenerationConfig, + WhisperConfig, + WhisperForConditionalGeneration, +) + + +class WhisperNativeExportTest(unittest.TestCase): + def setUp(self): + torch.manual_seed(0) + config = WhisperConfig( + vocab_size=256, + num_mel_bins=80, + d_model=64, + encoder_layers=1, + encoder_attention_heads=4, + encoder_ffn_dim=128, + # Two layers also pins cross_kv's multi-output ordering + # (delegate outputs follow computation order). + decoder_layers=2, + decoder_attention_heads=4, + decoder_ffn_dim=128, + max_source_positions=50, + max_target_positions=16, + decoder_start_token_id=2, + pad_token_id=0, + eos_token_id=2, + bos_token_id=1, + ) + self.model = WhisperForConditionalGeneration(config).eval() + self.model.generation_config.forced_decoder_ids = [(1, 10)] + self.features = example_input_features(config) + ( + self.encoder, + self.cross_kv, + self.decoder, + self.programs, + ) = build_exported_programs(self.model, self.features) + + def test_random_weights_loads_checkpoint_generation_config(self): + with tempfile.TemporaryDirectory() as path: + self.model.config.save_pretrained(path) + GenerationConfig(forced_decoder_ids=[(1, 10)]).save_pretrained(path) + + loaded = _load(path, 0, 0, random_weights=True) + + self.assertEqual(prompt_ids(loaded), [2, 10]) + + def test_random_weights_warns_when_generation_config_is_missing(self): + with tempfile.TemporaryDirectory() as path: + self.model.config.save_pretrained(path) + + with self.assertLogs( + "executorch.backends.native.examples.audio.export_hf_whisper", + level="WARNING", + ) as logs: + loaded = _load(path, 0, 0, random_weights=True) + + self.assertIn("No generation config found", logs.output[0]) + self.assertEqual(prompt_ids(loaded), [2]) + + def test_random_weights_rejects_malformed_generation_config(self): + with tempfile.TemporaryDirectory() as path: + self.model.config.save_pretrained(path) + with open(os.path.join(path, "generation_config.json"), "w") as f: + f.write("{invalid json") + + with self.assertRaises(OSError): + _load(path, 0, 0, random_weights=True) + + def test_dump_vectors_without_eos(self): + self.model.config.eos_token_id = None + with tempfile.TemporaryDirectory() as path: + _dump_test_vectors(path, self.model, self.encoder, self.features, 1) + + with open(os.path.join(path, "expected_ids.txt")) as f: + ids = [int(line) for line in f] + with open(os.path.join(path, "input_features.bin"), "rb") as f: + features = f.read() + + self.assertEqual(ids[:2], [2, 10]) + self.assertEqual(len(ids), 3) + self.assertEqual(features, self.features.contiguous().numpy().tobytes()) + + def test_prompt_ids_from_generation_config(self): + self.assertEqual(prompt_ids(self.model), [2, 10]) + + def test_prompt_ids_detects_language_without_shifting_forced_tokens(self): + self.model.generation_config.forced_decoder_ids = [(1, None), (2, 10)] + self.model.generation_config.lang_to_id = {"<|en|>": 11, "<|fr|>": 12} + language = int(self.model.detect_language(input_features=self.features).item()) + + self.assertEqual(prompt_ids(self.model, self.features), [2, language, 10]) + + def test_prompt_ids_requires_features_for_language_detection(self): + self.model.generation_config.forced_decoder_ids = [(1, None), (2, 10)] + + with self.assertRaisesRegex(ValueError, "input_features required"): + prompt_ids(self.model) + + def test_prompt_ids_rejects_gaps(self): + self.model.generation_config.forced_decoder_ids = [(2, 10)] + + with self.assertRaisesRegex(ValueError, "contiguous prefix"): + prompt_ids(self.model, self.features) + + def test_greedy_transcript_respects_token_and_context_limits(self): + prefix = prompt_ids(self.model) + capacity = self.model.config.max_target_positions + for budget in (0, 1, 32): + with self.subTest(max_new_tokens=budget): + decoder = WhisperDecoder(self.model).eval() + transcript = _greedy_transcript( + self.encoder, + self.cross_kv, + decoder, + self.features, + prefix, + set(), + budget, + ) + + self.assertEqual(transcript[: len(prefix)], prefix) + self.assertEqual(len(transcript), min(len(prefix) + budget, capacity)) + unused_start = len(transcript) - 1 if budget else 0 + for layer in decoder.self_cache.layers: + self.assertEqual( + torch.count_nonzero(layer.keys[:, :, unused_start:]).item(), 0 + ) + self.assertEqual( + torch.count_nonzero(layer.values[:, :, unused_start:]).item(), 0 + ) + + def test_all_methods_wholly_delegate(self): + for name, ep in self.programs.items(): + with self.subTest(method=name): + edge = to_edge_transform_and_lower( + ep, + transform_passes=get_default_passes(), + partitioner=[ + NativePartitioner( + external_constants_tag=None, _serialize_as_ptn=True + ) + ], + compile_config=get_default_compile_config(), + ) + gm = edge.exported_program().graph_module + self.assertEqual(len(get_lowered_submodules(gm)), 1) + residual = [ + node.target + for node in gm.graph.nodes + if node.op == "call_function" + and node.target not in (executorch_call_delegate, operator.getitem) + ] + self.assertEqual(residual, []) + + def test_exported_decoder_cache_is_zero_initialized(self): + buffers = list(self.programs["decoder"].buffers()) + self.assertEqual(len(buffers), 2 * self.model.config.decoder_layers) + for buffer in buffers: + self.assertEqual(torch.count_nonzero(buffer).item(), 0) + + def test_exported_programs_match_eager(self): + with torch.no_grad(): + expected_hidden = self.encoder(self.features) + actual_hidden = self.programs["encoder"].module()(self.features) + if isinstance(actual_hidden, (tuple, list)): + actual_hidden = actual_hidden[0] + torch.testing.assert_close(actual_hidden, expected_hidden) + + expected_cross = self.cross_kv(expected_hidden) + actual_cross = self.programs["cross_kv"].module()(expected_hidden) + self.assertEqual(len(actual_cross), len(expected_cross)) + for actual_t, expected_t in zip(actual_cross, expected_cross): + torch.testing.assert_close(actual_t, expected_t) + + # Prefill two tokens, then decode single tokens against the + # live cache; eager and exported modules keep separate caches. + steps = [ + (torch.tensor([[3, 4]]), torch.tensor([0, 1])), + (torch.tensor([[5]]), torch.tensor([2])), + (torch.tensor([[6]]), torch.tensor([3])), + ] + exported_decoder = self.programs["decoder"].module() + for ids, positions in steps: + expected = self.decoder(ids, positions, *expected_cross) + actual = exported_decoder(ids, positions, *actual_cross) + if isinstance(actual, (tuple, list)): + actual = actual[0] + torch.testing.assert_close(actual, expected) + + def test_three_method_package_saves(self): + # Regression: each method lifts its own scalar constants under the + # same auto-generated names; the package must still save. + program = to_native(self.programs) + self.assertEqual(program.methods, {"encoder", "cross_kv", "decoder"}) + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "whisper.ptn") + program.save(path) + self.assertGreater(os.path.getsize(path), 0) + + +if __name__ == "__main__": + unittest.main()