Is it possible to smooth FC2 input?

[MaxwellWjj]

Hi Developers,

Recently when I try to apply smoothquant in my side (Qwen3-1.7B), I found that the FC2 (or the down_proj in Qwen-like definition) is not concluded in the smoothed layers. However, I observed that the static per-tensor scaling factors in this layer input, can be extremely large if no smooth is applied.

Layer 0: {'q_proj_input': 0.009227362204724409, 'o_proj_input': 0.021776574803149606, 'gate_input': 0.010765255905511811, 'down_input': 0.15748031496062992}

Layer 1: {'q_proj_input': 0.008427657480314961, 'o_proj_input': 0.011441929133858268, 'gate_input': 0.015071358267716535, 'down_input': 1.236220472440945}

Layer 2: {'q_proj_input': 0.009781003937007874, 'o_proj_input': 0.018331692913385825, 'gate_input': 0.023375984251968504, 'down_input': 133.03937007874015}

...

Layer 26: {'q_proj_input': 0.03297244094488189, 'o_proj_input': 2.031496062992126, 'gate_input': 0.022637795275590553, 'down_input': 11.21259842519685}

Layer 27: {'q_proj_input': 0.03641732283464567, 'o_proj_input': 3.0078740157480315, 'gate_input': 0.035679133858267716, 'down_input': 23.433070866141733}

As you can see, the down_input here refers to the per-tensor scale in down_proj (should be the same to fc2 in OPT). When accumulated through layers, the down_input scale becomes extremely large, which means the outliers here explode! Then of course, the final ppl does not look acceptable (from original ~16 to quantized ~90)

If we can apply smooth to this layer, I believe the result can improve a lot. May I know if you have tried to implement that? Many thanks!

[MaxwellWjj]

OK guys... I just have a try. Just simply modify the smooth_lm() function of smooth.py:

...
elif isinstance(module, (LlamaDecoderLayer, MistralDecoderLayer, Qwen2DecoderLayer, Qwen3DecoderLayer)):
    attn_ln = module.input_layernorm  # attention forward norm
    qkv = [
        module.self_attn.q_proj,
        module.self_attn.k_proj,
        module.self_attn.v_proj,
    ]

    qkv_input_scales = scales[name + ".self_attn.q_proj"]
    smooth_ln_fcs_llama_like(attn_ln, qkv, qkv_input_scales, alpha)

    ffn_ln = module.post_attention_layernorm  # feed forward norm
    fcs = [module.mlp.gate_proj, module.mlp.up_proj]
    fcs_input_scales = scales[name + ".mlp.gate_proj"]

    smooth_ln_fcs_llama_like(ffn_ln, fcs, fcs_input_scales, alpha)
    
    # ------------------------------------- ADD HERE! -------------------------------------------------
    # try to smooth qwen3 MLP's down_proj
    # check if scales has scales[name + ".mlp.down_proj"]
    if name + ".mlp.down_proj" in scales: 
        fc1 = module.mlp.up_proj
        fc2 = module.mlp.down_proj
        fc2_input_scales = scales[name + ".mlp.down_proj"]
    
        smooth_fc_fc_llama_like(fc1, fc2, fc2_input_scales, alpha)
    else:
        raise RuntimeError("Cannot find down_proj scales for Qwen3 MLP smoothing.")
    # ------------------------------------- ADD HERE! -------------------------------------------------
...

# ------------------------------------- Implementation of smooth_fc_fc_llama_like -----------------------------
@torch.no_grad()
def smooth_fc_fc_llama_like(fc1, fc2, act_scales, alpha=0.5):
    # for qwen3 MLP up_proj smoothing -> need to adjust gate_proj and up_proj
    assert isinstance(fc1, nn.Linear) and isinstance(fc2, nn.Linear)
    assert fc1.out_features == fc2.in_features == act_scales.numel()
    device, dtype = fc2.weight.device, fc2.weight.dtype
    act_scales = act_scales.to(device=device, dtype=dtype)
    weight_scales = fc2.weight.abs().max(dim=0, keepdim=True)[0]
    weight_scales = weight_scales.max(dim=0)[0].clamp(min=1e-5)
    print("act_scales shape:", act_scales.shape)
    print("weight_scales shape:", weight_scales.shape)
    scales = (
        (act_scales.pow(alpha) / weight_scales.pow(1 - alpha))
        .clamp(min=1e-5)
        .to(device)
        .to(dtype)
    )
    
    fc1.weight.div_(scales.view(-1, 1))
    fc2.weight.mul_(scales.view(1, -1))
# ------------------------------------- Implementation of smooth_fc_fc_llama_like -----------------------------

You don't even need to modify the gen_smq_act_scales() function in generate_act_scales.py, because it natually supports generating the scales[name + ".mlp.down_proj"] (or say fc2), since it is applied to all nn.Linear. But you have not used it before.

Then I set the alpha = 0.5, the ppl improves from ~95 to ~20. Although it still has drop from original ~16, I think this accuracy drop can be mitigated by finding a optimal alpha.

Also I have to mention that, the SwiGLU-based MLP in Llama and Llama-like provides the potential to smooth down_proj. The input of SiLU function (gate_proj) is different from up_proj which I set as the previous layer (just assume it as the layernorm). For some models using GELU, it may not be able to smooth down_proj since the GELU is an element-wise non-linear function.

Maybe you can modify this small change and push a new version to support more models :)