From 929b428314d19d95888cf1c4201706132907d179 Mon Sep 17 00:00:00 2001 From: SID <99672439+SID-6921@users.noreply.github.com> Date: Sat, 26 Sep 2026 18:45:41 -0400 Subject: [PATCH 1/6] fix(autocast): stream calibration batches instead of materializing all of them ReferenceRunner.run() wraps the model with MARK_ALL, so every intermediate tensor becomes a graph output. It then called Comparator.run(), which executes every calibration sample up front and retains one full-graph activation dump per sample; those results were copied a second time into all_batch_data before aggregation. Peak memory therefore grew linearly with the number of calibration samples, making real (non-random) calibration data unusable past a few dozen samples. Use Comparator's streaming mode and fold each batch into the running per-tensor statistics as it arrives, so peak memory is independent of the sample count. Aggregated absmax/min/max values and the single-batch raw-array return are unchanged. Measured on a model with ~14 MiB of activations per sample: samples before after 1 8.2 MiB 8.1 MiB 8 128.4 MiB 24.2 MiB 32 448.5 MiB 32.3 MiB 64 897.4 MiB 23.5 MiB Fixes #2337 Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 1 + modelopt/onnx/autocast/referencerunner.py | 158 +++++++++--------- .../onnx/autocast/test_referencerunner.py | 36 ++++ 3 files changed, 118 insertions(+), 77 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4f621df5dea..4134d82f2ed 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -94,6 +94,7 @@ Changelog - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected. - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. +- Fix ONNX AutoCast reference inference retaining every calibration sample's activations at once. The reference run marks every intermediate tensor as a graph output, so holding one full-graph activation dump per sample made peak memory scale with the number of samples in ``calibration_data`` and could exhaust host memory well below the 512 samples ``CalibrationDataProvider`` documents. Samples are now streamed and folded into the running per-tensor statistics as they arrive, so peak memory is independent of the sample count. The aggregated statistics are unchanged. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py index 4db54b3b366..46d346cbc86 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,32 @@ 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. - """ - 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 + @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: + """Seeds per-tensor statistics from the first batch.""" + stats = OrderedDict() + for name, data in batch_data.items(): + absmax, min_val, max_val = self._batch_stats(data) + stats[name] = TensorStats( + absmax=absmax, min_val=min_val, max_val=max_val, 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 +302,38 @@ 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 + self._fold_tensor_stats(stats, batch_data) finally: # Capture the output before restoring stdout captured_output = string_buffer.getvalue() @@ -317,34 +342,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..9cd46996c31 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,37 @@ 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 = [] + concurrent = [] + + def spy(method): + def wrapped(self, *args): + batch_refs.append(weakref.ref(args[-1])) + gc.collect() + concurrent.append(sum(ref() is not None for ref in batch_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), + ) + 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}" From 532e8aff3dc00a01e624f080d52f79d5286349ff Mon Sep 17 00:00:00 2001 From: SID <99672439+SID-6921@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:13:53 -0400 Subject: [PATCH 2/6] Address review: trim the changelog entry and track activations in the test The changelog entry carried root-cause detail that AGENTS.md keeps out of CHANGELOG.rst; it is now one sentence describing the user-visible effect. The streaming test tracked only the per-batch mapping. It now also weak-refs one activation array from each batch, since the arrays are the memory and could outlive the mapping that carried them. Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 2 +- tests/unit/onnx/autocast/test_referencerunner.py | 14 +++++++++++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4134d82f2ed..a55bbf28553 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -94,7 +94,7 @@ Changelog - Fix ``megatron_generate`` dropping the VLM vision inputs (``pixel_values`` / ``image_grid_thw`` / ``image_sizes``) after the first generated token when KV-cache decoding is off, including the automatic fallback under sequence parallelism, which made generation silently ignore the image. No other ModelOpt feature is affected. - Fix two issues in the vLLM offline hidden-state dump (``examples/speculative_decoding/collect_hidden_states/compute_hidden_states_vllm.py``) that only surface on large runs. **Resume:** the filter that skips conversations whose ``.pt`` already exists now runs with ``load_from_cache_file=False``. It depends on on-disk state, which is not part of the fingerprint ``datasets`` computes from the function and the dataset, so with a persistent HF cache reused across a resumed or requeued run the cached "keep everything" result from an earlier run was replayed and the dump re-generated and overwrote conversations it had already finished (observed: tens of thousands of ``.pt`` rewritten while the output count stayed flat). **Staging:** generation is now chunked (``--save-chunk-size``, default 256), so each chunk is saved and its staged hidden states freed before the next chunk is generated. Previously the whole dataset was generated before anything was saved, which kept every conversation staged in the connector's ``shared_storage_path`` (``/dev/shm``, i.e. RAM, by default) at once and exhausted it partway through large dumps. Chunking also makes the dump incrementally durable, so an interrupted run keeps its finished conversations and resumes from them. The save path now also frees each conversation's staged hidden states in a ``finally``, so a conversation skipped mid-loop (e.g. a short ``loss_mask``) can no longer leak its staging file, and conversation ids are validated as plain filenames before being used to build output paths. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. -- Fix ONNX AutoCast reference inference retaining every calibration sample's activations at once. The reference run marks every intermediate tensor as a graph output, so holding one full-graph activation dump per sample made peak memory scale with the number of samples in ``calibration_data`` and could exhaust host memory well below the 512 samples ``CalibrationDataProvider`` documents. Samples are now streamed and folded into the running per-tensor statistics as they arrive, so peak memory is independent of the sample count. The aggregated statistics are unchanged. +- 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``. 0.47.0 (2026-09-23) ^^^^^^^^^^^^^^^^^^^ diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 9cd46996c31..26dd8841138 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -445,14 +445,19 @@ def test_multi_batch_aggregation_statistics(reference_runner): 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 = [] - concurrent = [] + batch_refs, array_refs = [], [] + concurrent, array_concurrent = [], [] def spy(method): def wrapped(self, *args): - batch_refs.append(weakref.ref(args[-1])) + batch_data = args[-1] + # Track an activation alongside the mapping: the arrays are the memory, and + # they could outlive the mapping that carried them. + batch_refs.append(weakref.ref(batch_data)) + array_refs.append(weakref.ref(next(iter(batch_data.values())))) 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 @@ -475,3 +480,6 @@ def wrapped(self, *args): 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}" + ) From e5b05c2c6a4615d18dd6868ae2f4ec91b206d584 Mon Sep 17 00:00:00 2001 From: SID <99672439+SID-6921@users.noreply.github.com> Date: Sat, 26 Sep 2026 19:24:40 -0400 Subject: [PATCH 3/6] Track an output activation, not the input feed, in the streaming test batch_data holds the input feed before the Comparator outputs, so the first value was X1 rather than the Y1 activation. Marking every tensor as an output is what makes a batch big, so the outputs are what the assertion needs to watch: retaining only the output arrays passed the previous check and fails this one. Signed-off-by: SID <99672439+SID-6921@users.noreply.github.com> --- tests/unit/onnx/autocast/test_referencerunner.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 26dd8841138..0c561c14a46 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -451,10 +451,11 @@ def test_multi_batch_does_not_accumulate_batches(reference_runner, monkeypatch): def spy(method): def wrapped(self, *args): batch_data = args[-1] - # Track an activation alongside the mapping: the arrays are the memory, and - # they could outlive the mapping that carried them. + # 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(next(iter(batch_data.values())))) + 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)) From 50c5a005e97bb863e9b1e238648a5dc1f9f30489 Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:58:26 -0400 Subject: [PATCH 4/6] Assert aggregated stats in the streaming regression test The streaming test exercised the memory-liveness machinery but discarded ReferenceRunner.run()'s actual return value, so a broken fold (e.g. one that never updates the running stats) would pass silently as long as batches stayed unretained. Store the result and assert the known X1/X2/Y1/Y2 extrema; verified these fail against a no-op _fold_tensor_stats and pass against the real one. Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- tests/unit/onnx/autocast/test_referencerunner.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 0c561c14a46..12473e81afd 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -477,10 +477,24 @@ def wrapped(self, *args): X1=np.full((1, 3), i, dtype=np.float32), X2=np.full((1, 3), i + 1, dtype=np.float32), ) - reference_runner.run(temp_dir) + 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 From 3c2449b05cee60cd878cc0f104f06b0c01f88a6f Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Mon, 26 Oct 2026 15:43:41 -0400 Subject: [PATCH 5/6] Move CHANGELOG entry to the 0.49.0 section Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2e65384f1ef..bf4f939f293 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -60,6 +60,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``. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ @@ -254,7 +255,6 @@ Changelog - Fix YAML config I/O decoding with the locale codepage instead of UTF-8, which made a config containing any non-ASCII byte fail to load on a machine whose locale is not UTF-8 (notably Windows, where the default is cp1252). ``modelopt/recipe/loader.py``, the two ONNX autotune state files, the two transformers config readers, the distill config and the puzzletron profile now pass ``encoding="utf-8"`` explicitly. Only the YAML config paths are covered: these are the files most likely to carry non-ASCII text in comments, model names or paths, and the only ones read inside a user's process. - Hybrid (e.g. Nemotron-H) checkpoints saved by the ``examples/megatron_bridge`` scripts now record their layer spec in ``run_config.yaml`` in a form that reloads, so they can be converted to HuggingFace; a checkpoint saved by an earlier release still needs its ``model.hybrid_stack_spec`` block replaced by hand. - Fix ONNX export of quantized models failing with ``Unsupported op_type for real weight quantization: Transpose`` when a weight reaches its consumer transposed, as in MaxViT's attention blocks. The same export now also quantizes each weight along the axis its Q/DQ pair declares, correcting silently wrong ``dq_only`` output when that axis differed from the consumer's default. -- 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 the ``evaluation`` agent skill's LiveCodeBench recipe scoring correct solutions that use binary standard I/O (``sys.stdin.buffer`` / ``sys.stdout.buffer``) as failures: the recipe now pins a fixed LiveCodeBench judge in the NeMo-Skills container. LiveCodeBench scores change as a result, so re-run earlier baselines before comparing. - Fix GPTQ with Megatron data parallelism: each data-parallel replica of a module now builds its Hessian from the whole calibration set instead of only its own shard, so all replicas reach the same quantized weights. - Fix ``examples/llm_sparsity/weight_sparsity/finetune.py`` saving a sharded state dict under FSDP2, the default FSDP version in Transformers 5, which made ``export_hf_ckpt.py`` fail to restore it. Re-run finetuning to produce a loadable ``finetuned_modelopt_state.pth``. From 3ed75ee21bfd0f492a9077db5e3c9ff7255781b0 Mon Sep 17 00:00:00 2001 From: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> Date: Wed, 28 Oct 2026 12:42:53 -0400 Subject: [PATCH 6/6] Fix NaN in first calibration batch poisoning streamed stats permanently meenchen flagged on review: seeding stats directly from the first batch (instead of folding it through the same sentinel comparison as every later batch) means a NaN in that first batch sets absmax/min_val/max_val to NaN, and no later finite batch can ever clear it. Python's max(acc, x) only replaces acc when x > acc, and any comparison against NaN is False, so once acc is NaN it stays NaN regardless of what streams in afterward. The pre-streaming code never had this problem: it always seeded from -inf/inf/-inf sentinels and folded every batch, including the first, through max()/min(). A NaN batch_absmax loses its comparison against the sentinel (or any real running value) and is silently skipped, exactly like it is for every batch after the first. Fix: _init_tensor_stats now only seeds the sentinels (no longer reads the first batch's values), and the batch loop folds every batch, including the first, through _fold_tensor_stats. Added test_multi_batch_nan_first_batch_does_not_poison_stats, reproducing meenchen's exact scenario (NaN first batch, then an activation of 10.0) and asserting the aggregated stat is finite. Verified with a standalone simulation of old/buggy/fixed aggregation against meenchen's exact numbers (NaN batch then 70000): old and fixed both report 70000.0, the buggy version reports NaN. Could not run the real test suite in this environment (local onnx install has a package-metadata/compiled-extension version mismatch unrelated to this change; TensorProto.FLOAT4E2M1 import fails even after reinstalling onnx 1.23.2). Signed-off-by: Siddhardha Nanda <99672439+SID-6921@users.noreply.github.com> --- modelopt/onnx/autocast/referencerunner.py | 16 +++++++--- .../onnx/autocast/test_referencerunner.py | 30 +++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/modelopt/onnx/autocast/referencerunner.py b/modelopt/onnx/autocast/referencerunner.py index 46d346cbc86..a70ffa74b33 100644 --- a/modelopt/onnx/autocast/referencerunner.py +++ b/modelopt/onnx/autocast/referencerunner.py @@ -251,12 +251,18 @@ def _batch_stats(data: np.ndarray) -> tuple[float, float, float]: return np.max(np.abs(data)), np.min(data), np.max(data) def _init_tensor_stats(self, batch_data: OrderedDict) -> OrderedDict: - """Seeds per-tensor statistics from the first batch.""" + """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. + """ stats = OrderedDict() for name, data in batch_data.items(): - absmax, min_val, max_val = self._batch_stats(data) stats[name] = TensorStats( - absmax=absmax, min_val=min_val, max_val=max_val, shape=data.shape + absmax=-np.inf, min_val=np.inf, max_val=-np.inf, shape=data.shape ) return stats @@ -333,7 +339,9 @@ def run(self, inputs=None): # first batch's raw arrays are only safe to return while it is # the only batch. first_batch = None - self._fold_tensor_stats(stats, batch_data) + # 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() diff --git a/tests/unit/onnx/autocast/test_referencerunner.py b/tests/unit/onnx/autocast/test_referencerunner.py index 12473e81afd..fc4128c054f 100644 --- a/tests/unit/onnx/autocast/test_referencerunner.py +++ b/tests/unit/onnx/autocast/test_referencerunner.py @@ -498,3 +498,33 @@ def wrapped(self, *args): 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