Repository navigation
Qualcomm AI Engine Direct - [GenAI Pipeline Phase 2] PRB1 - Compilation foundation - #22846
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22846
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 6748db8 with merge base e5acd53 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@pytorchbot label "release notes: qualcomm" |
b547d15 to
dcc47e3
Compare
|
@claude Review this pr , and share a high level and technical summary on how this improves underlying things and sets up foundation |
|
Claude finished @psiddh's task in 2m 2s —— View job PR review: GenAI Pipeline Phase 2, PRB1 (compilation foundation)
High-level summaryPhase 1 built the pipeline's structure (stages, strategies, adapters, typed configs), but compilation was a stub: How it improves things:
This leaves Stream B ready for PR-B2 (multi-graph lowering, weight sharing, sharding, spill-fill). The Technical walkthrough
The layering is sound. The adapter compiles exactly one graph. Grouping several graphs into one Review findings1. Once real specs are passed, the pipeline would always fail the target check (the SoC is a string, not the enum). 2. 3. 4. 5. Docstring nits.
6. CI. The one failure, VerdictThe foundation looks solid: clear naming, a reasonable decision on what is frozen, and the two bug fixes are real. I'd fix #1 before PR-B2 connects real compile specs to the adapter, and #2 before any CLI uses I didn't run the test suite: the sandbox didn't approve Python commands, so this review is from reading the code. Findings #1 and #2 are confirmed by tracing the call sites and the |
dcc47e3 to
6748db8
Compare
|
Regarding Claude's comments:
@psiddh Kindly take a look when you get chance, I think this PR should be GTG now. |
…paration and Quantization (#23050) ## Summary First PR of Phase 2 Stream A. This PR makes the GenAI Pipeline runnable for LLM and MLLM models through the front half of the flow: ```text CLI -> model lookup -> source transforms -> model preparation -> dataset-backed calibration -> quantization -> save QDQ model ``` `examples/` and the legacy `llama.py` path remain unchanged. Compilation and device inference remain in Stream B / PR-A2. ## What's Included ### CLI Entry Point: `cli.py` Adds a GenAI Pipeline CLI that resolves model, dataset, and quantization options, builds `PipelineContext`, and invokes model preparation and quantization. FP16 omits quantization; QAT, embedding quantization, and attention sink are rejected for now. The rejected feature will be raised in following PR. ### Registry-Backed Model lookup: `model_lookup.py` Maps a user-facing model name to the model-specific inputs consumed by the pipeline. Callers resolve one registry config, then derive from that same entry: - **Model construction:** component and graph keyed constructors, including decoder decode/quantize graphs, optional prefill graphs, and multimodal embedding variants. - **Source preparation:** component-keyed state-dict weight transforms and module G2G transforms, plus the Hugging Face state-dict loader when needed. - **Pipeline configuration:** component sharding, quantization dtypes and recipe classes, and LLM or MLLM loader/quantizer adapters. This centralizes model-family branching at the registry boundary; the model preparation and quantization stages consume only the derived component and graph keyed configuration. ### LLM and Multimodal LLM model components: `model_components/` Adds pipeline import surfaces for decoder, token embedding, and encoder components over existing `examples` implementations. ### Module G2G Transforms: `source_transform/`: Collects existing `llama.py` transformations into reusable pipeline functions: - Weight transforms. - Model graph-to-graph (G2G) transforms. ### LLM and MLLM Model Loading and Preparation: `strategies/model_preparation/`: `ExecuTorchModelPreparationStrategy` implements the shared preparation flow: 1. Load one weighted module per component and graph variant. 2. Read per-graph metadata. 3. Build model-native example inputs for each graph. 4. Select one weight-sharing module per component. 5. Apply component-level module transforms. 6. Load the tokenizer. 7. Build the model-specific inference helper used by quantization, if any. 8. Optionally export the tokenizer for runtime. 9. Extract a tokenizer chat template, falling back to `extra_options`. `LLMLoaderAdapter` handles text decoders; `MLLMLoaderAdapter` handles multimodal models. The default adapter owns shared tokenizer and transforms logic. ### Dataset stack: `datasets/` Adds typed dataset options, lookup, loaders, calibration adapters, and LLM/MLLM collators. It supports random fallback data, lm-eval samples, JSON messages, Hugging Face chat datasets, and multimodal messages. `generate_calibration_data()` now receives `example_inputs` explicitly because collators require the calibration-graph signature. ### Quantization: `strategies/quantization/`: `ExecuTorchQuantizationStrategy` implements native PTQ for component and graph keyed models: 1. Create one QNN quantizer per component graph. 2. Export every graph variant and prepare it with observers. 3. Initialize observers for deploy graphs. 4. Build real calibration inputs through the calibration data adapter. 5. Run PTQ calibration on graphs that consume real data. 6. Convert prepared graphs to QDQ modules. 7. Propagate encodings from quantize graphs to deploy graphs. 8. Remove quantize-only graphs from the output. Decoder and token embedding use separate quantize and deploy graphs; encoders use one shared graph. Encodings are propagated before lowering. Quantizer creation belongs to the strategy; LLM/MLLM adapters own only model-family-specific export, preparation, conversion, and calibration. ### Quantization Helpers: `quant_utilities.py`: Provides QNN quantizer construction, recipe application, QDQ saving, logits/KV-cache attributes, and encoding propagation. KV-cache override is enabled only when `n_cache_layers` is provided. ### Quantization Output `QuantizationOutputConfig` now returns component and graph keyed `GraphBundle` objects. Each bundle carries the quantized graph module, export inputs, metadata, and optional quantized IO dtypes. Quantize-only graphs are removed after their encodings have been propagated, so downstream compilation sees only deployment graphs. ## Tests Covers CLI/config construction, source transforms, datasets and collators, loader adapters, model preparation, quantization adapters and strategy, and preparation/quantization stage integration. ## PR Review Checklist - All dependencies are injected via constructor with sensible defaults - Yes. - All external calls are behind injectable interfaces - Yes - Hugging Face, `TokenizerWrapper`, `torch.export`, PT2E, QNN quantizer construction, dataset loading, and calibration execution sit behind adapters or lookup-created callables. - Yes. - Unit tests cover the new public behavior - Yes. - Docstrings on public classes and methods - Yes. - Type annotations on function signatures - Yes. - Logging follows the strategy in the LLD - Yes; info for stage-level progress, debug for per-step details, warning for degraded/normalized CLI behavior. No existing behavior changed - Yes - `examples/` is not modified and `llama.py` remains the reference path. - Yes ## Related PRs Phase 2 flow. PR-A1 and PR-B1 are independent and can land in either order. PR-B2 depends on PR-B1. PR-A2 depends on PR-A1 and PR-B2, and closes Phase 2 by wiring the runner and adding the parity test that gates Phase 3 deletion of the legacy flow. ```text PR-A1 (this PR) ────────────────┐ ├─► PR-A2 PR-B1 ───────────► PR-B2 ───────┘ ``` Phase 1 (merged): - PR 1: Core data model, engine routing & exceptions: #20409 - PR 2: Strategy interfaces & stage wrappers: #20795 - PR 3: Pipeline orchestrator: #21149 - PR 4: Adapter interfaces, default implementations & dataset providers: #21751 - PR 5: Model preparation & quantization strategy implementations: #21899 - PR 6: Compilation & inference strategy implementations: #22284 - PR 7: Integration & E2E tests: cancelled, will be raised after Phase 2 is done. Phase 2: - PR A1: Model registry, model_lookup, CLI, Multimodal LLM & LLM model preparation, dataset-backed calibration, component/graph-aware quantization, and encoding reconciliation. [This PR]. - PR B1: Compilation foundation: #22846 - PR B2: Multi-graph lowering, weight sharing, sharding, and spill-fill. Depends on PR-B1. - PR A2: Device-runner adapter, pipeline-runner wiring, README, and E2E parity test. Depends on PR-A1 and PR-B2. ## Test plan ### Run only tests added in this PR: ``` bash python -m pytest \ backends/qualcomm/genai_pipeline/tests/strategies/model_preparation/ \ backends/qualcomm/genai_pipeline/tests/strategies/quantization/ \ backends/qualcomm/genai_pipeline/tests/source_transform/ \ -v ``` Result: ``` log 126 passed, 2 subtests passed in 11.59s ``` ### Run all genai_pipeline tests: ``` bash python -m pytest backends/qualcomm/genai_pipeline/tests/ -v ``` Result: ``` log 238 passed, 13 subtests passed in 14.37s ``` ### Run all tests with coverage: ``` bash python -m pytest backends/qualcomm/genai_pipeline/tests/ \ --cov=backends/qualcomm/genai_pipeline \ --cov-config=backends/qualcomm/.coveragerc \ --cov-report=term-missing ``` Result: ``` log 238 passed, 13 subtests passed in 17.41s ``` ### Confirm the legacy flow is unaffected: ``` bash python -c " from executorch.examples.qualcomm.oss_scripts.llama.llama import _build_parser, export_llama from executorch.examples.qualcomm.oss_scripts.llama.wrappers import MultiModalManager, HybridAttentionSinkEvictor print('LEGACY IMPORTS OK') " ``` Result: ``` log `LEGACY IMPORTS OK` ```
Summary
First PR of Phase 2 and the start of Stream B (compilation). Phase 1 added the skeleton, interfaces, stage wrappers, typed configs with stubbed logic. This PR adds the compilation groundwork the rest of Stream B builds on and makes
DefaultCompilerAdapteractually compile a graph to a.pte.Nothing under
examples/changes. The legacyllama.pyflow is untouched and stays CI's reference; the two coexist until Phase 3, whose deletion of the legacy path is gated on a parity test added in PR-A2.What's included
Two separate name spaces:
artifact_keys.py,graph_names.pyAn artifact key names a
.ptefile; a graph name names a method inside one. A hybrid decoder produces one.pte(text_decoder) holding two methods (kv_forward,prefill_forward), so these are deliberately distinct key spaces.pte_paths, so a compiled bundle reaches inference without translation. PlusDECODE_QDQ_FILENAME, a filename (a.pt2the SQNR eval reads, never seen by the runner) and therefore deliberately absent fromALL_ARTIFACT_KEYS.DECODER_GRAPH_NAMESis decode first, and the order matters:kvmode builds no prefill graph, and decode is the authoritative source of the artifact's constant methods.examples/.../decoder_constants.pyrather than importing them, so this package does not depend on the example scripts. Agreement tests keep them in step; the legacy copies go away with the legacy flow.graph_bundle.py:GraphBundle, the frozen per-graph unit quantization hands to compilation (module,inputs,meta,quant_io_dtypes,modality_inputs,executorch_config){graph_name: GraphBundle}map rather than several parallel dicts that can drift apart.frozen=Truestops in-place edits across a stage boundary. Immutability stops at the field boundary though,convert_pt2emutatesmetain place and the docstring says so, so the guarantee is not overread.quant_io_dtypesis the one thing compilation cannot derive: the graph-boundary dtypes follow from the quantization recipe's bit widths.Nonemeans quantization was skipped.__post_init__requires both keys or neither, because the tagger indexes both unconditionally so a partial mapping is aKeyErrorat lowering, not a partially-quantized boundary.CompilationInputConfig.graphsis the consumer side, added here; PR-A1 adds the producer side, PR-B2 is the first stage to read it. Defined here so both streams import one definition.compilation/compile_spec_builder.py:QnnCompileSpecBuilder,resolve_soc_model,resolve_backend_typegenerate_htp_compiler_spec/generate_gpu_compiler_specwithgenerate_qnn_executorch_compiler_specbehind onebuild();llama.pyrepeats that branch three times today.resolve_soc_modelconvertsstr→QcomChipset.PipelineContextcarries the SoC as a string so nothing in the pipeline imports the QNN schema; lowering needs the enum. Per earlier review this belongs in the adapter layer.resolve_backend_typeuses a lookup map built from the enum itself, sinceQnnExecuTorchBackendType.__str__already yields the CLI name so no hardcoded table, and a bad name raisesValueErrorlisting the valid ones.enable_x86_64is constructor state, not a per-call flag: the emulator supports neither weight sharing nor shared buffers, which is exactly whyllama.pywritesnot args.enable_x86_64at all three call sites. Hardcoding either on would break the x86 CI run without failing anything locally.shared_bufferis on, whileControlArgs.shared_bufferdefaults off to mirrorllama.py's parser. A caller driving the builder from aControlArgsmust pass it explicitly. Documented onbuild()and pinned by two tests.control_args.py: typed bridge fromPipelineContextto the existing LLM componentsargparse.Namespace(args.max_seq_len,args.enable_x86_64, …), so this is a dataclass carrying the same 65 attributes, buildable from aPipelineContext.argparse.Namespace, and that is load-bearing:QnnConfig.load_configdispatches onisinstance(config, argparse.Namespace), so a plain dataclass would break the device path.build_parser()derives anargparseparser from the same fields, so a CLI entry point stops maintaining a second default table.llama.pyparser by a test, since this is a second configuration surface and drift would be silent. Two fields are explicitly exempt (train_config,lr_configdefault toNonerather than to YAML paths underexamples/), with a test pinning that the exemption is deliberate.artifact_paths:List[Path]→Dict[str, Path](CompilationResult,CompilationOutputConfig,InferenceInputConfig,DeviceRunnerAdapter,DefaultDeviceRunnerAdapter)Keyed because the runner addresses artifacts individually and which ones exist varies by model, so position carries no reliable meaning. Absent artifacts are omitted, never mapped to
Noneso the runner tests key membership to decide whether a model is multimodal. This is the frozen B→A2 interface.DefaultCompilerAdapterimplementation: replaces PR6'sNotImplementedErrorMemoryPlanningPass(alloc_graph_input=False, alloc_graph_output=False)+BuildQuantIo(), since with a shared buffer the caller supplies graph I/O addresses from RPC memory), writes the.pte, and returns it under theartifact_keythe caller names keyed by the caller because a 1:1 adapter cannot tell a vision encoder from a text decoder when both arrive asmodel.soc_modelandbackend_typeare validated againstcompile_specsrather than trusted. Lowering takes its target only from the specs, so these parameters would otherwise be decorative: a caller could validate ops for one SoC and compile for another and see nothing until the artifact ran on device. The target is read back out of the specs withflatbuffer_to_optionand a disagreement raises. Specs with no QNN entry skip the check..pteis the strategy's job, so the adapter stays a 1:1 wrapper) and spill-fill sizing (a property of the group, computed from the lowered program). Both land in PR-B2.Two bugs found along the way (worth a look, not specific to this feature)
DefaultDeviceRunnerAdapter.push_artifactsdid[str(p) for p in artifact_paths]. Once that is a dict, iteration yields keys, so it would have pushed the strings"text_decoder","tok_embedding"… instead of paths. Now.values(), matchingEvalBase._get_adb. This is why the type change could not be a pure annotation edit.CompilationInputConfigimportedCompileSpecfromexecutorch.exir.backend.compile_spec, which does not exist (it iscompile_spec_schema). Hidden at runtime byTYPE_CHECKING; mypy flags it.Tests: 6 new files, 8 updated. Beyond the unit tests,
DefaultCompilerAdapterwas run against the real stack (a small module compiled through QNN to a non-empty.pteon the x86 path); that takes ~15s so the committed adapter tests mock lowering, while the builder tests keep two real-API cases (specs really areCompileSpecs;use_multi_contexts+online_prepareis rejected rather than masked).PR Review Checklist
examples/files touched;llama.pyand the wrappers import unchanged).Related PRs
Phase 2 flow.
PR-A1andPR-B1are independent and can land in either order.PR-B2depends onPR-B1.PR-A2depends on all three, and closes Phase 2 by wiring the runner and adding the parity test that gates the Phase-3 deletion of the legacy flow.Phase 1 (merged):
Phase 2:
model_lookup, CLI,LLMQuantizerAdapter, multi-graph quantization & encoding reconciliation. ConsumesControlArgs,GraphBundleandgraph_namesfrom this PR, but does not depend on it landing first; pendingllama_stories_260kE2E parity test. Depends on PR-A1 and PR-B2; pendingTest plan
Run only tests added in this PR:
Result:
Run all
genai_pipelinetests:Result:
Run all tests with coverage:
Result:
Confirm the legacy flow is unaffected:
Result: