fix(autoround): prevent input_capture_hook from accumulating GPU memory during optimization

[xesdiny]

Summary

Fixes #3023

AutoRoundModifier.input_capture_hook is registered on all decoding layers during
on_calibration_start() but is never disabled 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=N
this retains N hidden_states tensors per outer iteration, causing
~N × hidden_states_size MB of GPU memory growth per iteration.

At gradient_accumulate_steps=8, hidden_size=2048, seq_len=6144 (bf16):
+204 MB/iter → ~40 GB accumulated at ITERS=200 → OOM on 44 GB GPUs.

Changes

• Add _consumed_layers: set PrivateAttr to track layers processed by apply_autoround()
• input_capture_hook: use setdefault() for calibration-phase initialization; exit immediately for consumed layers
• apply_autoround(): mark layer consumed before pop()-ing its inputs
• on_calibration_end(): clear _consumed_layers

Verification

After fix, on 8 × RTX 5880 (44 GB), Qwen/Qwen3.5-35B-A3B, gradient_accumulate_steps=8:

Before: Δalloc = +204 MB/iter across all 8 ranks
After:  Δalloc = +0.0 MB/iter across all 8 ranks (20 iterations verified)

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 21da54bc-e573-497a-8f22-76114b9a3d7f

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🔴 1 of 2 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require two reviews	👀 reviews
🟢	Require one maintainer review	—

🔴 Require two reviews

Waiting for

• #approved-reviews-by >= 2

This rule is failing.

PRs labelled "two-reviews" must have at least two approving reviews before merging.

• #approved-reviews-by >= 2
• #changes-requested-reviews-by = 0

---

Show 1 satisfied protection

🟢 Require one maintainer review

All PRs must have at least one approving review from a maintainer before merging.

• #changes-requested-reviews-by = 0
• any of:
  • approved-reviews-by=yiliu30
  • approved-reviews-by=HDCharles
  • approved-reviews-by=brian-dellabetta
  • approved-reviews-by=dsikka
  • approved-reviews-by=kylesayrs

[gemini-code-assist]

Code Review

This pull request introduces a _consumed_layers set to track and short-circuit input capture hooks once a layer's inputs are consumed during AutoRound calibration, effectively preventing GPU memory growth during optimization loops. The feedback suggests specifying the element type for the _consumed_layers set as set[str] to improve type safety and maintain consistency with other private attributes.

[gemini-code-assist]

To improve type safety and maintain consistency with other private attributes (such as _all_module_input), please specify the element type for the _consumed_layers set as set[str].

Suggested change

	_consumed_layers: set = PrivateAttr(default_factory=set)
	_consumed_layers: set[str] = PrivateAttr(default_factory=set)

[xesdiny]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request introduces a mechanism to track consumed layers during AutoRound calibration by adding a _consumed_layers set. This ensures that the input capture hook becomes a no-op once a layer's inputs are processed, preventing GPU memory growth from accumulated hidden states. The reviewer suggested using .pop(..., None) instead of .pop() when retrieving captured inputs to avoid a potential KeyError if calibration data is missing, and instead raising a more descriptive RuntimeError.

[gemini-code-assist]

Using .pop() with only the key will raise a KeyError if the layer's temporary name is not present in _all_module_input (for example, if calibration data was empty or the forward pass did not run for this layer). Using .pop(..., None) and raising a descriptive RuntimeError provides a clearer error message for troubleshooting.

Suggested change

	self._consumed_layers.add(decoding_layer._tmp_name)
	cur_inputs = self._all_module_input.pop(decoding_layer._tmp_name)
	self._consumed_layers.add(decoding_layer._tmp_name)
	cur_inputs = self._all_module_input.pop(decoding_layer._tmp_name, None)
	if not cur_inputs:
	raise RuntimeError(
	f"No calibration inputs captured for layer {decoding_layer._tmp_name}. "
	"This can happen if calibration data is missing or the forward pass did not execute."
	)

[yiliu30]

Hi @xesdiny, good catch! Rather than adding type information to _consumed_layers, I’d prefer to remove the input-capture hook for each layer as soon as its calibration inputs have been collected, before AutoRound tuning begins.

The hooks for later layers can remain active until their own calibration pass. This way, we avoid retaining optimization inputs without introducing additional state for consumed layers.

[xesdiny]

Thanks for the reviews! Updated the implementation in cb3cbc0:

Per @yiliu30: replaced the _consumed_layers guard approach with early hook removal. Each layer's forward_pre hook is now stored by name in _capture_hooks at registration time, then removed via self.remove_hooks({handle}) immediately before quantize_block is called — so the hook is gone before the SignSGD optimization loop starts. Hooks for later layers remain active until their own calibration pass. This eliminates the extra tracking state entirely.

Local validation on 8× L20 (Qwen3.5-35B-A3B, gradient_accumulate_steps=8, iters=200, LIMIT_TO_LAYERS=2):

[MEM-ITER]   1  Δalloc=+0.0MB  alloc=9846MB
[MEM-ITER]  50  Δalloc=+0.0MB  alloc=9846MB
[MEM-ITER] 100  Δalloc=+0.0MB  alloc=9846MB
[MEM-ITER] 150  Δalloc=+0.0MB  alloc=9846MB
[MEM-ITER] 200  Δalloc=+0.0MB  alloc=9846MB

[yiliu30]

Thanks @xesdiny! This fix LGTM. The failing UTs are unrelated to this change, so please wait for the fix to land on the main branch.

[xesdiny]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request introduces hook management in the AutoRoundModifier to prevent re-populating calibration inputs during the SignSGD optimization loop. It tracks registered input-capture hooks in a new _capture_hooks private attribute, removes each layer's hook before its optimization begins, and raises a RuntimeError if no calibration inputs are captured. The review feedback suggests improving type safety by using a more specific type hint (dict[str, Any]) for the _capture_hooks attribute.

[gemini-code-assist]

For better type safety and consistency with other private attributes in this class (such as _all_module_input), consider using a more specific type hint for _capture_hooks, such as dict[str, Any].

Suggested change

	_capture_hooks: dict = PrivateAttr(default_factory=dict)
	_capture_hooks: dict[str, Any] = PrivateAttr(default_factory=dict)