Skip to content

Add an ability to annotate a model with layer annotations.#2361

Open
yuslepukhin wants to merge 6 commits intomainfrom
yuslepukhin/convert_layer_annotate
Open

Add an ability to annotate a model with layer annotations.#2361
yuslepukhin wants to merge 6 commits intomainfrom
yuslepukhin/convert_layer_annotate

Conversation

@yuslepukhin
Copy link
Member

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 CaptureLayerAnnotations pass, 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:

  • Added the CaptureLayerAnnotations pass in olive/passes/pytorch/capture_layer_annotations.py, which stores layer annotation mappings in model attributes for downstream ONNX conversion passes.
  • Registered the new CaptureLayerAnnotations pass in olive/olive_config.json for pipeline integration.

ONNX conversion and annotation application:

  • Modified olive/passes/onnx/conversion.py and olive/passes/onnx/model_builder.py to detect and apply layer annotations during ONNX model conversion, annotating nodes whose names match configured substrings. [1] [2]
  • Introduced olive/passes/onnx/layer_annotation.py with functions to annotate both ONNX IR and ModelProto graphs, including recursive annotation of control-flow nodes.

Testing and validation:

  • Added test/passes/pytorch/test_capture_layer_annotations.py to verify correct storage, propagation, and validation of layer annotation mappings, including deep-copy semantics and config validation.

Checklist before requesting a review

  • Add unit tests for this change.
  • Make sure all tests can pass.
  • [?] Update documents if necessary.
  • Lint and apply fixes to your code by running lintrunner -a
  • Is this a user-facing change? If yes, give a description of this change to be included in the release notes.

(Optional) Issue link

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 CaptureLayerAnnotations PyTorch pass to persist layer_annotations into model_attributes.
  • Added ONNX-side annotation utilities (annotate_ir_model / annotate_proto_model) and integrated them into OnnxConversion and ModelBuilder.
  • 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants