[Major] Add CPU offloading support for apply_scale, apply_clip, pseud…

…o_quantize_model_weight, real_quantize_model_weight
trotsky1997 · Jul 1, 2023 · d32095a · d32095a
1 parent 95cd9c2
commit d32095a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/awq/quantize/auto_clip.py b/awq/quantize/auto_clip.py
@@ -75,9 +75,11 @@ def auto_clip_block(module,
         # due to qk bmm, it is hard to clip precisely
         if any([_ in name for _ in ["q_", "k_", "query", "key", "Wqkv"]]):
             continue
+        named_linears[name].cuda()
         max_val = auto_clip_layer(
             named_linears[name].weight, input_feat[name], n_bit=w_bit, q_config=q_config)
         clip_list.append((name, max_val))
+        named_linears[name].cpu()
     return clip_list