Let AWQ mappings absorb scales into a sliced projection

[Ar4ikov]

SUMMARY:

AWQMapping.smooth_layer is a single module name, and _apply_smoothing folds 1/s
into that one module's weight. That covers every architecture in the registry today,
because in all of them the balance layers' input is produced by one module (a norm, or
a preceding projection).

It cannot express an AdaLN-modulated block. There the balance layers see

norm(x) * (1 + scale) + shift

and scale / shift are slices of one shared projection's output, not modules of
their own. Concretely, in MiniMax-H3's diffusion transformer
(MiniMaxH3Transformer3DModel, diffusers#14355):

shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb)
norm_hidden_states = self.norm1(hidden_states) * (
    1.0 + scale_msa.index_select(0, adaln_indices)
) + shift_msa.index_select(0, adaln_indices)

There is no module whose weight can absorb 1/s, so AWQ has nothing to fold into and
these models fall back to default_mappings, which do not match.

The absorption does exist, and it is exact:

y / s = (norm(x) / s) * (1 + scale) + (shift / s)

norm(x)/s comes from dividing the norm's weight, and shift/s from dividing the rows
of the shared projection that emit the shift chunk. scale is left alone: it multiplies
elementwise, and the 1 + term is already handled by the norm weight, so scaling scale
as well would double-apply the correction.

This PR adds the smallest thing that can express that:

• SlicedSmoothTarget(layer, chunk_index, num_chunks, repeat) names a contiguous row
  slice of another projection's output, in the layout produced by
  view(-1, num_chunks * hidden) followed by chunk(num_chunks, -1), tiled repeat
  times.
• AWQMapping.extra_smooth_targets — an optional list of them, defaulting to empty.
• absorb_sliced_scales() applies 1/s to those rows, weight and bias.
• Mappings for MiniMaxH3Transformer3DModel.

I deliberately did not widen smooth_layer's type to a union. That would ripple
through match_modules_set, ResolvedMapping, _check_layers_are_compatible and every
registered architecture for no benefit, since the primary smooth layer stays a single
module in this case too. The extra targets are resolved the way
activation_hook_target already is, by dotted path, with one difference documented in
the dataclass: they resolve against the module holding smooth_layer, not against the
balance layers' ancestor. In an AdaLN block the balance layers sit under attn / ff
while the shared projection is a sibling of the norm, so resolving against the ancestor
would not find it.

Two things left out on purpose:

• No ff.net.0.proj -> ff.net.2 mapping for MiniMax-H3. That projection is a fused
  SwiGLU, and the existing Linear-smooth branch folds 1/s into the last
  out_features, which here are the gate half. The block output is not linear in the
  gate, so that fold would not be equivalence preserving. Smoothing the up half needs a
  sliced target on the smooth side, which is a separate change; the same gap applies to
  any fused-SwiGLU architecture, not just this one.
• No diffusion calibration path. oneshot feeds tokenized text, while a DiT takes
  (hidden_states, temb, adaln_indices, rotary_emb, attention_mask) that only a real
  sampling run produces. So the mapping this PR registers is not reachable end to end
  from llm-compressor today. I would rather land the mapping mechanism separately from
  the data pipeline, but if you would prefer this to wait until a diffusion calibration
  source exists, say so and I will hold it.

TEST PLAN:

New tests/llmcompressor/modifiers/transform/awq/test_sliced_smooth.py, six unit tests
on a synthetic AdaLN block, so nothing here depends on diffusers or on a downloaded
checkpoint:

$ pytest tests/llmcompressor/modifiers/transform/awq/ -v -m unit
...
test_sliced_smooth.py::test_absorb_sliced_scales_preserves_output PASSED                  [ 89%]
test_sliced_smooth.py::test_absorbing_the_scale_chunk_is_not_equivalent PASSED            [ 91%]
test_sliced_smooth.py::test_absorb_sliced_scales_only_touches_its_own_chunk PASSED        [ 93%]
test_sliced_smooth.py::test_absorb_sliced_scales_rejects_shape_mismatch PASSED            [ 95%]
test_sliced_smooth.py::test_extra_smooth_targets_resolve_against_the_smooth_layer_parent PASSED [ 97%]
test_sliced_smooth.py::test_minimax_h3_mappings_are_registered PASSED                     [100%]

================ 46 passed, 3 deselected, 14 warnings in 3.34s =================

That run covers the whole transform/awq directory, so the 40 pre-existing tests in
test_base.py and test_dynamic_mappings.py are in it and still pass: the new dataclass
field is optional and defaults to empty, so no registered architecture changes behaviour.
Environment: torch 2.13.0+cu130, transformers 5.14.1, compressed-tensors
0.17.2.a20260731, CPU only, no GPU needed.

