Add deduplication pass for initializer tensors (#66) #67

AbhishekHerbertSamuel · 2025-06-05T05:53:07Z

Summary

This PR adds a new graph transformation pass: DeduplicateInitializersPass.

It removes duplicate initializer tensors (typically model weights) based on a unique fingerprint derived from:

Tensor byte content (tobytes())
Data type (dtype)
Shape

All redundant initializers are removed, and nodes referencing them are updated to use the canonical (first-seen) tensor.

Implementation Details

Fingerprints are tracked using a dictionary: (tobytes, dtype, shape) → name
Redundant initializers are removed using graph.initializers.pop(...)
Node inputs are updated via node.replace_input_with(...) for correctness and safety

Benefits

Reduces memory and file size by eliminating duplicated weight tensors
Simplifies graph structure for downstream optimization and export

File Added

src/onnx_ir/passes/common/deduplicate_initializers.py

Closes

Closes #66

justinchuby · 2025-06-05T06:04:09Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+        # Iterate over all initializers in the graph
+        for initializer in list(graph.initializers.values()):
+            key = (
+                initializer.const_value.tobytes(),              # Content fingerprint


This is memory consuming and thus highly inefficient. Consider comparing the dtype and shape first, and only compare values when you need to

justinchuby · 2025-06-05T06:04:59Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+from onnx_ir.passes.base import GraphTransformPass
+
+
+class DeduplicateInitializersPass(GraphTransformPass):


I understand this may be generated by some AIs. Please ensure the class names etc. are correct, and follow coding style from other files in this directory.

justinchuby · 2025-06-05T06:06:36Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+                seen[key] = initializer.name
+
+        # Update node inputs to use the canonical initializer names
+        for node in graph:


you probably need to check nodes from the subgraphs too. You may use the ir.traversal recursive iterator for this.

Hi @justinchuby,
I’ve addressed your feedback in the latest commit:

Optimized memory usage by grouping by (dtype, shape) before comparing tobytes()

Used iterate_graph(graph) to handle nodes in subgraphs as well

Let me know if any further changes are needed. Thanks again for the thoughtful review!

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

justinchuby

It’s fine to use an AI for contribution. Please ensure however that the code actually works

AbhishekHerbertSamuel · 2025-06-05T14:25:38Z

Thank you for the feedback, Justin. I'll check if it works and then only send it here.

…bgraph traversal Address reviewer feedback: - Optimized memory by grouping by dtype and shape before comparing values - Used iterate_graph to handle subgraphs - Validated on normal and subgraph models; deduplication works as expected Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel · 2025-06-06T06:32:17Z

Hi Justin,

Thanks again for your feedback! I've verified that the updated implementation works as intended. Here's the test setup and output: (I ran the test locally, didn't push it here)

Local file path for the test: /Users/abhishekherbertsamuel/ir-py/src/test_local_dedup.py

Test code:
import numpy as np
from onnx_ir._core import Graph, Node, Tensor, Value
from onnx_ir.passes.common.deduplicate_initializers import DeduplicateInitializersPass
def test_normal_and_subgraph_dedup():
print("\n=== TEST: Normal Graph and Subgraph Deduplication ===")

# Shared tensor content
arr = np.array([1, 2, 3])
t1 = Tensor(arr)
t2 = Tensor(arr.copy())  # clone with same content

# Main graph values
v1 = Value(name="w1", const_value=t1)
v2 = Value(name="w2", const_value=t2)

# Subgraph has its own separate Value object (same tensor, new graph-safe instance)
sub_tensor = Tensor(arr.copy())
sub_val = Value(name="w3", const_value=sub_tensor)

# Subgraph node and graph
sub_node = Node("", "Conv", inputs=[sub_val], outputs=[])
subgraph = Graph(
    inputs=[],
    outputs=[],
    nodes=[sub_node],
    initializers=[sub_val],
    name="subgraph"
)

# Main graph node
main_node = Node("", "Add", inputs=[v1, v2], outputs=[])

# Attach subgraph manually to the node (mimics nested block structure)
main_node.blocks = [subgraph]

# Construct main graph
parent_graph = Graph(
    inputs=[],
    outputs=[],
    nodes=[main_node],
    initializers=[v1, v2],
    name="main_graph"
)

print("Before Deduplication:")
print("Main Graph Initializers:", list(parent_graph.initializers.keys()))
print("Main Node inputs:", [v.name for v in main_node.inputs])
print("Subgraph Initializers:", list(subgraph.initializers.keys()))
print("Subgraph Node inputs:", [v.name for v in sub_node.inputs])

# Apply deduplication
DeduplicateInitializersPass().apply(parent_graph)

print("\nAfter Deduplication:")
print("Main Graph Initializers:", list(parent_graph.initializers.keys()))
print("Main Node inputs:", [v.name for v in main_node.inputs])
print("Subgraph Initializers:", list(subgraph.initializers.keys()))
print("Subgraph Node inputs:", [v.name for v in sub_node.inputs])

if name == "main":
test_normal_and_subgraph_dedup()

Test Screenshot: (Have uploaded it here)

If I have missed out on anything, please let me know.

With regards,
Abhishek Herbert Samuel

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel · 2025-06-07T18:57:11Z

Hi @justinchuby,

I've pushed the finalized implementation and test as separate, signed commits. The following have been addressed:

DeduplicateInitializersPass: Added under passes/common, follows repo conventions, uses (dtype, shape) → {tobytes: name} grouping for memory efficiency, and traverses all subgraphs via RecursiveGraphIterator.

Test coverage: A dedicated unittest verifies correct deduplication in the main graph and ensures subgraphs remain isolated.

Coding standards: Followed the structure and documentation style of other passes (e.g., topological_sort.py).

Commit signed: Used -s with a clean message summarizing the functionality.

I have also attached a screenshot of the unit test which passed successfully on my local copy of this repository.

Please let me know if any final changes are needed. Thanks again for your guidance and mentorship throughout this PR!

Best,
Abhishek Herbert Samuel

justinchuby · 2025-06-07T19:22:07Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+from onnx_ir.traversal import RecursiveGraphIterator
+
+
+class DeduplicateInitializersPass:


Suggested change

class DeduplicateInitializersPass:

class DeduplicateInitializersPass(ir.passes.InPlacePass):

please subclass ir.passes.InPlacePass. You can use https://github.com/AbhishekHerbertSamuel/ir-py/blob/ef46092b5f10303bb9fe126eef0f5b44585e3b16/src/onnx_ir/passes/common/constant_manipulation.py#L23 as an example.

justinchuby · 2025-06-07T19:24:07Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+    using RecursiveGraphIterator.
+    """
+
+    def apply(self, graph: Graph) -> Graph:


Please implement the call method. The first argument should be an ir.Model. You may use other passes in this directory as examples. Be sure to import modules only: https://google.github.io/styleguide/pyguide.html#224-decision

justinchuby · 2025-06-07T19:25:15Z

src/onnx_ir/passes/common/deduplicate_initializers_test.py

+import unittest
+import numpy as np
+
+from onnx_ir._core import Tensor, Value, Node, Graph


Please import modules only. You may use https://github.com/onnx/ir-py/blob/main/src/onnx_ir/passes/common/unused_removal_test.py as an example

justinchuby · 2025-06-07T19:30:43Z

Please feel free to ask questions when you are going through the code base or need help understanding parts of the code. It would be helpful to take a look at other existing passes and usages to ensure they are implemented in a similar style.

justinchuby · 2025-06-07T19:33:27Z

My concern with this pass in particular is that we are using the full bytes in the look up table. This is memory intensive. I wonder if there is a good (efficient) hash method that can be apply to the bytes content, and use the hash value in the look up table. Only when the hash matches do we compare the actual bytes.

AbhishekHerbertSamuel · 2025-06-08T04:06:24Z

Hi @justinchuby,
Thanks a lot for your detailed feedback :)

I’ll update the class to inherit from ir.passes.InPlacePass as suggested and move the main logic into the call method, following the repo’s conventions (like in constant_manipulation.py).
I’ll also change the test imports to follow the module-only import guideline — thanks for pointing me to the correct example!

Regarding the memory concern:
You're absolutely right — using tobytes() directly is memory-intensive. I’ll switch to using sha256 to hash the tensor bytes first, which helps group potential duplicates quickly. Then, to avoid any risk of false positives from rare hash collisions, I’ll still compare the full bytes only when the hashes match. This keeps things memory-efficient while still being safe and accurate. Thanks again for the suggestion!

Will push the changes shortly. Please let me know if I missed anything else. Appreciate your guidance!

Warm regards,
Abhishek Herbert Samuel

- Implemented DeduplicateInitializersPass to remove redundant initializers with identical shape, dtype, and values within individual graphs. - Ensured deduplication is confined to the same graph scope (no cross-subgraph merging). - Added unit tests covering: - Exact duplicates - Different shapes/dtypes - Scalars - Multiple duplicates - Non-deduplicable distinct values - Removed subgraph-related tests due to ONNX serialization behavior omitting their initializers. Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel · 2025-06-09T10:21:08Z

Hi @justinchuby,
I've pushed the finalized version of DeduplicateInitializersPass along with a focused set of unit tests. The current tests comprehensively validate deduplication behavior across various scenarios—shape, dtype, scalar, and value uniqueness.

Tests involving subgraph initializers were removed, as ONNX drops those during serialization, making them unreliable to assert against. Let me know if you'd like a different strategy for subgraph coverage.

Thanks again for your guidance throughout!

Warm regards,
Abhishek Herbert Samuel

codecov · 2025-06-09T10:23:27Z

Codecov Report

Attention: Patch coverage is 76.47059% with 12 lines in your changes missing coverage. Please review.

Project coverage is 73.71%. Comparing base (6656096) to head (8f6dbdc).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
.../onnx_ir/passes/common/deduplicate_initializers.py	76.47%	5 Missing and 7 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #67      +/-   ##
==========================================
+ Coverage   73.57%   73.71%   +0.14%     
==========================================
  Files          37       38       +1     
  Lines        4492     4543      +51     
  Branches      902      915      +13     
==========================================
+ Hits         3305     3349      +44     
- Misses        858      861       +3     
- Partials      329      333       +4

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

src/onnx_ir/passes/common/deduplicate_initializers.py

src/onnx_ir/passes/common/deduplicate_initializers_test.py

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel · 2025-06-11T08:54:29Z

Hi @justinchuby, tried to make the remaining changes based on the CI workflow's results. While my tests ran locally and are fine (6/6 correct), the codecov bot and workflows were not triggered here automatically as it was the previous time. Is that ok or a sign of an error?

With regards,
Abhishek Herbert Samuel

src/onnx_ir/passes/common/deduplicate_initializers.py

src/onnx_ir/passes/common/deduplicate_initializers_test.py

justinchuby · 2025-06-11T21:20:19Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+                        break  # only break when deduplication is successful
+                else:
+                    # no matching content found: append as a new entry
+                    group[content_hash].append((initializer.name, content))


You may store the values instead in the hash table, and use const_val.tobytes() for comparison. This way the bytes do not stay in the memory or take up space?

Suggested change

group[content_hash].append((initializer.name, content))

group[content_hash].append(const_val)

justinchuby · 2025-06-11T21:21:35Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+                        if initializer.name is not None:
+                            graph.initializers.pop(initializer.name)
+                        break  # only break when deduplication is successful
+                else:


justinchuby · 2025-06-11T21:22:04Z

Thank you! could you also take a look at the lint errors?

AbhishekHerbertSamuel · 2025-06-12T05:27:46Z

Hi @justinchuby, will work on it and send the corrected code here:)

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel · 2025-06-12T11:54:11Z

Hi @justinchuby, have committed it with the lint issues and byte optimization addressed. Please find attached the screenshot here, showing the absence of lint errors in my local system.

Do let me know if there are any other changes :)

justinchuby · 2025-06-12T15:16:14Z

src/onnx_ir/passes/common/deduplicate_initializers.py

@@ -27,16 +27,16 @@ class DeduplicateInitializersPass(onnx_ir.passes.InPlacePass):

    def call(self, model: onnx_ir.Model) -> onnx_ir.passes.PassResult:


Can you import onnx_ir as ir and use it as such, to stay consistent with the rest of the code base?

justinchuby · 2025-06-12T16:11:47Z

Thanks! I will do a more detailed review soon

AbhishekHerbertSamuel · 2025-06-12T17:31:01Z

Sure @justinchuby, will fix it and maintain code consistency :)

xadupre · 2025-06-12T18:06:57Z

src/onnx_ir/passes/common/deduplicate_initializers.py

+                continue  # Skip if initializer has no constant value
+            dtype = const_val.dtype.name
+            shape = tuple(int(dim) if isinstance(dim, int) else -1 for dim in const_val.shape)
+            content = const_val.tobytes()


This is going to be very slow on big tensors. Is there a way to compare rawdata directly to save some time?

After discussion it seems a good idea to avoid comparing big tensors at all. @AbhishekHerbertSamuel Could you limit the size to 1024 values? You can find the element count of the tensor with tensor.size

Threshold 1024 is a very small number in production, there is case where model size is reduced from 150MB to 100MB by tying weights, maybe we can parameterize this.

And when comparing those constants, we can compare shape and dtype first, instead of comparing the raw data directly, this will save tremendous time.

Thanks. I agree that’s a good idea

Copilot

Pull Request Overview

Adds a new graph transformation pass to remove duplicate initializer tensors by hashing their content and updating node inputs to the canonical tensor.

Introduces DeduplicateInitializersPass to group initializers by dtype, shape, and content hash, remove duplicates, and rewrite node inputs.
Implements content‐based deduplication with SHA-256 fingerprinting and exact byte comparison to avoid collisions.
Provides unit tests covering identical, shape/dtype differences, scalar, multiple duplicates, and unique-value scenarios.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
src/onnx_ir/passes/common/deduplicate_initializers.py	Implements new pass to deduplicate initializer tensors
src/onnx_ir/passes/common/deduplicate_initializers_test.py	Adds unit tests for various deduplication scenarios

Comments suppressed due to low confidence (3)

src/onnx_ir/passes/common/deduplicate_initializers.py:31

[nitpick] The name 'seen' is ambiguous; consider renaming it to something more descriptive like 'initializer_groups' or 'fingerprint_groups'.

seen: dict[tuple[str, tuple[int, ...]], dict[str, list[onnx_ir.Value]]] = {}

src/onnx_ir/passes/common/deduplicate_initializers.py:33

[nitpick] The variable 'name_map' could be renamed to 'duplicate_to_canonical' or similar to clarify its purpose.

name_map = {}

src/onnx_ir/passes/common/deduplicate_initializers.py:67

There are no tests covering deduplication behavior within nested subgraphs; consider adding a test case to verify that duplicate initializers in subgraphs are also handled.

for node in onnx_ir.traversal.RecursiveGraphIterator(graph):

src/onnx_ir/passes/common/deduplicate_initializers.py

AbhishekHerbertSamuel · 2025-06-13T04:44:13Z

Thank you @xadupre @inisis @justinchuby for the feedback. Will make the requested changes and ensure that the PR is ready to be merged.

…nd size limit - Avoids comparing large tensors >1024 elements to reduce performance overhead - Compares shape and dtype before accessing tensor content - Adds test coverage for subgraph deduplication (If node branches) - Passes all linters: ruff, mypy, editorconfig Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel · 2025-06-13T18:15:33Z

@xadupre @justinchuby @inisis I have made the requested changes. Please check and let me know if it's ready for merging or if other changes need to be made prior to that. Thank you once again :)

AbhishekHerbertSamuel requested review from titaiwangms and a team as code owners June 5, 2025 05:53

AbhishekHerbertSamuel mentioned this pull request Jun 5, 2025

Create a tensor de-duplication pass #66

Open

justinchuby reviewed Jun 5, 2025

View reviewed changes

AbhishekHerbertSamuel added 2 commits June 5, 2025 13:27

Add deduplication pass for initializer tensors (onnx#66)

159e89a

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

Address feedback: optimize tensor fingerprinting and traverse subgraphs

ae8f078

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel force-pushed the add-deduplicate-initializers-pass branch from f99fa0c to ae8f078 Compare June 5, 2025 07:57

justinchuby requested changes Jun 5, 2025

View reviewed changes

Add DeduplicateInitializersPass and test covering graph and subgraph

ef46092

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

justinchuby reviewed Jun 7, 2025

View reviewed changes

github-advanced-security bot found potential problems Jun 9, 2025

View reviewed changes

justinchuby self-assigned this Jun 9, 2025

Finalize DeduplicateInitializersPass implementation and test coverage

6b3e0b7

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

AbhishekHerbertSamuel force-pushed the add-deduplicate-initializers-pass branch from a00be10 to 6b3e0b7 Compare June 11, 2025 08:40

AbhishekHerbertSamuel added 2 commits June 11, 2025 14:13

Finalize DeduplicateInitializersPass implementation and test coverage

b318c63

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

Finalize DeduplicateInitializersPass implementation and test coverage

5507fd3

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

github-advanced-security bot found potential problems Jun 11, 2025

View reviewed changes

Merge branch 'main' into add-deduplicate-initializers-pass

5f1c634

justinchuby reviewed Jun 11, 2025

View reviewed changes

AbhishekHerbertSamuel added 2 commits June 12, 2025 15:13

Address ruff and lint errors in DeduplicateInitializersPass

5dcc28e

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

Address ruff and lint errors in DeduplicateInitializersPass

a3a350b

Signed-off-by: Abhishek Herbert Samuel <[email protected]>

justinchuby reviewed Jun 12, 2025

View reviewed changes

xadupre reviewed Jun 12, 2025

View reviewed changes

justinchuby requested a review from Copilot June 13, 2025 03:08

Copilot AI reviewed Jun 13, 2025

View reviewed changes

src/onnx_ir/passes/common/deduplicate_initializers.py Show resolved Hide resolved

src/onnx_ir/passes/common/deduplicate_initializers.py Outdated Show resolved Hide resolved

		from onnx_ir.passes.base import GraphTransformPass


		class DeduplicateInitializersPass(GraphTransformPass):

		from onnx_ir.traversal import RecursiveGraphIterator


		class DeduplicateInitializersPass:

	class DeduplicateInitializersPass:
	class DeduplicateInitializersPass(ir.passes.InPlacePass):

	group[content_hash].append((initializer.name, content))
	group[content_hash].append(const_val)

		@@ -27,16 +27,16 @@ class DeduplicateInitializersPass(onnx_ir.passes.InPlacePass):

		def call(self, model: onnx_ir.Model) -> onnx_ir.passes.PassResult:

Add deduplication pass for initializer tensors (#66) #67

Are you sure you want to change the base?

Add deduplication pass for initializer tensors (#66) #67

Uh oh!

Conversation

AbhishekHerbertSamuel commented Jun 5, 2025

Summary

Implementation Details

Benefits

File Added

Closes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

justinchuby left a comment

Choose a reason for hiding this comment

Uh oh!

AbhishekHerbertSamuel commented Jun 5, 2025

Uh oh!

AbhishekHerbertSamuel commented Jun 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

AbhishekHerbertSamuel commented Jun 7, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

justinchuby commented Jun 7, 2025

Uh oh!

justinchuby commented Jun 7, 2025

Uh oh!

AbhishekHerbertSamuel commented Jun 8, 2025

Uh oh!

AbhishekHerbertSamuel commented Jun 9, 2025

Uh oh!

codecov bot commented Jun 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AbhishekHerbertSamuel commented Jun 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

justinchuby commented Jun 11, 2025

Uh oh!

AbhishekHerbertSamuel commented Jun 12, 2025

Uh oh!

AbhishekHerbertSamuel commented Jun 12, 2025

Uh oh!

Choose a reason for hiding this comment

AbhishekHerbertSamuel commented Jun 6, 2025 •

edited

Loading

codecov bot commented Jun 9, 2025 •

edited

Loading

AbhishekHerbertSamuel commented Jun 11, 2025 •

edited

Loading