From 35ddefab2e401c86eaf0998edc64317c46a55a1d Mon Sep 17 00:00:00 2001 From: Ram Sathyavageeswaran Date: Wed, 1 Jul 2026 21:35:16 +0000 Subject: [PATCH] fix: promote None-initialized kwarg to list in forward_capture When an optional kwarg (e.g. a sparse-attention mask) is absent on the first calibration sample, the initialization branch stores it as None. If a later sample delivers a real Tensor for the same key, the append branch called None.append() / None.extend(), raising: AttributeError: 'NoneType' object has no attribute 'extend' Fix: before calling .append()/.extend(), check whether the stored value is None and promote it to an empty list first. Fixes: #1950 Signed-off-by: Ram Sathyavageeswaran --- auto_round/calibration/hooks.py | 3 + .../core/test_forward_capture_none_kwarg.py | 137 ++++++++++++++++++ 2 files changed, 140 insertions(+) create mode 100644 test/test_cpu/core/test_forward_capture_none_kwarg.py diff --git a/auto_round/calibration/hooks.py b/auto_round/calibration/hooks.py index 7ebbf1fe0..84559b945 100644 --- a/auto_round/calibration/hooks.py +++ b/auto_round/calibration/hooks.py @@ -119,6 +119,9 @@ def forward_capture(m, hidden_states=None, *positional_inputs, **kwargs): state.inputs[name][key] = [stored, raw_new] continue new_data = to_device(new_data, device=torch.device("cpu")) + # Guard against None-initialized kwargs that arrive as tensors on later samples (#1950). + if state.inputs[name][key] is None: + state.inputs[name][key] = [] if state.quantizer.batch_size <= 1: state.inputs[name][key].append(new_data) else: diff --git a/test/test_cpu/core/test_forward_capture_none_kwarg.py b/test/test_cpu/core/test_forward_capture_none_kwarg.py new file mode 100644 index 000000000..005f99df5 --- /dev/null +++ b/test/test_cpu/core/test_forward_capture_none_kwarg.py @@ -0,0 +1,137 @@ +"""Regression tests for https://github.com/intel/auto-round/issues/1950. + +``forward_capture`` raised ``AttributeError: 'NoneType' object has no attribute +'extend'`` when an optional kwarg was ``None`` on the first calibration sample +but a real ``Tensor`` on a subsequent one. The fix promotes the stored ``None`` +to an empty list before appending. +""" + +from functools import partial +from types import SimpleNamespace + +import torch + +from auto_round.calibration.hooks import make_block_forward_func + + +def _make_state(batch_size=1): + """Return a minimal state stub accepted by ``make_block_forward_func``. + + Args: + batch_size (int): Number of samples per calibration batch. + + Returns: + SimpleNamespace: State object with ``inputs``, ``quantizer``, + ``model_context``, ``has_variable_block_shape``, + ``blocks_requiring_input_ids``, and ``_should_stop_cache_forward``. + """ + quantizer = SimpleNamespace(batch_size=batch_size, batch_dim=None) + model_context = SimpleNamespace( + shared_cache_keys=("position_ids", "cache_position", "position_embeddings", "cu_seqlens"), + ) + return SimpleNamespace( + inputs={}, + quantizer=quantizer, + model_context=model_context, + has_variable_block_shape=False, + blocks_requiring_input_ids=[], + _should_stop_cache_forward=lambda name: False, + ) + + +class _FakeModule(torch.nn.Module): + """Minimal decoder-block stub: passes ``hidden_states`` through unchanged.""" + + def __init__(self): + super().__init__() + + def orig_forward(self, hidden_states, **kwargs): + """Identity forward used as the underlying block implementation.""" + return hidden_states + + +def _attach_capture(state, name, module): + """Attach a ``forward_capture`` closure to *module* for block *name*. + + Mirrors the wiring done by ``replace_forward_with_hooks`` so tests + exercise the same code path as production calibration. + + Args: + state: Calibration state stub (from ``_make_state``). + name (str): Block name key used in ``state.inputs``. + module (_FakeModule): Module to instrument. + + Returns: + _FakeModule: The same module with ``forward`` replaced. + """ + fn = make_block_forward_func(state, name) + module.forward = partial(fn, module) + return module + + +def test_none_then_tensor_kwarg_batch_size_1(): + """``batch_size=1``: None-initialized kwarg must not crash on later Tensor sample. + + Sequence: + 1. First forward call — no ``optional_mask`` kwarg at all. + 2. State is mutated to simulate the None-initialization path + (kwarg was ``None`` on its first appearance). + 3. Second forward call delivers a real Tensor for ``optional_mask`` + — must not raise ``AttributeError``. + """ + state = _make_state(batch_size=1) + name = "decoder.layers.0" + module = _attach_capture(state, name, _FakeModule()) + + hidden = torch.randn(1, 4, 8) + module(hidden) + + # Simulate None stored by the initialization branch. + state.inputs[name]["optional_mask"] = None + + module(hidden, optional_mask=torch.ones(1, 4, 4)) + + stored = state.inputs[name].get("optional_mask") + assert isinstance(stored, list), f"expected list, got {type(stored)}" + assert len(stored) == 1 + + +def test_none_then_tensor_kwarg_batch_size_gt1(): + """``batch_size > 1``: None-initialized kwarg must not crash on later Tensor sample. + + The ``batch_size > 1`` path calls ``.extend()`` rather than ``.append()``, + exercising the second crash site from the original bug. + """ + state = _make_state(batch_size=2) + state.quantizer.batch_dim = 0 + name = "decoder.layers.0" + module = _attach_capture(state, name, _FakeModule()) + + hidden = torch.randn(2, 4, 8) + module(hidden) + + state.inputs[name]["optional_mask"] = None + + module(hidden, optional_mask=torch.ones(2, 4, 4)) + + stored = state.inputs[name].get("optional_mask") + assert isinstance(stored, list), f"expected list, got {type(stored)}" + # torch.split(..., 1, dim=0) on a batch-2 tensor yields 2 chunks. + assert len(stored) == 2 + + +def test_normal_tensor_kwarg_unaffected(): + """Kwargs that are always Tensors continue to accumulate correctly after the fix.""" + state = _make_state(batch_size=1) + name = "decoder.layers.0" + module = _attach_capture(state, name, _FakeModule()) + + hidden = torch.randn(1, 4, 8) + mask = torch.ones(1, 4, 4) + + module(hidden, attention_mask=mask) + module(hidden, attention_mask=mask) + + stored = state.inputs[name].get("attention_mask") + assert isinstance(stored, list) + assert len(stored) == 2