Add RunningAbsMaxSmoothQuantObserver for memory-efficient calibration (#3946)

[jcaip]

Summary:

Add a memory-efficient SmoothQuant observer that uses running per-channel
absmax instead of storing all calibration inputs. This reduces calibration
memory from O(N x features) to O(features), preventing RAM spikes and OOM
kills when calibrating on large datasets.

• Add RunningAbsMaxSmoothQuantObserver class in core.py
• Add use_running_absmax config option to SmoothQuantConfig
• Export the new observer from the module

Differential Revision: D94260071

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3946

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 04581a1 with merge base 960f307 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@jcaip has exported this pull request. If you are a Meta employee, you can view the originating Diff in D94260071.

[namgyu-youn]

Not sure why test script is located inside package. Shouldn't we update https://github.com/pytorch/ao/blob/main/test/prototype/test_smoothquant.py instead?

[jcaip]

Yes, will tell Hossein to do that before merging

[namgyu-youn]

Like this new logic overally, but do we actually need RunningAbsMaxSmoothQuantObserver? To me, minimal SmoothQuantObserver update (even though its bc-breaking) looks better.

Also, RunningAbsMaxSmoothQuantObserver reminds me AQT (AffineQuantizedTensor) nightmare. I am not sure if we really need abstracted notation.

[jcaip]

Yeah agree SmoothQuantObserver abstraction is kind of funky.

I think the best way this mess is to define a ObserverTensor class that we can then modify appropriately. That way we don't need to have SmoothQuantObserver, SmoothQuantObservedLinear and can just use a single SmoothQuantObserverTensor.

Do you know if smoothquant API is used much anywhere? The only thing that concerns me is that by default replacing SmoothQuantObserver with RunningAbsMaxSmoothQuantObserver is that we need two passes through the data for RunningAbsMaxSmoothQuantObserver, so existing workflows might fail.

We do need a way to test numerics though between SmoothQuantObserver and RunningAbsMaxSmoothQuantObserver, so even if we make RunningAbsMaxSmoothQuantObserver we'll still need to store this code somewhere. I think it's better to leave the cleanup as a follow up.

[namgyu-youn]

I think the best way this mess is to define a ObserverTensor class that we can then modify appropriately. That way we don't need to have SmoothQuantObserver, SmoothQuantObservedLinear and can just use a single SmoothQuantObserverTensor.

So are you considering common observer class? More generally, quantization subclass (Float8Tensor/Int8Tensor) should not contain observer-related ops (like act_pre_scale)? Actually I was working for this topic, so it would be great if you can look into it — #3925.

[namgyu-youn]

Do you know if smoothquant API is used much anywhere? The only thing that concerns me is that by default replacing SmoothQuantObserver with RunningAbsMaxSmoothQuantObserver is that we need two passes through the data for RunningAbsMaxSmoothQuantObserver, so existing workflows might fail.

Yeah I am also interested in this topic. For example, because AWQ uses greedy search

ao/torchao/prototype/awq/core.py

Lines 70 to 86 in 5045d76

	for i in range(self.scale_options):
	ratio = i * 1 / self.scale_options
	scales = x_max.pow(ratio).to(self.weight.dtype).clamp(min=1e-4).view(-1)
	if best_scales is None:
	best_scales = torch.ones_like(scales)
	scales = scales / (scales.max() * scales.min()).sqrt()
	config_handler = _QUANTIZE_CONFIG_HANDLER[type(self.base_config)]
	dummy_mod = DummyModule(self.weight * scales)
	quant_mod = config_handler(dummy_mod, self.base_config)
	w = quant_mod.weight
	orig_out = F.linear(acc, self.weight, self.bias)
	q_out = F.linear(acc / scales, w, self.bias)
	loss = (orig_out - q_out).pow(2).mean().item()
	if loss < best_loss:
	best_scales = scales
	best_loss = loss

, we can try better searching algorithm like binary/tenary. It's GPTQ, but I already observed tenary is 2x faster than greedy: ModelCloud/GPTQModel#2419

Wondering if you have any candidates for these numerical techniques, what about your thought?

[namgyu-youn]

We do need a way to test numerics though between SmoothQuantObserver and RunningAbsMaxSmoothQuantObserver, so even if we make RunningAbsMaxSmoothQuantObserver we'll still need to store this code somewhere. I think it's better to leave the cleanup as a follow up.

I think in this early stage, we don't need multiple-configs to store all of them. Also I don't want to see RunningAbsMaxSmoothQuantObserver anywhere.

Instead, we can just run perf using https://github.com/pytorch/ao/blob/main/benchmarks/quantization/calibration_based/measure_accuracy_and_performance.sh and update minimally across AWQ/GPTQ/SmoothQuant I feel.

[jcaip]

@claude can you fix lint please

[claude]

Claude finished @jcaip's task in 10m 33s —— View job

---

Fixed lint issues

• Identify linting tools and configuration — repo uses ruff (v0.11.6) for linting and formatting
• Run linter to find issues — CI ruff format --check failed on torchao/quantization/quantize_/common/quantization_step.py
• Fix formatting in quantization_step.py
• Commit and push fixes

Changes made

Applied ruff format style fixes to torchao/quantization/quantize_/common/quantization_step.py:

1. Collapsed single-line docstring closing """ onto the same line (e.g., """Insert observers before running calibration flow.""")
2. Added blank line after each docstring expression (standard ruff format)
3. Wrapped long enum value assignments that exceeded 88-char line limit with parentheses:
   • PREPARE_FOR_SMOOTHQUANT_SMOOTHING_FACTOR
   • PREPARE_FOR_SMOOTHQUANT_ACTIVATION_SCALES

Note: I was unable to install/run ruff locally due to permission restrictions, so I applied the formatting rules manually based on Black/ruff conventions. If CI still fails, the formatting may need further adjustment.
| Branch

[meta-codesync]

@jcaip has imported this pull request. If you are a Meta employee, you can view this in D94260071.

[namgyu-youn]

@jcaip @jerryzh168 do you have any GPU profiling e2e results for this PR? If I remember correctly, there wasn't a profiling benchmark at the time — just wanted to check if that's been added or if there are any benchmark numbers available I am missing.