Skip to content

AWQ Apply Scales Bugfix when smooth layer output length doesn't match balance layer input length #1451

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 21, 2025
Merged

AWQ Apply Scales Bugfix when smooth layer output length doesn't match balance layer input length #1451

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
575d12f
p1
brian-dellabetta May 16, 2025
35a0399
more to do
brian-dellabetta May 16, 2025
58a715d
only scale subset with valid shape
brian-dellabetta May 19, 2025
876bed0
update to skip qkv_proj->o_proj mappings on shape condition
brian-dellabetta May 20, 2025
5f06c06
style fixes
brian-dellabetta May 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 31 additions & 12 deletions src/llmcompressor/modifiers/awq/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -310,14 +310,24 @@ def _set_resolved_mappings(self, model: Module) -> None:
if not balance_layer:
continue

# exclude v_proj/o_proj mappings whose shapes are incompatible
# exclude v_proj->o_proj mappings whose shapes are incompatible
# https://github.com/mit-han-lab/llm-awq/pull/67#issuecomment-1681632777
if (
".v_proj" in layer_name
and ".o_proj" in balance_name
and isinstance(smooth_layer, torch.nn.Linear)
isinstance(smooth_layer, torch.nn.Linear)
and isinstance(balance_layer, torch.nn.Linear)
and smooth_layer.weight.shape != balance_layer.weight.shape
and ".o_proj" in balance_name
and (
(
".v_proj" in layer_name
and smooth_layer.out_features
!= balance_layer.in_features
)
or (
".qkv_proj" in layer_name
and smooth_layer.out_features
!= 3 * balance_layer.in_features
)
)
):
num_skipped_oproj_mappings += 1
continue
Expand Down Expand Up @@ -466,33 +476,42 @@ def _apply_smoothing(self, model: Module) -> None:
inp, w_mean, x_mean, module2inspect, balance_layers, fp16_output
)

scales = best_scales

@torch.no_grad()
def smooth(module):
with align_module_device(module):
scales = best_scales.to(module.weight.device)
if module in balance_layers:
module.weight.mul_(scales.view(1, -1).to(module.weight.device))
update_offload_parameter(
module,
"weight",
module.weight.mul_(scales.view(1, -1)),
)
elif module == smooth_layer:
if module.weight.ndim == 1:
update_offload_parameter(
module,
"weight",
module.weight.div(scales.to(module.weight.device)),
module.weight.div_(scales),
)
else:
# NOTE: edge case when smooth layer number of out_features
# is not equal to balance layer number of in_features
# e.g. when fused qkv_proj is used to smooth o_proj
# in this case, default to scaling the last output features
# because the desired smooth layer is v_proj
# https://github.com/casper-hansen/AutoAWQ/blob/main/awq/quantize/scale.py#L123
update_offload_parameter(
module,
"weight",
module.weight.div(
scales.view(-1, 1).to(module.weight.device)
module.weight[-scales.size(0) :].div_(
scales.view(-1, 1)
),
)
if hasattr(module, "bias") and module.bias is not None:
update_offload_parameter(
module,
"bias",
module.bias.div(scales.to(module.bias.device)),
module.bias.div_(scales),
)

parent = get_fsdp_parent(mapping.smooth_name, model)
Expand Down