From 7e195c1f1d248a5f83270b43ef8dd5e3994b1394 Mon Sep 17 00:00:00 2001 From: xesdiny Date: Wed, 12 Aug 2026 14:30:30 +0800 Subject: [PATCH 1/4] fix(autoround): prevent input_capture_hook from accumulating during optimization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AutoRoundModifier registers `input_capture_hook` as a forward_pre hook on all decoding layers, but never disables it before the SignSGD optimization loop. During optimization, BlockForwardRunner.forward() calls decoding_layer.forward() on every mini-batch, triggering the hook each time. With gradient_accumulate_steps=8 this appends 8 × hidden_states tensors per outer iteration into _all_module_input; at ITERS=200 this accumulates ~200 × gradient_accumulate_steps × hidden_states_size bytes of GPU memory, causing OOM on models with large hidden states. Fix: introduce `_consumed_layers` (a PrivateAttr set) to track which layers have had their inputs consumed by apply_autoround(). The hook uses setdefault() for calibration-phase initialization but exits immediately for consumed layers, preventing any accumulation during optimization. apply_autoround() marks the layer consumed before pop()ing its inputs so that concurrent hook firings also no-op. on_calibration_end() clears the set for clean re-use. Verified: GPU allocated memory growth drops from +204 MB/iter to +0 MB/iter across 20 iterations on an 8-GPU setup with gradient_accumulate_steps=8. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Xesdiny --- src/llmcompressor/modifiers/autoround/base.py | 21 +++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/src/llmcompressor/modifiers/autoround/base.py b/src/llmcompressor/modifiers/autoround/base.py index 33b950aa1f..65c87e4d46 100644 --- a/src/llmcompressor/modifiers/autoround/base.py +++ b/src/llmcompressor/modifiers/autoround/base.py @@ -177,6 +177,7 @@ class AutoRoundModifier(Modifier, QuantizationMixin): # private variables _all_module_input: dict[str, list[tuple]] = PrivateAttr(default_factory=dict) _q_input: torch.Tensor | None = PrivateAttr(default=None) + _consumed_layers: set = PrivateAttr(default_factory=set) def on_initialize(self, state: State, **kwargs) -> bool: """ @@ -225,9 +226,17 @@ def start_calibration(self, model: torch.nn.Module): model.apply(enable_quantization) # quantize at the same time as calibrate def input_capture_hook(self, module, args, kwargs): - if module._tmp_name not in self._all_module_input: - self._all_module_input[module._tmp_name] = [] - self._all_module_input[module._tmp_name].append((args, kwargs)) + name = module._tmp_name + # After apply_autoround consumes a layer's inputs, it is added to + # _consumed_layers so this hook becomes a no-op for subsequent forward + # calls during the SignSGD optimization loop. Without this guard the + # hook fires on every mini-batch forward (gradient_accumulate_steps × + # iters calls) and holds the hidden_states tensor alive, causing + # ~(gradient_accumulate_steps × hidden_states_size) MB of GPU memory + # growth per outer iteration. + if name in self._consumed_layers: + return + self._all_module_input.setdefault(name, []).append((args, kwargs)) def on_calibration_start(self, state: State, event: Event, **kwargs): # register quantization calibration hooks @@ -313,7 +322,10 @@ def apply_autoround(self, state, modules): ar.batch_dim = 0 first_param = next(decoding_layer.parameters()) device = first_param.device - cur_inputs = self._all_module_input[decoding_layer._tmp_name] + # Mark consumed before pop so the hook immediately becomes a no-op + # for any forward call triggered during the optimization loop. + self._consumed_layers.add(decoding_layer._tmp_name) + cur_inputs = self._all_module_input.pop(decoding_layer._tmp_name) self._set_attention_masks(ar, decoding_layer, cur_inputs) decoding_layer.tuning_device = device # Only hand device placement to AutoRound when the caller explicitly @@ -359,6 +371,7 @@ def on_calibration_end(self, state: State, event: Event, **kwargs): self._remove_temporary_names(state.model) self.remove_hooks() self._q_input = None + self._consumed_layers.clear() def get_unquantized_layer_names(self, wrapped_model: torch.nn.Module) -> list[str]: unquantized_layers = [] From a5b6c5b23ce93d431f6c66cfff365b44eba24a61 Mon Sep 17 00:00:00 2001 From: xesdiny Date: Wed, 12 Aug 2026 15:15:52 +0800 Subject: [PATCH 2/4] fix(autoround): use set[str] type annotation for _consumed_layers Addresses Gemini Code Assist suggestion: specifying the element type improves type safety and is consistent with _all_module_input which already uses a parameterized type annotation. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Xesdiny --- src/llmcompressor/modifiers/autoround/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llmcompressor/modifiers/autoround/base.py b/src/llmcompressor/modifiers/autoround/base.py index 65c87e4d46..60cfbb5812 100644 --- a/src/llmcompressor/modifiers/autoround/base.py +++ b/src/llmcompressor/modifiers/autoround/base.py @@ -177,7 +177,7 @@ class AutoRoundModifier(Modifier, QuantizationMixin): # private variables _all_module_input: dict[str, list[tuple]] = PrivateAttr(default_factory=dict) _q_input: torch.Tensor | None = PrivateAttr(default=None) - _consumed_layers: set = PrivateAttr(default_factory=set) + _consumed_layers: set[str] = PrivateAttr(default_factory=set) def on_initialize(self, state: State, **kwargs) -> bool: """ From 27e4a58f12a40f966c424480871515dbbc758f1d Mon Sep 17 00:00:00 2001 From: xesdiny Date: Tue, 18 Aug 2026 18:48:23 +0800 Subject: [PATCH 3/4] refactor(autoround): replace _consumed_layers guard with early hook removal MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per reviewer feedback (@yiliu30): instead of tracking consumed layers in a set and gating the hook with an early return, remove each layer's input-capture hook via remove_hooks() before AutoRound optimization begins. Hooks for subsequent layers remain active until their own calibration pass. This eliminates the extra _consumed_layers state entirely: _capture_hooks maps each decoding layer name to its RemovableHandle registered in on_calibration_start, and the handle is popped + removed immediately before quantize_block is called. Also per @gemini-code-assist: changed .pop(key) to .pop(key, None) with a descriptive RuntimeError on missing calibration inputs. Validated on 8× L20 (gradient_accumulate_steps=8, iters=200, Qwen3.5-35B-A3B): Δalloc=+0.0 MB across all 200 iterations on all 8 ranks, alloc locked at 9846 MB throughout. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Xesdiny --- src/llmcompressor/modifiers/autoround/base.py | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/src/llmcompressor/modifiers/autoround/base.py b/src/llmcompressor/modifiers/autoround/base.py index 60cfbb5812..5f0a3a430f 100644 --- a/src/llmcompressor/modifiers/autoround/base.py +++ b/src/llmcompressor/modifiers/autoround/base.py @@ -177,7 +177,7 @@ class AutoRoundModifier(Modifier, QuantizationMixin): # private variables _all_module_input: dict[str, list[tuple]] = PrivateAttr(default_factory=dict) _q_input: torch.Tensor | None = PrivateAttr(default=None) - _consumed_layers: set[str] = PrivateAttr(default_factory=set) + _capture_hooks: dict = PrivateAttr(default_factory=dict) def on_initialize(self, state: State, **kwargs) -> bool: """ @@ -227,15 +227,6 @@ def start_calibration(self, model: torch.nn.Module): def input_capture_hook(self, module, args, kwargs): name = module._tmp_name - # After apply_autoround consumes a layer's inputs, it is added to - # _consumed_layers so this hook becomes a no-op for subsequent forward - # calls during the SignSGD optimization loop. Without this guard the - # hook fires on every mini-batch forward (gradient_accumulate_steps × - # iters calls) and holds the hidden_states tensor alive, causing - # ~(gradient_accumulate_steps × hidden_states_size) MB of GPU memory - # growth per outer iteration. - if name in self._consumed_layers: - return self._all_module_input.setdefault(name, []).append((args, kwargs)) def on_calibration_start(self, state: State, event: Event, **kwargs): @@ -244,10 +235,10 @@ def on_calibration_start(self, state: State, event: Event, **kwargs): self.start_calibration(state.model) for _, module in state.model.named_modules(): if self._is_decoding_layer(module): - # register input capture hook for decoding layers - self.register_hook( + handle = self.register_hook( module, self.input_capture_hook, "forward_pre", with_kwargs=True ) + self._capture_hooks[module._tmp_name] = handle def on_sequential_epoch_end( self, state: State, event: Event, modules: list[torch.nn.Module], **kwargs @@ -322,10 +313,19 @@ def apply_autoround(self, state, modules): ar.batch_dim = 0 first_param = next(decoding_layer.parameters()) device = first_param.device - # Mark consumed before pop so the hook immediately becomes a no-op - # for any forward call triggered during the optimization loop. - self._consumed_layers.add(decoding_layer._tmp_name) - cur_inputs = self._all_module_input.pop(decoding_layer._tmp_name) + # Remove this layer's input-capture hook before optimization begins + # so subsequent forward passes during the SignSGD loop do not + # re-populate _all_module_input with hidden_states tensors. + layer_name = decoding_layer._tmp_name + if layer_name in self._capture_hooks: + self.remove_hooks({self._capture_hooks.pop(layer_name)}) + cur_inputs = self._all_module_input.pop(layer_name, None) + if not cur_inputs: + raise RuntimeError( + f"No calibration inputs captured for layer {layer_name}. " + "This can happen if calibration data is missing or the " + "forward pass did not execute for this layer." + ) self._set_attention_masks(ar, decoding_layer, cur_inputs) decoding_layer.tuning_device = device # Only hand device placement to AutoRound when the caller explicitly @@ -371,7 +371,7 @@ def on_calibration_end(self, state: State, event: Event, **kwargs): self._remove_temporary_names(state.model) self.remove_hooks() self._q_input = None - self._consumed_layers.clear() + self._capture_hooks.clear() def get_unquantized_layer_names(self, wrapped_model: torch.nn.Module) -> list[str]: unquantized_layers = [] From e2f6fc51865ad40d7718caf70fad722a7a767130 Mon Sep 17 00:00:00 2001 From: Xesdiny Date: Mon, 17 Aug 2026 10:46:32 +0800 Subject: [PATCH 4/4] fix(autoround): add specific type hint dict[str, Any] for _capture_hooks Per Gemini code review feedback: use a more specific type annotation for _capture_hooks to match the style of other private attributes (e.g. _all_module_input: dict[str, list[tuple]]). Signed-off-by: Xesdiny --- src/llmcompressor/modifiers/autoround/base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/llmcompressor/modifiers/autoround/base.py b/src/llmcompressor/modifiers/autoround/base.py index 5f0a3a430f..cd824af1ac 100644 --- a/src/llmcompressor/modifiers/autoround/base.py +++ b/src/llmcompressor/modifiers/autoround/base.py @@ -1,5 +1,6 @@ import os from contextlib import contextmanager +from typing import Any import torch import torch.nn as nn @@ -177,7 +178,7 @@ class AutoRoundModifier(Modifier, QuantizationMixin): # private variables _all_module_input: dict[str, list[tuple]] = PrivateAttr(default_factory=dict) _q_input: torch.Tensor | None = PrivateAttr(default=None) - _capture_hooks: dict = PrivateAttr(default_factory=dict) + _capture_hooks: dict[str, Any] = PrivateAttr(default_factory=dict) def on_initialize(self, state: State, **kwargs) -> bool: """