fix(autocast): stream calibration batches instead of materializing all of them

CHANGELOG.rst

@@ -64,6 +64,7 @@ Changelog
 - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged.
 - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions.
 - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation.
+- Fix ONNX AutoCast exhausting host memory on large calibration sets: statistics are now aggregated as each sample runs, so peak memory no longer grows with the number of samples in ``calibration_data``.
 - Fix ``fsdp2_aware_weight_update`` masking an error raised during its unshard (e.g. a CUDA OOM) with a spurious ``UnboundLocalError``.
 
 0.48.0 (2026-10-xx)

modelopt/onnx/autocast/referencerunner.py

@@ -69,6 +69,32 @@ def size(self):
         return result
 
 
+class _FeedDictRecorder:
+    """Wraps a data loader, remembering the feed dict of the batch currently in flight.
+
+    Comparator's streaming mode yields outputs only, so the graph inputs for each batch
+    have to be captured as the loader is consumed.
+    """
+
+    def __init__(self, data_loader):
+        self._data_loader = data_loader
+        self.last_feed_dict = {}
+
+    @property
+    def input_metadata(self):
+        return self._data_loader.input_metadata
+
+    @input_metadata.setter
+    def input_metadata(self, value):
+        # Plain lists have no such attribute; Comparator tolerates the AttributeError.
+        self._data_loader.input_metadata = value
+
+    def __iter__(self):
+        for feed_dict in self._data_loader:
+            self.last_feed_dict = feed_dict
+            yield feed_dict
+
+
 class ReferenceRunner:
     """A class to run ONNX models with ONNXRuntime for reference inference."""
 
@@ -217,54 +243,38 @@ def _get_ort_runner(self, model):
             raise
         return runners, model_temp_dir
 
-    def _aggregate_tensor_stats(self, all_batch_data: list[OrderedDict]) -> OrderedDict:
-        """Aggregate tensor statistics across multiple batches.
-
-        Args:
-            all_batch_data: List of dictionaries containing tensor data for each batch.
-
-        Returns:
-            OrderedDict mapping tensor names to TensorStats objects.
+    @staticmethod
+    def _batch_stats(data: np.ndarray) -> tuple[float, float, float]:
+        """Returns (absmax, min, max) of one tensor, treating an empty tensor as zeros."""
+        if data.size == 0:
+            return 0, 0, 0
+        return np.max(np.abs(data)), np.min(data), np.max(data)
+
+    def _init_tensor_stats(self, batch_data: OrderedDict) -> OrderedDict:
+        """Creates per-tensor statistics seeded from +/-inf sentinels.
+
+        Seeding from the sentinels and always folding the first batch through
+        _fold_tensor_stats (rather than taking its values directly) keeps a NaN in
+        the first batch from poisoning the running stat: max(-inf, nan) stays -inf,
+        the same way max(running_value, nan) already ignores NaN in every later
+        batch.
         """
-        if len(all_batch_data) == 1:
-            # Single batch - return raw data for backward compatibility
-            return all_batch_data[0]
-
-        logger.info(f"Aggregating statistics across {len(all_batch_data)} batches...")
-
-        aggregated = OrderedDict()
-        tensor_names = all_batch_data[0].keys()
-
-        for name in tensor_names:
-            absmax = -np.inf
-            min_val = np.inf
-            max_val = -np.inf
-            shape = None
-
-            for batch_data in all_batch_data:
-                if name not in batch_data:
-                    continue
-                data = batch_data[name]
-                if shape is None:
-                    shape = data.shape
-
-                batch_absmax = np.max(np.abs(data)) if data.size > 0 else 0
-                batch_min = np.min(data) if data.size > 0 else 0
-                batch_max = np.max(data) if data.size > 0 else 0
-
-                absmax = max(absmax, batch_absmax)
-                min_val = min(min_val, batch_min)
-                max_val = max(max_val, batch_max)
-
-            if shape is not None:
-                aggregated[name] = TensorStats(
-                    absmax=absmax,
-                    min_val=min_val,
-                    max_val=max_val,
-                    shape=shape,
-                )
-
-        return aggregated
+        stats = OrderedDict()
+        for name, data in batch_data.items():
+            stats[name] = TensorStats(
+                absmax=-np.inf, min_val=np.inf, max_val=-np.inf, shape=data.shape
+            )
+        return stats
+
+    def _fold_tensor_stats(self, stats: OrderedDict, batch_data: OrderedDict) -> None:
+        """Folds one batch into the running statistics in place."""
+        for name, tensor_stats in stats.items():
+            if name not in batch_data:
+                continue
+            absmax, min_val, max_val = self._batch_stats(batch_data[name])
+            tensor_stats.absmax = max(tensor_stats.absmax, absmax)
+            tensor_stats.min_val = min(tensor_stats.min_val, min_val)
+            tensor_stats.max_val = max(tensor_stats.max_val, max_val)
 
     def run(self, inputs=None):
         """Run FP32 inference with provided or random inputs.
