[5676209][ONNX][Autocast] Add support for single npz file with multiple samples (#815)

## 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:
```sh
$ 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*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
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](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes

## Additional Information
Equivalent support is already included in the quantization workflow:
https://github.com/NVIDIA/Model-Optimizer/blob/1cc8e6bf3917f61500e81d4ded0af5d5a00e2e25/modelopt/onnx/quantization/calib_utils.py#L50


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## 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.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>

Author: gcunhase

CHANGELOG.rst

@@ -13,6 +13,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 - Add standalone type inference option (``--use_standalone_type_inference``) in ONNX AutoCast as an alternative to ONNX's ``infer_shapes``. This experimental feature performs type-only inference without shape inference, useful as a workaround when shape inference fails or to avoid unnecessary shape inference overhead.
 - Add support for Kimi K2 Thinking model quantization from the original int4 checkpoint.
 - Add support for ``params`` constraint based automatic neural architecture search in Minitron pruning (``mcore_minitron``) as an alternative to manual pruning (using ``export_config``). See `examples/pruning/README.md <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/pruning>`_ for more details on its usage.
+- Add support for calibration data with multiple samples in ``npz`` format in the ONNX Autocast workflow.
 
 0.41 (2026-01-19)
 ^^^^^^^^^^^^^^^^^

modelopt/onnx/autocast/referencerunner.py

@@ -30,6 +30,7 @@
 import onnx
 
 from modelopt.onnx.autocast.logging_config import configure_logging, logger
+from modelopt.onnx.quantization.calib_utils import CalibrationDataProvider
 from modelopt.onnx.quantization.ort_utils import _prepare_ep_list
 
 configure_logging()
@@ -70,7 +71,13 @@ def _load_inputs_from_json(self, input_data_path):
 
     def _load_inputs_from_npz(self, input_data_path):
         """Load inputs from NPZ format."""
-        return [np.load(input_data_path)]
+        calib_data = np.load(input_data_path)
+
+        if isinstance(calib_data, np.lib.npyio.NpzFile):
+            # Wrap data into a CalibDataProvider to support a single NPZ file containing data from multiple batches
+            data_loader = {key: calib_data[key] for key in calib_data.files}
+            return CalibrationDataProvider(self.model, data_loader).calibration_data_list
+        return [calib_data]
 
     def _validate_inputs(self, data_loader):
         """Validate that input names and shapes match the model."""

modelopt/onnx/quantization/calib_utils.py

@@ -38,7 +38,7 @@ class CalibrationDataProvider(CalibrationDataReader):
 
     def __init__(
         self,
-        onnx_path: str,
+        onnx_path: str | onnx.ModelProto,
         calibration_data: CalibrationDataType,
         calibration_shapes: str | None = None,
     ):
@@ -58,7 +58,7 @@ def __init__(
         logger.info("Setting up CalibrationDataProvider for calibration")
         # Tensor data is not required to generate the calibration data
         # So even if the model has external data, we don't need to load them here
-        onnx_model = onnx.load(onnx_path)
+        onnx_model = onnx.load(onnx_path) if isinstance(onnx_path, str) else onnx_path
         input_names = get_input_names(onnx_model)
         input_shapes = {} if calibration_shapes is None else parse_shapes_spec(calibration_shapes)
         inferred_input_shapes = get_input_shapes(onnx_model)