(module,
w_bit, q_config,
input_feat)
| 85 | |
| 86 | @torch.no_grad() |
| 87 | def auto_clip_block(module, |
| 88 | w_bit, q_config, |
| 89 | input_feat): |
| 90 | |
| 91 | named_linears = {name: m for name, |
| 92 | m in module.named_modules() if isinstance(m, nn.Linear)} |
| 93 | |
| 94 | clip_list = [] |
| 95 | for name in named_linears: |
| 96 | # due to qk bmm, it is hard to clip precisely |
| 97 | if any([_ in name for _ in ["q_", "k_", "query", "key", "Wqkv"]]): |
| 98 | continue |
| 99 | named_linears[name].cuda() |
| 100 | |
| 101 | max_val, min_val = auto_2clip_layer( |
| 102 | named_linears[name].weight, input_feat[name], n_bit=w_bit, q_config=q_config) |
| 103 | |
| 104 | clip_list.append((name, max_val, min_val)) |
| 105 | |
| 106 | named_linears[name].cpu() |
| 107 | return clip_list |
| 108 | |
| 109 | @torch.no_grad() |
| 110 | def run_clip( |
no test coverage detected