The load-bearing one is test_absorb_sliced_scales_preserves_output: fold 1/s into the
norm weight and the shift rows, scale the balance layer's input columns by s, and the
block's output must be unchanged. In float64 it matches to rtol=atol=1e-12.

test_absorbing_the_scale_chunk_is_not_equivalent is the control. It does the same thing
but folds into the scale chunk instead of shift, and asserts the output changes.
Without it the first test would pass for the wrong reason, since a no-op would also
"preserve" the output.

One caveat on the usual "does the test fail without the patch" check: for a new API it
fails only with ImportError, which proves nothing. That control test is the real guard,
which is why I wrote it.

I also ran the derivation against the real module rather than only the synthetic one.
Building MiniMaxH3TransformerBlock from diffusers' minimax-h3 branch, hidden 64,
float64, and applying the same fold:

max abs diff : 2.220e-16
max rel diff : 6.876e-17
control (wrong fold) rel diff: 1.752e-02

Checks on the changed files:

$ ruff check src/llmcompressor/modifiers/transform/awq/ tests/llmcompressor/modifiers/transform/awq/test_sliced_smooth.py
All checks passed!

ruff format --check reports one pre-existing reformat in base.py (an assert at
line ~789) and one in test_base.py, neither of which this PR touches; my local ruff is
0.15.18 and evidently differs from the pinned one, so I left both alone rather than add
unrelated churn.

Context: I hit this quantizing MiniMax-H3's 33B joint video/audio DiT to fit on 2x24GB.
adaln_proj is 13.0B of the 33.1B parameters there, which is also why keeping it out of
the quantized set matters: the diffusers docstring notes that a rounding applied before
its silu biases every block's modulation identically at every sampling step, so the error
accumulates coherently along the denoising trajectory.

@brian-dellabetta, since you built the mapping mechanism this extends in #2526 and #2451,
and @kylesayrs as a second pair of eyes on the _apply_smoothing change. Happy to reshape
the API if you would rather express this differently, or to hold it until a diffusion
calibration source lands.

