Add an ability to annotate a model with layer annotations.#2361
Open
yuslepukhin wants to merge 6 commits intomainfrom
Open
Add an ability to annotate a model with layer annotations.#2361yuslepukhin wants to merge 6 commits intomainfrom
yuslepukhin wants to merge 6 commits intomainfrom
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a mechanism to capture user-provided “layer annotation” mappings on PyTorch/HF model handlers and apply them during ONNX export/model-building by writing per-node metadata_props["layer_ann"] based on node-name substring matches.
Changes:
- Added
CaptureLayerAnnotationsPyTorch pass to persistlayer_annotationsintomodel_attributes. - Added ONNX-side annotation utilities (
annotate_ir_model/annotate_proto_model) and integrated them intoOnnxConversionandModelBuilder. - Added unit tests for the new PyTorch pass and registered the pass in
olive_config.json.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
test/passes/pytorch/test_capture_layer_annotations.py |
Adds tests validating CaptureLayerAnnotations stores mappings into model_attributes and basic config validation behavior. |
olive/passes/pytorch/capture_layer_annotations.py |
Introduces CaptureLayerAnnotations pass to attach layer_annotations to the model handler. |
olive/passes/onnx/model_builder.py |
Applies layer_annotations to generated ONNX ModelProto by annotating nodes’ metadata props before saving. |
olive/passes/onnx/layer_annotation.py |
Implements substring-based node annotation for both onnxscript IR graphs and ONNX ModelProto, including recursion into subgraphs. |
olive/passes/onnx/conversion.py |
Applies layer_annotations during IR-based ONNX conversion before serializing to an Olive model. |
olive/olive_config.json |
Registers CaptureLayerAnnotations for pipeline integration. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe your changes
This pull request introduces a new mechanism for capturing and propagating layer annotation metadata through the ONNX conversion pipeline. The primary addition is the
CaptureLayerAnnotationspass, which enables users to specify mappings of layer names to node-name substrings; these mappings are then used to annotate ONNX nodes with layer information. The implementation ensures that these annotations are preserved and applied consistently across both ONNX IR and ONNX ModelProto representations. Comprehensive tests have been added to validate the new functionality.Layer annotation capture and propagation:
CaptureLayerAnnotationspass inolive/passes/pytorch/capture_layer_annotations.py, which stores layer annotation mappings in model attributes for downstream ONNX conversion passes.CaptureLayerAnnotationspass inolive/olive_config.jsonfor pipeline integration.ONNX conversion and annotation application:
olive/passes/onnx/conversion.pyandolive/passes/onnx/model_builder.pyto detect and apply layer annotations during ONNX model conversion, annotating nodes whose names match configured substrings. [1] [2]olive/passes/onnx/layer_annotation.pywith functions to annotate both ONNX IR and ModelProto graphs, including recursive annotation of control-flow nodes.Testing and validation:
test/passes/pytorch/test_capture_layer_annotations.pyto verify correct storage, propagation, and validation of layer annotation mappings, including deep-copy semantics and config validation.Checklist before requesting a review
lintrunner -a(Optional) Issue link