[Bug]: forward_capture AttributeError ('NoneType' has no attribute 'extend') during MoE calibration when a kwarg is None-then-Tensor

[pasta-paul]

Problem Description

When AutoRound calibrates a Mixture-of-Experts model (here GLM-5.2, arch GlmMoeDsaForCausalLM, 256 routed experts) driven via llm-compressor's AutoRoundModifier, calibration crashes during the input-capture forward with:

File ".../auto_round/utils/model.py", line 2199, in forward
    return base_hook(m, hidden_states, *positional_inputs, **kwargs)
File ".../auto_round/calibration/hooks.py", line 126, in forward_capture
    state.inputs[name][key].extend(
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'NoneType' object has no attribute 'extend'

Root cause

In auto_round/calibration/hooks.py forward_capture, the initialization branch stores a raw value for a key that is None or a shared-cache key (paraphrasing 0.13.0):

if key not in state.inputs[name].keys():            # initialization
    data = to_device(kwargs[key], device=torch.device("cpu"))
    if data is None or key in state.model_context.shared_cache_keys:
        state.inputs[name][key] = data              # <-- stores raw None for a NON-shared key
        continue
    ...

Later, on a subsequent sample where the same key arrives as a Tensor, the append branch runs and does:

else:  # append cache inputs
    ...
    if state.quantizer.batch_size <= 1:
        state.inputs[name][key].append(new_data)
    else:
        state.inputs[name][key].extend(            # line 126 -> AttributeError: None has no .extend
            list(torch.split(new_data, 1, dim=state.quantizer.batch_dim)))

So the failure occurs whenever a forward kwarg is None on the first calibration sample but a Tensor on a later sample for a key that is not in shared_cache_keys. The init path stored a scalar None (not a per-sample list), and the append/extend path assumes a list. This is common for MoE / sparse-attention models where optional kwargs (e.g. a mask or routing-related tensor) are None for some sequences and present for others.

Suggested fix

Treat a non-shared None key as per-sample list storage at init, OR guard the append branch. Minimal guard at the append site:

else:  # append cache inputs
    ...
    if not isinstance(state.inputs[name].get(key), list):
        # key was first seen as None (non-shared); promote to per-sample list
        prev = state.inputs[name].get(key)
        state.inputs[name][key] = [] if prev is None else [prev]
    if state.quantizer.batch_size <= 1:
        state.inputs[name][key].append(new_data)
    else:
        state.inputs[name][key].extend(list(torch.split(new_data, 1, dim=state.quantizer.batch_dim)))

(Or, at init, store [data] instead of raw data when data is None and key not in shared_cache_keys.)

Environment

• auto-round 0.13.0
• llm-compressor 0.12.0 (AutoRoundModifier(targets=["Linear"], scheme="W4A16", iters=200), moe_calibrate_all_experts=True)
• transformers 5.10.1, torch 2.12.0, CUDA 12.8
• Model: GLM-5.2 (GlmMoeDsaForCausalLM, 256 routed experts), bf16 source

Notes

Happy to open a PR with the fix if the maintainers prefer one of the two approaches above. The crash reproduces deterministically at the first MoE layer's calibration; it is independent of bit-width (seen at W4A16).

[yiliu30]

Thanks @pasta-paul , Good point. Assign it to you :)
In many models it is stable, but this path can still be Optional[Tensor] and vary across calibration samples.
So for safety we should handle None -> Tensor transitions. The key detail is preserving per-sample alignment:

• if existing value is raw None, normalize to [None] (not []) before append/extend
• [] drops the earlier sample slot and can misalign later indexing/replay

[ramkrishs]

The fix for this is available in #1979. It promotes a None-initialized kwarg to an empty list before calling .append() / .extend(), covering both batch_size <= 1 and batch_size > 1 paths.