[5676209][ONNX][Autocast] Add support for single npz file with multiple samples

[gcunhase]

What does this PR do?

Type of change: New feature

Overview: Currently, Autocast only supports calibration data with shape matching the model's input. This PR adds support for calibration data with shape that is a multiple of the model's input. It does so by re-arranging the data as such that it contains multiple samples with shape matching the model's input.

Simplified example:

• ONNX input: [1, 3, 224, 224]
• Calibration data: [10, 3, 224, 224]
• Calibration data with multiple samples: [1, 3, 224, 224] * 10

Usage

Single npz file with multiple samples:

$ python -m modelopt.onnx.autocast --onnx_path=$MODEL_NAME.onnx --calibration_data=calib_data_10.npz

Testing

See bug 5676209.

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes
• Did you write any new necessary tests?: No
• Did you add or update any necessary documentation?: No
• Did you update Changelog?: Yes

Additional Information

Equivalent support is already included in the quantization workflow:

Model-Optimizer/modelopt/onnx/quantization/calib_utils.py

Line 50 in 1cc8e6b

Ex. If a model has input shapes like {"sample": (2, 4, 64, 64), "timestep": (1,),

Summary by CodeRabbit

• New Features
  • CalibrationDataProvider now accepts both file paths and pre-loaded ONNX models as input.
  • NPZ calibration files can now provide multiple batches for calibration workflows.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

📝 Walkthrough

Walkthrough

These changes modify calibration data handling in ONNX quantization. CalibrationDataProvider now accepts pre-loaded ONNX ModelProto objects in addition to file paths, and the reference runner wraps NPZ calibration data within CalibrationDataProvider to enable batch handling from single files.

Changes

Cohort / File(s)	Summary
NPZ Data Wrapping modelopt/onnx/autocast/referencerunner.py	Added CalibrationDataProvider import; modified _load_inputs_from_npz to wrap NPZ data in CalibrationDataProvider and return calibration_data_list instead of single array for multi-batch support.
Model Input Flexibility modelopt/onnx/quantization/calib_utils.py	Widened CalibrationDataProvider.init parameter type from str to `str

Sequence Diagram

sequenceDiagram
    participant RR as Reference Runner
    participant CDP as CalibrationDataProvider
    participant ONNX as ONNX Model

    rect rgba(100, 149, 237, 0.5)
    Note over RR,ONNX: New Flow: NPZ Data with CalibrationDataProvider
    RR->>RR: Load NPZ file
    RR->>CDP: Wrap NPZ data + model path
    CDP->>ONNX: Load model from path (if string)
    ONNX-->>CDP: ModelProto returned
    CDP-->>RR: CalibrationDataProvider with batches
    end

    rect rgba(144, 238, 144, 0.5)
    Note over RR,ONNX: Alternative: Pre-loaded Model
    RR->>RR: Load NPZ file + ModelProto
    RR->>CDP: Wrap NPZ data + ModelProto
    CDP->>CDP: Use ModelProto directly (skip loading)
    CDP-->>RR: CalibrationDataProvider with batches
    end

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: adding support for single NPZ files with multiple samples in ONNX Autocast.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.29%. Comparing base (4f4558a) to head (a496175).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #815      +/-   ##
==========================================
+ Coverage   74.17%   74.29%   +0.12%     
==========================================
  Files         192      192              
  Lines       19246    19251       +5     
==========================================
+ Hits        14276    14303      +27     
+ Misses       4970     4948      -22     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.