@@ -298,17 +308,40 @@ def run(self, inputs=None):
 
         # Load the modified model and create an inference session
         runners, model_temp_dir = self._get_ort_runner(modify_outputs)
+        num_batches = 0
+        first_batch = None
+        stats = None
         try:
             # Comparator is used despite the fact that we are using ONNXRuntime
             # because it provides the ability to generate random inputs using DataLoader
-            data_loader = self._load_inputs(inputs)
+            data_loader = _FeedDictRecorder(self._load_inputs(inputs))
 
             # Temporarily redirect stdout to suppress Comparator.run() output
             stdout = sys.stdout
             string_buffer = io.StringIO()
             sys.stdout = string_buffer
             try:
-                results = Comparator.run(runners, data_loader=data_loader)
+                # Stream one batch at a time and fold it into the running statistics
+                # immediately. Under MARK_ALL every intermediate tensor is a graph output,
+                # so materializing every batch would make peak memory scale with the
+                # number of calibration samples.
+                for result in Comparator.run(runners, data_loader=data_loader, streaming=True):
+                    batch_data = OrderedDict()
+                    batch_data.update(data_loader.last_feed_dict)
+                    batch_data.update(result[0][1][0])
+
+                    num_batches += 1
+                    if num_batches == 1:
+                        stats = self._init_tensor_stats(batch_data)
+                        first_batch = batch_data
+                    else:
+                        # Runner output buffers may be reused across iterations, so the
+                        # first batch's raw arrays are only safe to return while it is
+                        # the only batch.
+                        first_batch = None
+                    # Fold every batch, including the first, through the same sentinel
+                    # comparison so a NaN batch can never poison the running statistic.
+                    self._fold_tensor_stats(stats, batch_data)
             finally:
                 # Capture the output before restoring stdout
                 captured_output = string_buffer.getvalue()
@@ -317,34 +350,13 @@ def run(self, inputs=None):
             if model_temp_dir is not None:
                 model_temp_dir.cleanup()
 
-        if not results:
+        if num_batches == 0:
             logger.error(f"ONNXRuntime execution failed with output:\n{captured_output}")
             raise Exception("ONNXRuntime failed to run, see logs for details")
 
-        # Collect all batch data (inputs + outputs)
-        all_batch_data = []
-        runner_results = results[0][1]  # Get all iteration results for the first runner
-        data_loader_iter = iter(data_loader)
-
-        for iter_idx, iter_result in enumerate(runner_results):
-            output_dict = OrderedDict(iter_result)
+        if num_batches == 1:
+            # Single batch - return raw data for backward compatibility
+            return first_batch
 
-            # Get corresponding input data
-            try:
-                input_data = next(data_loader_iter)
-            except StopIteration:
-                # If data_loader is exhausted, it might be a DataLoader that generates random data
-                input_data = {}
-
-            # Combine inputs and outputs for this batch
-            batch_dict = OrderedDict()
-            batch_dict.update(input_data)
-            batch_dict.update(output_dict)
-            all_batch_data.append(batch_dict)
-
-        num_batches = len(all_batch_data)
-        if num_batches > 1:
-            logger.info(f"Processed {num_batches} batches of calibration data")
-
-        # Aggregate statistics across all batches
-        return self._aggregate_tensor_stats(all_batch_data)
+        logger.info(f"Aggregated statistics across {num_batches} batches of calibration data")
+        return stats

tests/unit/onnx/autocast/test_referencerunner.py

@@ -13,9 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import gc
 import json
 import os
 import tempfile
+import weakref
 from collections import OrderedDict
 from unittest.mock import Mock
 
