Skip to content

ONNX multi-exit export + end-to-end integration tests - #4

Merged
sohams25 merged 1 commit into
mainfrom
feat-onnx-and-integration
Jun 6, 2026
Merged

ONNX multi-exit export + end-to-end integration tests#4
sohams25 merged 1 commit into
mainfrom
feat-onnx-and-integration

Conversation

@sohams25

@sohams25 sohams25 commented Jun 6, 2026

Copy link
Copy Markdown
Owner

Removes the last shipped limitation (ONNX) and adds comprehensive isolation+integration testing.

ONNX export

  • export_to_onnx / earlyon export write a static multi-output graph (one output per exit + final). ONNX has no portable per-sample control flow, so routing stays at runtime — the graph computes every exit and the caller picks the first confident one (honest trade: portability over the early-exit compute saving).
  • Legacy tracer (dynamo=False) traces the static all-exits forward cleanly. Works for conv (4D) and transformer (3D-token) heads, with a dynamic batch axis. Saves/restores the model's train/eval mode (torch.onnx.export otherwise silently leaves it in train).
  • onnxruntime added to the dev extra to verify exported graphs.

Testing — in isolation and together

  • Isolation (tests/test_onnx.py): export cnn + token models, outputs match torch all-exits under onnxruntime, dynamic vs static batch, mode preservation, output-count tracks exits.
  • Together (tests/test_integration.py): full pipelines — build → train → calibrate (confidence + entropy) → analyze → benchmark → save/load → ONNX — across a factory backbone (cifar_resnet) and a custom-wrapped transformer, plus a CUDA end-to-end run.

Test plan

  • ruff · black · isort clean
  • mypy --strict — 0 errors
  • pytest -m "not gpu" — 159 passed, 97% coverage
  • pytest -m gpu — 5 passed (RTX 4050)
  • CI green on this PR

ONNX export (removes the documented limitation):
- export_to_onnx / `earlyon export` write a static multi-output graph (one
  output per exit + final). Routing stays at runtime — the graph computes every
  exit and the caller picks the first confident one.
- legacy tracer is required (the new exporter trips the wrapper's _is_compiling
  guard); the `dynamo` kwarg only exists on torch>=2.4, so it's applied only
  when present (older torch defaults to the legacy exporter anyway). Saves/
  restores the model's train/eval mode; suppresses the exporter's deprecation
  noise.
- `onnx` is a real dependency of the exporter (loaded at call time): declared as
  an `[onnx]` extra with a clear ImportError, and added to `dev` so CI runs the
  tests instead of skipping them.
- works for conv (4D) and transformer (3D-token) heads; dynamic-batch axis.

Testing — in isolation AND together:
- tests/test_onnx.py: export match vs torch (cnn+token), dynamic/static batch,
  mode preservation, output count, and the deployment contract — applying the
  confidence/entropy routing rule to the ONNX outputs reproduces the torch
  inference decision + prediction.
- tests/test_integration.py: full pipelines (build → train → calibrate conf/
  entropy → analyze → benchmark → save/load → ONNX) across a factory backbone
  and a custom transformer, plus a CUDA run.

164 CPU + 5 GPU tests, 96% coverage.
@sohams25
sohams25 force-pushed the feat-onnx-and-integration branch from 3549157 to 079c39e Compare June 6, 2026 07:37

@sohams25 sohams25 left a comment

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

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

Review — ✅ ready to merge

CI green across 3.10 / 3.11 / 3.12. This run installs onnx via the dev extra, so the ONNX tests genuinely run (not skip), validating export on CI's torch.

/review + a 4-dimension adversarial review found 11 findings (2 HIGH) — both HIGH were real ship-blockers and are fixed:

  • dynamo=False kwarg only exists on torch>=2.4 (pyproject allows >=2.0): now applied only when the param is present (older torch defaults to the legacy exporter anyway), so export works across the whole supported torch range.
  • onnx was an undeclared hard runtime dep (legacy exporter loads it; tests guarded on onnxruntime, so a clean runner would have failed, not skipped): now an [onnx] extra with a clear ImportError, added to dev, and tests importorskip("onnx").

Also addressed: deployment-contract test (apply the routing rule to ONNX outputs → matches torch inference, confidence + entropy), deprecation-warning suppression, dynamic-batch token coverage, narrowed the static-reject exception, eval-before-adapter, and the CLI now echoes the traced input shape.

Tested in isolation and together: 164 CPU + 5 GPU, 96% coverage. Merging.

@sohams25
sohams25 merged commit 56cc825 into main Jun 6, 2026
3 checks passed
@sohams25
sohams25 deleted the feat-onnx-and-integration branch June 6, 2026 07:40
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.

1 participant