Enable ONNX external-data sharding - #503
Conversation
Weights that need a dtype cast are held as ir.LazyTensor and converted at serialization time, so saving alternates between casting a tensor and writing it, leaving both the CPU and the disk idle half the time. onnx_ir gained a max_workers option that overlaps materialization with disk writes and parallelizes both, with peak memory bounded independently of the worker count. Default ModelPackage.save to 8 workers and thread the option through. The option is feature-detected because it is not in onnx_ir 0.2.x, which the package still supports, so older installs keep saving serially instead of raising TypeError. Fix the progress bar for concurrent saving. onnx_ir serializes callbacks with a lock but no longer delivers them in index order, and the bar mutated total/set_description without any synchronization of its own. Guard the closure state with a lock and keep counting invocations rather than tracking indices. Measured on Qwen2.5-7B cast from bfloat16 to float16, 15.26GB of external data: 264.2s serial -> 176.0s with 8 workers (1.50x). 4 workers gives 1.45x and 16 gives 1.48x, so the default of 8 is near the knee. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
max_workers ships in onnx_ir 1.1.0, so depend on it directly and drop the inspect.signature fallback. Merge this after 1.1.0 is released. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Hardcoding 8 happened to match this machine's performance-core count, which was a coincidence rather than a reason. Measuring the phases separately shows the bottleneck is the storage device, not the CPU: casting bf16 to f16 runs at 12.2 GB/s while serial writes reach only 1.6 GB/s. A write-only scaling test saturates at 2 threads (2.1x) and does not improve through 24. So the speedup comes from overlapping the cast with the write, not from write parallelism. That is why the default deliberately does not scale with the core count -- something like cpu_count/2 would oversubscribe against torch's own intra-op pool (already 8 threads here) without making the disk faster. Keep a small constant, but clamp it by os.cpu_count() so small containers do not spawn more threads than they can run, and handle cpu_count() returning None. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
This PR improves ModelPackage.save() performance when saving ONNX external data by enabling concurrent materialization + disk writes via onnx_ir’s new max_workers option when available, while preserving compatibility with older onnx_ir versions through feature detection. It also updates the progress callback behavior and adds test coverage for serial/parallel equivalence and callback ordering/thread-safety.
Changes:
- Add
max_workers(default8) toModelPackage.save()and pass it toir.save(..., max_workers=...)when supported. - Add runtime feature detection for
ir.save(max_workers=...)to avoid breaking olderonnx_irinstalls. - Add tests validating byte-identical output, load roundtrip, callback correctness under out-of-order and concurrent invocations, and feature detection behavior.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| src/mobius/_model_package.py | Adds max_workers support with feature detection; makes progress callback state thread-safe. |
| src/mobius/_model_package_test.py | Adds TestParallelSave verifying serial/parallel equivalence and progress-callback correctness. |
Suppressed comments (1)
src/mobius/_model_package_test.py:163
- This test relies on a specific
__closure__cell index to fetch the tqdm bar. Closure ordering is an implementation detail and can change with refactors/Python versions, making the test fragile. Extract the bar by searching closure cells for the object that has tqdm methods (e.g.,set_description) instead of indexing.
_Tensor(),
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| ), | ||
| ) | ||
| # Closure state is private to the callback; assert via the bound bar. | ||
| bar = callback.__closure__[1].cell_contents |
| with lock: | ||
| if not total_set: | ||
| pbar.total = metadata.total | ||
| total_set = True | ||
| pbar.update() |
Sizing the pool by core count gets the tradeoff backwards. With torch pinned to one intra-op thread, i.e. a small or slow machine, 8 workers is the fastest configuration (1.55x) while 2 workers is slower than serial: more concurrency is needed to hide a slow serial cast, not less. cpu_count() // 2 would give 1 on a 2-core box, disabling the optimization exactly where it helps most. Oversubscription is not penalized in practice either. With full torch threads, everything from 4 to 32 workers lands within noise (1.37x-1.45x), because these threads block in tofile with the GIL released rather than competing for CPU. Drop the core-count clamp and keep a plain constant. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Forward max_shard_size_bytes to ir.save instead of limiting the option to safetensors. Update the CLI help and add a behavioral test that writes, discovers, and reloads three ONNX external-data shards. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
70350a0 to
91996c9
Compare
Do not expose or pass max_workers or ONNX max_shard_size_bytes from Mobius. Those APIs belong to the pending onnx_ir change, and calling them would make Mobius fail before that change is released. Require only onnx_ir 1.0.0 so the Mobius update is independent of the ir-py PR. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Forward max_shard_size_bytes for ONNX external data, an API already available in onnx_ir 1.0.0. Keep Mobius independent of the pending parallel-save change by not exposing or passing max_workers. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Performance Comparison
|
🏗️ Architecture Diff
falcon / model — 8 change(s)Op summary: 66 → 68 nodes --- base
+++ head
@@ -34,7 +34,8 @@
LayerNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Add
@@ -57,7 +58,8 @@
LayerNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
AddAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
gemma2 / model — 12 change(s)Op summary: 105 → 107 nodes --- base
+++ head
@@ -54,7 +54,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -87,7 +88,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
gemma4 (gemma4) / decoder — 55 change(s)Op summary: 125 → 127 nodes --- base
+++ head
@@ -64,7 +64,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -107,7 +108,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -119,7 +121,7 @@
RMSNormalization
Transpose
MatMul
-Constant
+CastLike
Div
Tanh
MulAdded nodes:
Removed nodes:
Connectivity changes:
Initializer changes:
gemma4 (gemma4) / embedding — 2 change(s)Op summary: 19 → 18 nodes --- base
+++ head
@@ -1,5 +1,4 @@
Gather
-Constant
Mul
Constant
EqualRemoved nodes:
Initializer changes:
gemma4_text / model — 55 change(s)Op summary: 128 → 129 nodes --- base
+++ head
@@ -1,5 +1,4 @@
Gather
-Constant
Mul
Gather
Gather
@@ -67,7 +66,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -110,7 +110,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -122,7 +123,7 @@
RMSNormalization
Transpose
MatMul
-Constant
+CastLike
Div
Tanh
MulAdded nodes:
Removed nodes:
Connectivity changes:
Initializer changes:
llama / model — 8 change(s)Op summary: 60 → 62 nodes --- base
+++ head
@@ -25,7 +25,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -48,7 +49,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
llama (static-cache) / model — 9 change(s)Op summary: 56 → 58 nodes --- base
+++ head
@@ -19,7 +19,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -44,7 +45,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
mamba (ssm-text-generation) / model — 23 change(s)Op summary: 94 → 98 nodes --- base
+++ head
@@ -9,7 +9,8 @@
Conv
Unsqueeze
Add
-Swish
+Sigmoid
+Mul
Transpose
Transpose
MatMul
@@ -40,7 +41,8 @@
Mul
Add
Unsqueeze
-Swish
+Sigmoid
+Mul
Mul
Transpose
MatMul
@@ -54,7 +56,8 @@
Conv
Unsqueeze
Add
-Swish
+Sigmoid
+Mul
Transpose
Transpose
MatMul
@@ -85,7 +88,8 @@
Mul
Add
Unsqueeze
-Swish
+Sigmoid
+Mul
Mul
Transpose
MatMulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
phi3 / model — 9 change(s)Op summary: 58 → 60 nodes --- base
+++ head
@@ -26,7 +26,8 @@
Transpose
MatMul
Split
-Swish
+Sigmoid
+Mul
Mul
Transpose
MatMul
@@ -48,7 +49,8 @@
Transpose
MatMul
Split
-Swish
+Sigmoid
+Mul
Mul
Transpose
MatMulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
phi3 (static-cache) / model — 10 change(s)Op summary: 54 → 56 nodes --- base
+++ head
@@ -20,7 +20,8 @@
Transpose
MatMul
Split
-Swish
+Sigmoid
+Mul
Mul
Transpose
MatMul
@@ -44,7 +45,8 @@
Transpose
MatMul
Split
-Swish
+Sigmoid
+Mul
Mul
Transpose
MatMulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
qwen / model — 8 change(s)Op summary: 60 → 62 nodes --- base
+++ head
@@ -25,7 +25,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -48,7 +49,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
qwen (static-cache) / model — 9 change(s)Op summary: 56 → 58 nodes --- base
+++ head
@@ -19,7 +19,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -44,7 +45,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
qwen2 / model — 8 change(s)Op summary: 60 → 62 nodes --- base
+++ head
@@ -25,7 +25,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -48,7 +49,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
qwen2 (static-cache) / model — 9 change(s)Op summary: 56 → 58 nodes --- base
+++ head
@@ -19,7 +19,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -44,7 +45,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
qwen2_moe / model — 295 change(s)Op summary: 214 → 224 nodes --- base
+++ head
@@ -41,66 +41,71 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -138,77 +143,82 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Transpose
-MatMul
-Sigmoid
-Mul
-Add
-Add
-RMSNormalization
-Transpose
-MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Transpose
+MatMul
+Sigmoid
+Mul
+Add
+Add
+RMSNormalization
+Transpose
+MatMulAdded nodes:
Removed nodes:
Connectivity changes:
qwen2_moe (static-cache) / model — 295 change(s)Op summary: 204 → 214 nodes --- base
+++ head
@@ -29,66 +29,71 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -128,77 +133,82 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Transpose
-MatMul
-Sigmoid
-Mul
-Add
-Add
-RMSNormalization
-Transpose
-MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Transpose
+MatMul
+Sigmoid
+Mul
+Add
+Add
+RMSNormalization
+Transpose
+MatMulAdded nodes:
Removed nodes:
Connectivity changes:
qwen3 / model — 9 change(s)Op summary: 72 → 74 nodes --- base
+++ head
@@ -31,7 +31,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -60,7 +61,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
qwen3 (static-cache) / model — 10 change(s)Op summary: 68 → 70 nodes --- base
+++ head
@@ -25,7 +25,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -56,7 +57,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
qwen3_5_moe (hybrid-text-generation) / model — 299 change(s)Op summary: 264 → 275 nodes --- base
+++ head
@@ -62,7 +62,8 @@
Reshape
RMSNormalization
Cast
-Swish
+Sigmoid
+Mul
Cast
Mul
CastLike
@@ -81,66 +82,71 @@
Softmax
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -187,78 +193,83 @@
Softmax
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Transpose
-MatMul
-Sigmoid
-Mul
-Add
-Add
-Add
-RMSNormalization
-Transpose
-MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Transpose
+MatMul
+Sigmoid
+Mul
+Add
+Add
+Add
+RMSNormalization
+Transpose
+MatMulAdded nodes:
Removed nodes:
Connectivity changes:
qwen3_5_text (hybrid-text-generation) / model — 18 change(s)Op summary: 126 → 129 nodes --- base
+++ head
@@ -62,7 +62,8 @@
Reshape
RMSNormalization
Cast
-Swish
+Sigmoid
+Mul
Cast
Mul
CastLike
@@ -76,7 +77,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -113,7 +115,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Connectivity changes:
qwen3_5_vl (hybrid-qwen-vl) / decoder — 19 change(s)Op summary: 149 → 152 nodes --- base
+++ head
@@ -84,7 +84,8 @@
Reshape
RMSNormalization
Cast
-Swish
+Sigmoid
+Mul
Cast
Mul
CastLike
@@ -98,7 +99,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -136,7 +138,8 @@
RMSNormalization
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Connectivity changes:
qwen3_moe / model — 236 change(s)Op summary: 194 → 202 nodes --- base
+++ head
@@ -44,51 +44,55 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -131,64 +135,68 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-CastLike
-Mul
-ReduceSum
-Mul
-Add
-Add
-RMSNormalization
-Transpose
-MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+CastLike
+Mul
+ReduceSum
+Mul
+Add
+Add
+RMSNormalization
+Transpose
+MatMulAdded nodes:
Removed nodes:
Connectivity changes:
qwen3_moe (static-cache) / model — 40 change(s)Op summary: 184 → 192 nodes --- base
+++ head
@@ -32,7 +32,8 @@
Div
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -46,7 +47,8 @@
Mul
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -61,7 +63,8 @@
Add
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -76,7 +79,8 @@
Add
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -121,7 +125,8 @@
Div
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -135,7 +140,8 @@
Mul
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -150,7 +156,8 @@
Add
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -165,7 +172,8 @@
Add
Transpose
MatMul
-Swish
+Sigmoid
+Mul
Transpose
MatMul
MulAdded nodes:
Removed nodes:
Connectivity changes:
qwen3_next (hybrid-text-generation) / model — 615 change(s)Op summary: 562 → 585 nodes --- base
+++ head
@@ -62,7 +62,8 @@
Reshape
RMSNormalization
Cast
-Swish
+Sigmoid
+Mul
Cast
Mul
CastLike
@@ -83,70 +84,75 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -205,7 +211,8 @@
Reshape
RMSNormalization
Cast
-Swish
+Sigmoid
+Mul
Cast
Mul
CastLike
@@ -226,70 +233,75 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -348,7 +360,8 @@
Reshape
RMSNormalization
Cast
-Swish
+Sigmoid
+Mul
Cast
Mul
CastLike
@@ -369,70 +382,75 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
Transpose
MatMul
Mul
@@ -481,82 +499,87 @@
Div
Transpose
MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Constant
-Equal
-Cast
-Mul
-ReduceSum
-Unsqueeze
-Mul
-Add
-Transpose
-MatMul
-Swish
-Transpose
-MatMul
-Mul
-Transpose
-MatMul
-Transpose
-MatMul
-Sigmoid
-Mul
-Add
-Add
-Add
-RMSNormalization
-Transpose
-MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Constant
+Equal
+Cast
+Mul
+ReduceSum
+Unsqueeze
+Mul
+Add
+Transpose
+MatMul
+Sigmoid
+Mul
+Transpose
+MatMul
+Mul
+Transpose
+MatMul
+Transpose
+MatMul
+Sigmoid
+Mul
+Add
+Add
+Add
+RMSNormalization
+Transpose
+MatMulAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
whisper (speech-to-text) / encoder — 2 change(s)Op summary: 50 → 50 nodes No op-sequence changes. Modified attributes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Serialize tqdm mutations and count callback invocations rather than relying on index order. This supports concurrent callbacks from newer onnx_ir versions while remaining compatible with the serial callback behavior in onnx_ir 1.0.0. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Include CallbackInfo.filename in the tqdm description so concurrent ONNX shard saves visibly identify the shard currently reporting progress. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Use CallbackInfo shard metadata from newer onnx_ir releases to create a fixed-position tqdm bar for each external-data file. Fall back to one global bar with onnx_ir 1.0.0, which does not expose per-shard counts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Name the default external data file model.onnx_data so sharding naturally produces model-00001-of-00003.onnx_data rather than inserting the shard suffix between .onnx and .data. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Restore model.onnx.data as the default external-data location. Newer onnx_ir releases preserve this compound suffix when numbering shards; the compatibility test also accepts the legacy 1.0.0 shard spelling. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: af83bc37-b9d6-4b94-be64-a4daa5a7ce40 Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Derive each tqdm row from the one-based shard number embedded in the external-data filename, so concurrent callbacks cannot reorder the displayed bars. Cover shard-two-first callback delivery while retaining the onnx_ir 1.0 single-bar fallback. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Summary
onnx_irversion from 0.2.1 to 1.0.0.--max-shard-sizeto ONNX external-data saves asmax_shard_size_bytes, an API already available inonnx_ir1.0.0.onnx_irversions it renders one fixed-position tqdm bar per shard; withonnx_ir1.0.0 it falls back to one global bar.max_workers.Validation
The ONNX sharding round-trip and callback tests pass against both the published
onnx_ir==1.0.0package and the current onnx/ir-py#476 branch.