@@ -439,3 +441,90 @@ def test_multi_batch_aggregation_statistics(reference_runner):
         assert results["Y2"].absmax == 70.0
         assert results["Y2"].min_val == -50.0
         assert results["Y2"].max_val == 70.0
+
+
+def test_multi_batch_does_not_accumulate_batches(reference_runner, monkeypatch):
+    """Batches are folded as they stream in, so peak memory is independent of batch count."""
+    batch_refs, array_refs = [], []
+    concurrent, array_concurrent = [], []
+
+    def spy(method):
+        def wrapped(self, *args):
+            batch_data = args[-1]
+            # Track an output activation alongside the mapping: marking every tensor as an
+            # output is what makes a batch big, and the arrays could outlive the mapping
+            # that carried them. batch_data holds the input feed first, so index by name.
+            batch_refs.append(weakref.ref(batch_data))
+            array_refs.append(weakref.ref(batch_data["Y1"]))
+            gc.collect()
+            concurrent.append(sum(ref() is not None for ref in batch_refs))
+            array_concurrent.append(sum(ref() is not None for ref in array_refs))
+            return method(self, *args)
+
+        return wrapped
+
+    monkeypatch.setattr(
+        ReferenceRunner, "_init_tensor_stats", spy(ReferenceRunner._init_tensor_stats)
+    )
+    monkeypatch.setattr(
+        ReferenceRunner, "_fold_tensor_stats", spy(ReferenceRunner._fold_tensor_stats)
+    )
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        for i in range(5):
+            np.savez(
+                os.path.join(temp_dir, f"batch_{i:03d}.npz"),
+                X1=np.full((1, 3), i, dtype=np.float32),
+                X2=np.full((1, 3), i + 1, dtype=np.float32),
+            )
+        results = reference_runner.run(temp_dir)
+
+    assert len(concurrent) == 5, "every batch should reach the aggregator"
+    assert max(concurrent) == 1, f"batches were retained instead of streamed: {concurrent}"
+    assert max(array_concurrent) == 1, (
+        f"batch activations were retained instead of streamed: {array_concurrent}"
+    )
+    # The streaming machinery is only worth testing if it still produces the right
+    # numbers: X1=0..4, X2=1..5, Y1=X1+X2, Y2=X1*X2.
+    assert results["X1"].absmax == 4.0
+    assert results["X1"].min_val == 0.0
+    assert results["X1"].max_val == 4.0
+    assert results["X2"].absmax == 5.0
+    assert results["X2"].min_val == 1.0
+    assert results["X2"].max_val == 5.0
+    assert results["Y1"].absmax == 9.0
+    assert results["Y1"].min_val == 1.0
+    assert results["Y1"].max_val == 9.0
+    assert results["Y2"].absmax == 20.0
+    assert results["Y2"].min_val == 0.0
+    assert results["Y2"].max_val == 20.0
+
+
+def test_multi_batch_nan_first_batch_does_not_poison_stats(reference_runner):
+    """A NaN in the first batch must not lock a tensor's running stats at NaN.
+
+    Seeding directly from the first batch (instead of sentinels folded through
+    _fold_tensor_stats) breaks this: max(nan, 10.0) is nan because Python's max
+    never replaces its first argument with a comparand that loses a NaN
+    comparison, so a later finite batch could never clear the contamination.
+    """
+    with tempfile.TemporaryDirectory() as temp_dir:
+        # Batch 1: Y1 = X1 + X2 contains a NaN.
+        np.savez(
+            os.path.join(temp_dir, "batch_001.npz"),
+            X1=np.array([[np.nan, 0.0, 0.0]], dtype=np.float32),
+            X2=np.array([[0.0, 0.0, 0.0]], dtype=np.float32),
+        )
+        # Batch 2: all finite, with a clear absmax/min/max to assert on.
+        np.savez(
+            os.path.join(temp_dir, "batch_002.npz"),
+            X1=np.array([[10.0, -3.0, 0.0]], dtype=np.float32),
+            X2=np.array([[0.0, 0.0, 0.0]], dtype=np.float32),
+        )
+
+        results = reference_runner.run(temp_dir)
+
+        assert not np.isnan(results["Y1"].absmax)
+        assert results["Y1"].absmax == 10.0
+        assert results["Y1"].min_val == -3.0
+        assert results["Y1"].max_val == 10.0