diff --git a/CHANGELOG.rst b/CHANGELOG.rst index c45d8025386..56118d49fe9 100755 --- a/CHANGELOG.rst +++ b/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) diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py index 4db54b3b366..a70ffa74b33 100644 --- a/modelopt/onnx/autocast/referencerunner.py +++ b/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 diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 5c5c3c00ab2..fc4128c054f 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/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