Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 01:04:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

NNCF do quant on GPU

Browse Source
This commit is contained in:
Disty0
2025-04-23 23:07:23 +03:00
parent 789d7e8a52
commit 778de4d295
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/model_quant_nncf.py
+10 -3
View File
@@ -10,7 +10,7 @@ from diffusers.utils import get_module_from_name
from accelerate import init_empty_weights
from accelerate.utils import CustomDtype
from modules import devices
from modules import devices, shared
torch_dtype_dict = {
@@ -59,7 +59,11 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c
else:
reduction_axes = [layer.weight.ndim - 1]
layer.weight.data = layer.weight.data.float()
if shared.opts.diffusers_offload_mode != "none":
return_device = layer.weight.data.device
else:
return_device = devices.device
layer.weight.data = layer.weight.data.to(devices.device, dtype=torch.float32)
if is_asym_mode:
level_low = 0
@@ -127,8 +131,11 @@ def nncf_compress_layer(layer, num_bits, is_asym_mode, torch_dtype=None, quant_c
result_dtype=torch_dtype
)
layer.register_pre_forward_operation(decompressor)
compressed_weight = decompressor.pack_weight(compressed_weight)
compressed_weight = compressed_weight.to(return_device)
decompressor = decompressor.to(return_device)
layer.register_pre_forward_operation(decompressor)
layer.weight.requires_grad = False
layer.weight.data = compressed_weight