[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: 96dc7d76-85ed-49cb-8379-97a74d462c30

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

🔴 2 of 2 protections blocking · waiting on 👀 reviews

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

🔴 Require one maintainer review

Waiting for any of

• approved-reviews-by=HDCharles
• approved-reviews-by=brian-dellabetta
• approved-reviews-by=dsikka
• approved-reviews-by=kylesayrs
• approved-reviews-by=yiliu30

This rule is failing.

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

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

🔴 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

[gemini-code-assist]

Code Review

This pull request introduces support for smoothing AdaLN-modulated blocks in AWQ by adding SlicedSmoothTarget and updating the mapping resolution and smoothing logic to handle extra smooth targets. This is particularly useful for architectures like MiniMaxH3Transformer3DModel where shared projections need to absorb scaling factors for specific slices. The review feedback suggests making absorb_sliced_scales more robust by handling the device transfer of the scales tensor internally, which also simplifies its caller.

[gemini-code-assist]

To make absorb_sliced_scales more robust and self-contained, it is safer to move the scales tensor to the same device as module.weight inside the function itself. This prevents potential device mismatch runtime errors if the function is called from other contexts or tests where scales is not pre-aligned to the module's device.

Suggested change

	hidden = scales.size(0)
	group = target.num_chunks * hidden
	scales = scales.to(module.weight.device)
	hidden = scales.size(0)
	group = target.num_chunks * hidden

[gemini-code-assist]

Once the device transfer is handled internally within absorb_sliced_scales, we can simplify the caller here by passing best_scales directly without the explicit .to(...) call.

                    absorb_sliced_scales(module, target, best_scales)

[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.

[brian-dellabetta]

Hi @Ar4ikov , thanks for raising this. We have hit this norm issue with other custom norm implementations which do

norm(x) * (1 + scale)

We can probably extend this to include for the + shift addition, so that we have a general solution that works for AWQ as well as smoothquant. Can you take a look and see if this can also be extended for your use case?

https://github.com/vllm-project/llm-compressor/blob/main/src/llmcompressor/modeling/offset_norm.py#L64

[Ar4ikov]

Thanks @brian-dellabetta, I looked at offset_norm.py. Short answer: the idea generalizes and I would like to build that, but it cannot be done by extending CalibrationOffsetNorm in place, because in the AdaLN case the thing that needs the offset treatment is a slice of a Linear's bias rather than a norm weight. Details, with numbers for both layouts I checked.

There are two distinct AdaLN shapes in diffusers, and they differ in exactly the way that matters here.

1. Modulation hoisted out of the norm, norm keeps its weight. MiniMax-H3: adaln_proj is a block-level module shared by norm1, norm2 and three modalities, and norm1 is a plain nn.RMSNorm with a learnable weight.

shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(temb)
norm_hidden_states = self.norm1(hidden_states) * (1.0 + scale_msa.index_select(0, adaln_indices)) \
                     + shift_msa.index_select(0, adaln_indices)

Here the offset_norm problem does not actually arise. The 1 + multiplies a runtime tensor, not the stored weight, so norm1.weight /= s is already correct: it scales normalize(x) * g, and that distributes over (1 + scale). Only the additive term needs anything, and it needs the shift rows of adaln_proj.linear divided by s. That is what this PR does, and on the real module in float64 it is exact:

max rel diff                   : 6.876e-17
control (fold into scale rows) : 1.752e-02

2. Modulation inside the norm module, norm has no weight. AdaLayerNormZero (Flux, SD3) and AdaLayerNormZeroSingle own their linear, and their norm is nn.LayerNorm(dim, elementwise_affine=False). So there is no weight to divide and the constant 1 has nowhere to go. This is where your offset idea is needed, applied to the scale chunk of the linear:

W[scale_rows] /= s
b[scale_rows]  = (1 + b[scale_rows]) / s - 1     # the offset_norm trick, on a bias slice
W[shift_rows] /= s
b[shift_rows] /= s

Verified the same way on diffusers.models.normalization.AdaLayerNormZero, float64:

unified fold rel diff : 7.490e-16
naive fold   rel diff : 1.177e+00      # dividing b[scale_rows] by s like the shift

The naive variant is off by 118%, so this really is load-bearing and not a rounding detail.

So a single rule covers all three families:

Layout	Multiplicative side	Additive side
Gemma, Qwen3Next: n(x) * (1 + w), static	offset_norm as today	none
MiniMax-H3: norm has weight, projection hoisted	norm.weight /= s	shift rows /= s
Flux, SD3: norm has no weight, projection inside	scale rows, W /= s and b = (1+b)/s - 1	shift rows /= s

Two things constrain where this can live.

The chunk layouts differ: AdaLayerNormZero is chunk(6), AdaLayerNormZeroSingle is chunk(3), AdaLayerNormContinuous is chunk(2), and H3 additionally tiles the whole 6-chunk group once per modality and reshapes with view(-1, 6 * hidden). That is why SlicedSmoothTarget in this PR is parameterized by (chunk_index, num_chunks, repeat) rather than hardcoding an offset. I have not tested AdaLayerNormContinuous; by the table above it should be num_chunks=2, but I would rather say I did not check it than imply I did.

More importantly, NormCalibrationModule is registered by norm class name and replaces the norm module. That reaches case 3, where the linear is inside the module, but it cannot reach case 2: H3's projection is a sibling of the norm, shared between the attention and MLP branches, so no norm-module replacement can see it. Whatever hosts this needs a pointer to a module that is not the norm, which is the one thing the current registry cannot express.

Happy to do the work either way you prefer:

• keep it at mapping level as here, and add an offset: bool on SlicedSmoothTarget so the scale chunk gets (1+b)/s - 1 instead of /s, which turns this into a general AdaLN solution and lets me add Flux and SD3 mappings alongside H3, or
• move the absorption into a shared helper that both AWQ and SmoothQuant call, if you want the SmoothQuant path covered in the same change. It is currently a private function in the AWQ modifier, and you did mention SmoothQuant, so say the word and I will hoist it.

One caveat that applies whatever we choose, and that I would rather flag than have found in review: none of this is reachable end to end from llm-compressor today, because oneshot has no way to produce DiT calibration inputs. So the mechanism would land ahead of its first real user. I am fine with that being a reason to hold the PR.

[brian-dellabetta]

Hi @Ar4ikov , I think I'm following -- I don't have much experience with diffusion models or the diffusers package, and we've not had many user posts regarding it.

That reaches case 3, where the linear is inside the module, but it cannot reach case 2: H3's projection is a sibling of the norm, shared between the attention and MLP branches, so no norm-module replacement can see it. Whatever hosts this needs a pointer to a module that is not the norm, which is the one thing the current registry cannot express.

Regarding NormCalibrationModule, the name is a bit misleading. It is applied generically to any module in model.children() that matches class name. We used to have an MoECalibrationModule that would linearize 3D experts so they could be treated as instances of torch.nn.Linear (it was removed due to enhancements in transformers v5). So you could create a calibration module for the grouped module itself, if by updating the forward passes and weights you can move to a purely linear op and back out to the original implementation. Would that cover all your bases here? We are attempting to shoehorn everything in this way because many of the algorithms (AWQ, SmoothQuant, SpinQuant) are based on a linear transformation of weights mapping to a linear transformation on the activations. Just hoping we can retain that assumption generally, rather than fixing at the modifier level