Function auto_2clip_layer

quantization/autoclip.py:14–83

(w, input_feat, n_bit, q_config,
                    n_grid=20,
                    max_shrink=0.5,
                    n_sample_token=512)

@torch.no_grad()
def auto_2clip_layer(w, input_feat, n_bit, q_config,
                    n_grid=20,
                    max_shrink=0.5,
                    n_sample_token=512):
    assert w.dim() == 2
    org_w_shape = w.shape
    # w           [co, ci]      -> [co, 1, n_group, group size]
    # input_feat  [n_token, ci] -> [1, n_token, n_group, group size]

    group_size = q_config["q_group_size"] if q_config["q_group_size"] > 0 else w.shape[1]

    input_feat = input_feat.view(-1, input_feat.shape[-1])
    input_feat = input_feat.reshape(1, input_feat.shape[0], -1, group_size)
    input_feat = input_feat[:, 0::input_feat.shape[1] // n_sample_token]
    w = w.reshape(w.shape[0], 1, -1, group_size)

    oc_batch_size = 256 if w.shape[0] % 256 == 0 else 64  # prevent OOM
    
    assert w.shape[0] % oc_batch_size == 0
    w_all = w
    best_max_val_all = []
    best_min_val_all = []
    for i_b in range(w.shape[0] // oc_batch_size):
        w = w_all[i_b * oc_batch_size: (i_b + 1) * oc_batch_size]

        org_max_val = w.amax(dim=-1, keepdim=True)  # co, 1, n_group, 1
        org_min_val = w.amin(dim=-1, keepdim=True)

        best_max_val = org_max_val.clone()
        best_min_val = org_min_val.clone()

        min_errs = torch.ones_like(org_max_val) * 1e9
        input_feat = input_feat.to(w.device)
        org_out = (input_feat * w).sum(dim=-1)  # co, n_token, n_group

        for i_s_p in range(int(max_shrink * n_grid)):
            max_val = org_max_val * (1 - i_s_p / n_grid)
            for i_s_n in range(int(max_shrink * n_grid)):
                min_val = org_min_val * (1 - i_s_n / n_grid)
                # min_val = - max_val
                cur_w = torch.clamp(w, min_val, max_val)
                if q_config["quant_type"] == "int":
                    q_w = pseudo_quantize_tensor(cur_w, n_bit=n_bit, zero_point=True, q_group_size=q_config['q_group_size'])
                elif q_config["quant_type"] == "nf3":
                    q_w = pseudo_quantize_n2f3_tensor(cur_w, q_group_size=q_config['q_group_size'])
                else:
                    quant_type = q_config["quant_type"]
                    raise ValueError(f"Has no support {quant_type}. Valid quant_type:[int, nf3]")
                    
                cur_out = (input_feat * q_w).sum(dim=-1)

                err = (cur_out - org_out).pow(2).mean(dim=1).view(min_errs.shape)

                del cur_w
                del cur_out
                cur_best_idx = err < min_errs
                min_errs[cur_best_idx] = err[cur_best_idx]
                best_max_val[cur_best_idx] = max_val[cur_best_idx]